app-wide display unit in specviz (#2048)

* basic app-wide display unit in specviz (not all plugins are updated yet)
* unit-aware select component which maps to glue-supported unit strings
* implement use_display_units option for get_data, get_subsets

Co-authored-by: Jesse Averbukh <javerbukh@gmail.com>
Co-authored-by: Kyle Conroy <kyleconroy@gmail.com>

CHANGES.rst

@@ -33,6 +33,8 @@ Mosviz
 Specviz
 ^^^^^^^
 
+* Re-enabled unit conversion support. [#2048]
+
 Specviz2d
 ^^^^^^^^^
 

docs/specviz/plugins.rst

@@ -118,15 +118,8 @@ To export the table into the notebook via the API, call
 Unit Conversion
 ===============
 
-.. note::
-
-    This plugin is temporarily disabled. Effort to improve it is being
-    tracked at https://github.com/spacetelescope/jdaviz/issues/1972 .
-
 The spectral flux density and spectral axis units can be converted
-using the Unit Conversion plugin.  The Spectrum1D object to be
-converted is the currently selected spectrum in the spectrum viewer :guilabel:`Data`
-icon in the viewer toolbar.
+using the Unit Conversion plugin.
 
 Select the frequency, wavelength, or energy unit in the
 :guilabel:`New Spectral Axis Unit` pulldown
@@ -135,11 +128,8 @@ Select the frequency, wavelength, or energy unit in the
 Select the flux density unit in the :guilabel:`New Flux Unit` pulldown
 (e.g., Jansky, W/(Hz/m2), ph/(Angstrom cm2 s)).
 
-The :guilabel:`Apply` button will convert the flux density and/or
-spectral axis units and create a new Spectrum1D object that
-is automatically switched to in the spectrum viewer.
-The name of the new Spectrum1D object is "_units_copy_" plus
-the flux and spectral units of the spectrum.
+Note that this affects the default units in all viewers and plugins, where applicable,
+but does not affect the underlying data.
 
 .. _line-lists:
 

jdaviz/app.py

@@ -9,9 +9,9 @@
 from ipywidgets import widget_serialization
 import ipyvue
 
+from astropy import units as u
 from astropy.nddata import CCDData, NDData
 from astropy.io import fits
-from astropy import units as u
 from astropy.coordinates import Angle
 from regions import PixCoord, CirclePixelRegion, RectanglePixelRegion, EllipsePixelRegion
 
@@ -27,7 +27,7 @@
 from glue.config import colormaps, data_translator
 from glue.config import settings as glue_settings
 from glue.core import BaseData, HubListener, Data, DataCollection
-from glue.core.link_helpers import LinkSame
+from glue.core.link_helpers import LinkSame, LinkSameWithUnits
 from glue.plugins.wcs_autolinking.wcs_autolinking import WCSLink, IncompatibleWCS
 from glue.core.message import (DataCollectionAddMessage,
                                DataCollectionDeleteMessage,
@@ -38,6 +38,7 @@
 from glue.core.subset import (Subset, RangeSubsetState, RoiSubsetState,
                               CompositeSubsetState, InvertState)
 from glue.core.roi import CircularROI, EllipticalROI, RectangularROI
+from glue.core.units import unit_converter
 from glue_astronomy.spectral_coordinates import SpectralCoordinates
 from glue_jupyter.app import JupyterApplication
 from glue_jupyter.common.toolbar_vuetify import read_icon
@@ -68,10 +69,51 @@
                  mask=['mask', 'dq'])
 
 
+@unit_converter('custom-jdaviz')
+class UnitConverterWithSpectral:
+
+    def equivalent_units(self, data, cid, units):
+        if cid.label == "flux":
+            eqv = u.spectral_density(1 * u.m)  # Value does not matter here.
+            list_of_units = set(list(map(str, u.Unit(units).find_equivalent_units(
+                include_prefix_units=True, equivalencies=eqv))) + [
+                    'Jy', 'mJy', 'uJy',
+                    'W / (m2 Hz)', 'W / (Hz m2)',  # Order is different in astropy v5.3
+                    'eV / (s m2 Hz)', 'eV / (Hz s m2)',
+                    'erg / (s cm2)',
+                    'erg / (s cm2 Angstrom)', 'erg / (Angstrom s cm2)',
+                    'erg / (s cm2 Hz)', 'erg / (Hz s cm2)',
+                    'ph / (s cm2 Angstrom)', 'ph / (Angstrom s cm2)',
+                    'ph / (s cm2 Hz)', 'ph / (Hz s cm2)'
+                ])
+        else:  # spectral axis
+            # prefer Hz over Bq and um over micron
+            exclude = {'Bq', 'micron'}
+            list_of_units = set(list(map(str, u.Unit(units).find_equivalent_units(
+                include_prefix_units=True, equivalencies=u.spectral())))) - exclude
+        return list_of_units
+
+    def to_unit(self, data, cid, values, original_units, target_units):
+        # Given a glue data object (data), a component ID (cid), the values
+        # to convert, and the original and target units of the values, this method
+        # should return the converted values. Note that original_units
+        # gives the units of the values array, which might not be the same
+        # as the original native units of the component in the data.
+        if cid.label == "flux":
+            spec = data.get_object(cls=Spectrum1D)
+            eqv = u.spectral_density(spec.spectral_axis)
+        else:  # spectral axis
+            eqv = u.spectral()
+        return (values * u.Unit(original_units)).to_value(u.Unit(target_units), equivalencies=eqv)
+
+
 # Set default opacity for data layers to 1 instead of 0.8 in
 # some glue-core versions
 glue_settings.DATA_ALPHA = 1
 
+# Enable spectrum unit conversion.
+glue_settings.UNIT_CONVERTER = 'custom-jdaviz'
+
 custom_components = {'j-tooltip': 'components/tooltip.vue',
                      'j-external-link': 'components/external_link.vue',
                      'j-docs-link': 'components/docs_link.vue',
@@ -446,7 +488,7 @@ def _link_new_data(self, reference_data=None, data_to_be_linked=None):
         if isinstance(linked_data.coords, SpectralCoordinates):
             wc_old = ref_data.world_component_ids[-1]
             wc_new = linked_data.world_component_ids[0]
-            self.data_collection.add_link(LinkSame(wc_old, wc_new))
+            self.data_collection.add_link(LinkSameWithUnits(wc_old, wc_new))
             return
 
         try:
@@ -492,8 +534,8 @@ def _link_new_data(self, reference_data=None, data_to_be_linked=None):
                 else:
                     continue
 
-                links.append(LinkSame(ref_data.pixel_component_ids[ref_index],
-                                      linked_data.pixel_component_ids[linked_index]))
+                links.append(LinkSameWithUnits(ref_data.pixel_component_ids[ref_index],
+                                               linked_data.pixel_component_ids[linked_index]))
 
             dc.add_link(links)
 
@@ -828,7 +870,8 @@ def get_subsets_from_viewer(self, viewer_reference, data_label=None, subset_type
         return regions
 
     def get_subsets(self, subset_name=None, spectral_only=False,
-                    spatial_only=False, object_only=False, simplify_spectral=True):
+                    spatial_only=False, object_only=False, simplify_spectral=True,
+                    use_display_units=False):
         """
         Returns all branches of glue subset tree in the form that subset plugin can recognize.
 
@@ -845,6 +888,8 @@ def get_subsets(self, subset_name=None, spectral_only=False,
             leave out the region class name and glue_state.
         simplify_spectral : bool
             Return a composite spectral subset collapsed to a simplified SpectralRegion.
+        use_display_units: bool, optional
+            Whether to convert to the display units defined in the <unit-conversion> plugin.
 
         Returns
         -------
@@ -863,15 +908,17 @@ def get_subsets(self, subset_name=None, spectral_only=False,
             if isinstance(subset.subset_state, CompositeSubsetState):
                 # Region composed of multiple ROI or Range subset
                 # objects that must be traversed
-                subset_region = self.get_sub_regions(subset.subset_state, simplify_spectral)
+                subset_region = self.get_sub_regions(subset.subset_state,
+                                                     simplify_spectral, use_display_units)
             elif isinstance(subset.subset_state, RoiSubsetState):
                 # 3D regions represented as a dict including an
                 # AstropyRegion object if possible
                 subset_region = self._get_roi_subset_definition(subset.subset_state)
             elif isinstance(subset.subset_state, RangeSubsetState):
                 # 2D regions represented as SpectralRegion objects
                 subset_region = self._get_range_subset_bounds(subset.subset_state,
-                                                              simplify_spectral)
+                                                              simplify_spectral,
+                                                              use_display_units)
             else:
                 # subset.subset_state can be an instance of MaskSubsetState
                 # or something else we do not know how to handle
@@ -945,17 +992,23 @@ def _remove_duplicate_bounds(self, spec_regions):
                 regions_no_dups += region
         return regions_no_dups
 
-    def _get_range_subset_bounds(self, subset_state, simplify_spectral=True):
+    def _get_range_subset_bounds(self, subset_state, simplify_spectral=True, use_display_units=False):
         # TODO: Use global display units
         # units = dc[0].data.coords.spectral_axis.unit
         viewer = self.get_viewer(self._jdaviz_helper. _default_spectrum_viewer_reference_name)
         data = viewer.data()
-        if viewer:
-            units = u.Unit(viewer.state.x_display_unit)
-        elif data and len(data) > 0 and isinstance(data[0], Spectrum1D):
+        if data and len(data) > 0 and isinstance(data[0], Spectrum1D):
             units = data[0].spectral_axis.unit
         else:
             raise ValueError("Unable to find spectral axis units")
+        if use_display_units:
+            # converting may result in flipping order (wavelength <-> frequency)
+            ret_units = self._get_display_unit('spectral')
+            subset_bounds = [(subset_state.lo * units).to(ret_units, u.spectral()),
+                             (subset_state.hi * units).to(ret_units, u.spectral())]
+            spec_region = SpectralRegion(min(subset_bounds), max(subset_bounds))
+        else:
+            spec_region = SpectralRegion(subset_state.lo * units, subset_state.hi * units)
 
         spec_region = SpectralRegion(subset_state.lo * units, subset_state.hi * units)
         if not simplify_spectral:
@@ -992,12 +1045,14 @@ def _get_roi_subset_definition(self, subset_state):
                  "region": roi_as_region,
                  "subset_state": subset_state}]
 
-    def get_sub_regions(self, subset_state, simplify_spectral=True):
+    def get_sub_regions(self, subset_state, simplify_spectral=True, use_display_units=False):
 
         if isinstance(subset_state, CompositeSubsetState):
             if subset_state and hasattr(subset_state, "state2") and subset_state.state2:
-                one = self.get_sub_regions(subset_state.state1, simplify_spectral)
-                two = self.get_sub_regions(subset_state.state2, simplify_spectral)
+                one = self.get_sub_regions(subset_state.state1,
+                                           simplify_spectral, use_display_units)
+                two = self.get_sub_regions(subset_state.state2,
+                                           simplify_spectral, use_display_units)
 
                 if isinstance(one, list) and "glue_state" in one[0]:
                     one[0]["glue_state"] = subset_state.__class__.__name__
@@ -1057,15 +1112,26 @@ def get_sub_regions(self, subset_state, simplify_spectral=True):
             else:
                 # This gets triggered in the InvertState case where state1
                 # is an object and state2 is None
-                return self.get_sub_regions(subset_state.state1, simplify_spectral)
+                return self.get_sub_regions(subset_state.state1,
+                                            simplify_spectral, use_display_units)
         elif subset_state is not None:
             # This is the leaf node of the glue subset state tree where
             # a subset_state is either ROI or Range.
             if isinstance(subset_state, RoiSubsetState):
                 return self._get_roi_subset_definition(subset_state)
 
             elif isinstance(subset_state, RangeSubsetState):
-                return self._get_range_subset_bounds(subset_state, simplify_spectral)
+                return self._get_range_subset_bounds(subset_state,
+                                                     simplify_spectral, use_display_units)
+
+    def _get_display_unit(self, axis):
+        if self._jdaviz_helper is None or self._jdaviz_helper.plugins.get('Unit Conversion') is None:  # noqa
+            raise ValueError("cannot detect unit conversion plugin")
+        try:
+            return getattr(self._jdaviz_helper.plugins.get('Unit Conversion')._obj,
+                           f'{axis}_unit_selected')
+        except AttributeError:
+            raise ValueError(f"could not find display unit for axis='{axis}'")
 
     def add_data(self, data, data_label=None, notify_done=True):
         """
@@ -1247,19 +1313,6 @@ def add_data_to_viewer(self, viewer_reference, data_label,
 
         self.set_data_visibility(viewer_item, data_label, visible=visible, replace=clear_other_data)
 
-    def _set_plot_axes_labels(self, viewer_id):
-        """
-        Sets the plot axes labels to be the units of the data to be loaded.
-
-        Parameters
-        ----------
-        viewer_id : str
-            The UUID associated with the desired viewer item.
-        """
-        viewer = self._viewer_by_id(viewer_id)
-
-        viewer.set_plot_axes()
-
     def remove_data_from_viewer(self, viewer_reference, data_label):
         """
         Removes a data set from the specified viewer.
@@ -1632,11 +1685,8 @@ def set_data_visibility(self, viewer_reference, data_label, visible=True, replac
         # active data.
         viewer_data_labels = [layer.layer.label for layer in viewer.layers]
         if len(viewer_data_labels) > 0 and getattr(self._jdaviz_helper, '_in_batch_load', 0) == 0:
-            active_data = self.data_collection[viewer_data_labels[0]]
-            if (hasattr(active_data, "_preferred_translation")
-                    and active_data._preferred_translation is not None):
-                self._set_plot_axes_labels(viewer_id)
-
+            # This "if" is nested on purpose to make parent "if" available
+            # for other configs in the future, as needed.
             if self.config == 'imviz':
                 viewer.on_limits_change()  # Trigger compass redraw
 

jdaviz/configs/cubeviz/helper.py

@@ -120,7 +120,8 @@ def specviz(self):
             self._specviz = Specviz(app=self.app)
         return self._specviz
 
-    def get_data(self, data_label=None, cls=None, subset_to_apply=None, function=None):
+    def get_data(self, data_label=None, cls=None, subset_to_apply=None, function=None,
+                 use_display_units=False):
         """
         Returns data with name equal to data_label of type cls with subsets applied from
         subset_to_apply.
@@ -151,7 +152,7 @@ def get_data(self, data_label=None, cls=None, subset_to_apply=None, function=Non
         elif function is False:
             function = None
         return self._get_data(data_label=data_label, cls=cls, subset_to_apply=subset_to_apply,
-                              function=function)
+                              function=function, use_display_units=use_display_units)
 
 
 def layer_is_cube_image_data(layer):

jdaviz/configs/default/plugins/line_lists/line_lists.py

@@ -199,7 +199,8 @@ def _on_viewer_data_changed(self, msg=None):
         self._on_spectrum_viewer_limits_changed()  # will also trigger _auto_slider_step
 
         # set the choices (and default) for the units for new custom lines
-        self.custom_unit_choices = create_spectral_equivalencies_list(viewer_data)
+        self.custom_unit_choices = create_spectral_equivalencies_list(
+            viewer_data.spectral_axis.unit)
         self.custom_unit = str(viewer_data.spectral_axis.unit)
 
     def _parse_redshift_msg(self, msg):
@@ -410,6 +411,10 @@ def vue_slider_reset(self, event):
 
     def _on_spectrum_viewer_limits_changed(self, event=None):
         sv = self.app.get_viewer(self._default_spectrum_viewer_reference_name)
+
+        if sv.state.x_min is None or sv.state.x_max is None:
+            return
+
         self.spectrum_viewer_min = float(sv.state.x_min)
         self.spectrum_viewer_max = float(sv.state.x_max)
 

jdaviz/configs/default/plugins/viewers.py

@@ -115,8 +115,6 @@ def _get_layer_info(layer):
                         return "mdi-chart-bell-curve", ""
             return "", suffix
 
-            return '', ''
-
         visible_layers = {}
         for layer in self.state.layers[::-1]:
             if layer.visible: