From 800cdb3038a98bf7e912d8e94d785577212d8c11 Mon Sep 17 00:00:00 2001
From: Thomas Robitaille <thomas.robitaille@gmail.com>
Date: Mon, 12 Aug 2024 14:34:41 +0100
Subject: [PATCH] Initial work on adding unit support for scatter viewer

---
 glue_jupyter/bqplot/scatter/layer_artist.py        | 14 +++++++++++++-
 glue_jupyter/bqplot/scatter/viewer.py              | 10 ++++++++--
 .../common/state_widgets/viewer_scatter.py         |  8 ++++++++
 .../common/state_widgets/viewer_scatter.vue        |  6 ++++++
 4 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/glue_jupyter/bqplot/scatter/layer_artist.py b/glue_jupyter/bqplot/scatter/layer_artist.py
index 6434f052..b4b9d251 100644
--- a/glue_jupyter/bqplot/scatter/layer_artist.py
+++ b/glue_jupyter/bqplot/scatter/layer_artist.py
@@ -9,6 +9,7 @@
 from glue.viewers.scatter.layer_artist import DensityMapLimits
 from glue.viewers.scatter.state import ScatterLayerState
 from glue_jupyter.bqplot.scatter.scatter_density_mark import GenericDensityMark
+from glue.core.units import UnitConverter
 
 from ...utils import colormap_to_hexlist, float_or_none
 from ..compatibility import ScatterGL, LinesGL
@@ -175,7 +176,6 @@ def _update_data(self):
                 x = ensure_numerical(self.layer[self._viewer_state.x_att].ravel())
                 if x.dtype.kind == "M":
                     x = datetime64_to_mpl(x)
-
         except (IncompatibleAttribute, IndexError):
             # The following includes a call to self.clear()
             self.disable_invalid_attributes(self._viewer_state.x_att)
@@ -195,6 +195,18 @@ def _update_data(self):
         else:
             self.enable()
 
+        converter = UnitConverter()
+
+        x = converter.to_unit(self._viewer_state.x_att.parent,
+                              self._viewer_state.x_att,
+                              x,
+                              self._viewer_state.x_display_unit)
+
+        y = converter.to_unit(self._viewer_state.y_att.parent,
+                              self._viewer_state.y_att,
+                              y,
+                              self._viewer_state.y_display_unit)
+
         if self.state.markers_visible:
 
             if self.state.density_map:
diff --git a/glue_jupyter/bqplot/scatter/viewer.py b/glue_jupyter/bqplot/scatter/viewer.py
index 7381053e..e1cf8230 100644
--- a/glue_jupyter/bqplot/scatter/viewer.py
+++ b/glue_jupyter/bqplot/scatter/viewer.py
@@ -36,7 +36,13 @@ def __init__(self, *args, **kwargs):
     def _update_axes(self, *args):
 
         if self.state.x_att is not None:
-            self.state.x_axislabel = str(self.state.x_att)
+            if self.state.x_display_unit:
+                self.state.x_axislabel = str(self.state.x_att) + f' [{self.state.x_display_unit}]'
+            else:
+                self.state.x_axislabel = str(self.state.x_att)
 
         if self.state.y_att is not None:
-            self.state.y_axislabel = str(self.state.y_att)
+            if self.state.y_display_unit:
+                self.state.y_axislabel = str(self.state.y_att) + f' [{self.state.y_display_unit}]'
+            else:
+                self.state.y_axislabel = str(self.state.y_att)
diff --git a/glue_jupyter/common/state_widgets/viewer_scatter.py b/glue_jupyter/common/state_widgets/viewer_scatter.py
index 018fd049..6a7fa68a 100644
--- a/glue_jupyter/common/state_widgets/viewer_scatter.py
+++ b/glue_jupyter/common/state_widgets/viewer_scatter.py
@@ -19,6 +19,12 @@ class ScatterViewerStateWidget(v.VuetifyTemplate):
     y_att_items = traitlets.List().tag(sync=True)
     y_att_selected = traitlets.Int(allow_none=True).tag(sync=True)
 
+    x_display_unit_items = traitlets.List().tag(sync=True)
+    x_display_unit_selected = traitlets.Int(allow_none=True).tag(sync=True)
+
+    y_display_unit_items = traitlets.List().tag(sync=True)
+    y_display_unit_selected = traitlets.Int(allow_none=True).tag(sync=True)
+
     def __init__(self, viewer_state):
 
         super().__init__()
@@ -28,3 +34,5 @@ def __init__(self, viewer_state):
 
         link_glue_choices(self, viewer_state, "x_att")
         link_glue_choices(self, viewer_state, "y_att")
+        link_glue_choices(self, viewer_state, 'x_display_unit')
+        link_glue_choices(self, viewer_state, 'y_display_unit')
diff --git a/glue_jupyter/common/state_widgets/viewer_scatter.vue b/glue_jupyter/common/state_widgets/viewer_scatter.vue
index 8c888e42..63f0b96e 100644
--- a/glue_jupyter/common/state_widgets/viewer_scatter.vue
+++ b/glue_jupyter/common/state_widgets/viewer_scatter.vue
@@ -6,6 +6,12 @@
         <div>
             <v-select label="y axis" :items="y_att_items" v-model="y_att_selected" hide-details />
         </div>
+        <div>
+            <v-select :items="x_display_unit_items" label="x axis units" v-model="x_display_unit_selected" hide-details/>
+        </div>
+        <div>
+            <v-select :items="y_display_unit_items" label="y axis units" v-model="y_display_unit_selected" hide-details/>
+        </div>
         <div>
             <v-subheader class="pl-0 slider-label">show axes</v-subheader>
             <v-switch v-model="glue_state.show_axes" hide-details style="margin-top: 0"/>