From 0abc22b4225b2853e1d98dceef879ec244f779e5 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Mon, 9 Sep 2024 15:42:41 -0400
Subject: [PATCH 01/22] WIP: refactor/simplify messaging and logic in unit
 conversion

---
 .../unit_conversion/unit_conversion.py        | 428 +++++++-----------
 jdaviz/core/events.py                         |   6 +
 jdaviz/core/template_mixin.py                 |   2 -
 jdaviz/core/validunits.py                     |   2 +-
 4 files changed, 179 insertions(+), 259 deletions(-)

diff --git a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
index ae7ec6093b..30d8f63307 100644
--- a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
+++ b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
@@ -4,7 +4,8 @@
 import numpy as np
 from traitlets import List, Unicode, observe, Bool
 
-from jdaviz.core.events import GlobalDisplayUnitChanged, AddDataToViewerMessage
+from jdaviz.configs.default.plugins.viewers import JdavizProfileView
+from jdaviz.core.events import GlobalDisplayUnitChanged, AddDataMessage
 from jdaviz.core.registries import tray_registry
 from jdaviz.core.template_mixin import (PluginTemplateMixin, UnitSelectPluginComponent,
                                         SelectPluginComponent, PluginUserApi)
@@ -32,6 +33,19 @@ def _valid_glue_display_unit(unit_str, sv, axis='x'):
     return choices_str[ind]
 
 
+def _flux_to_sb_unit(flux_unit, angle_unit):
+    if angle_unit not in ['pix2', 'sr']:
+        sb_unit = flux_unit
+    elif '(' in flux_unit:
+        pos = flux_unit.rfind(')')
+        sb_unit = flux_unit[:pos] + ' ' + angle_unit + flux_unit[pos:]
+    else:
+        # append angle if there are no parentheses
+        sb_unit = flux_unit + ' / ' + angle_unit
+
+    return sb_unit
+
+
 @tray_registry('g-unit-conversion', label="Unit Conversion",
                viewer_requirements='spectrum')
 class UnitConversion(PluginTemplateMixin):
@@ -47,12 +61,14 @@ class UnitConversion(PluginTemplateMixin):
     * :meth:`~jdaviz.core.template_mixin.PluginTemplateMixin.close_in_tray`
     * ``spectral_unit`` (:class:`~jdaviz.core.template_mixin.UnitSelectPluginComponent`):
       Global unit to use for all spectral axes.
-    * ``spectral_y_type`` (:class:`~jdaviz.core.template_mixin.SelectPluginComponent`):
-      Select the y-axis physical type for the spectrum-viewer.
     * ``flux_unit`` (:class:`~jdaviz.core.template_mixin.UnitSelectPluginComponent`):
       Global display unit for flux axis.
     * ``angle_unit`` (:class:`~jdaviz.core.template_mixin.UnitSelectPluginComponent`):
       Solid angle unit.
+    * ``sb_unit`` (str): Read-only property for the current surface brightness unit,
+      derived from the set values of ``flux_unit`` and ``angle_unit``.
+    * ``spectral_y_type`` (:class:`~jdaviz.core.template_mixin.SelectPluginComponent`):
+      Select the y-axis physical type for the spectrum-viewer (applicable only to Cubeviz).
     """
     template_file = __file__, "unit_conversion.vue"
 
@@ -78,7 +94,7 @@ class UnitConversion(PluginTemplateMixin):
     spectral_y_type_items = List().tag(sync=True)
     spectral_y_type_selected = Unicode().tag(sync=True)
 
-    # This is used a warning message if False. This can be changed from
+    # This shows an in-line warning message if False. This can be changed from
     # bool to unicode when we eventually handle inputing this value if it
     # doesn't exist in the FITS header
     pixar_sr_exists = Bool(True).tag(sync=True)
@@ -95,30 +111,31 @@ def __init__(self, *args, **kwargs):
             # this force all to sync?)
             self.disabled_msg = f'This plugin is temporarily disabled in {self.config}. Effort to improve it is being tracked at GitHub Issue 1972.'  # noqa
 
-        # TODO [markers]: existing markers need converting
-        self.spectrum_viewer.state.add_callback('x_display_unit',
-                                                self._on_glue_x_display_unit_changed)
-
-        self.spectrum_viewer.state.add_callback('y_display_unit',
-                                                self._on_glue_y_display_unit_changed)
-
-        self.session.hub.subscribe(self, AddDataToViewerMessage,
-                                   handler=self._find_and_convert_contour_units)
+        self.session.hub.subscribe(self, AddDataMessage,
+                                   handler=self._on_add_data_to_viewer)
 
         self.has_spectral = self.config in ('specviz', 'cubeviz', 'specviz2d', 'mosviz')
         self.spectral_unit = UnitSelectPluginComponent(self,
                                                        items='spectral_unit_items',
                                                        selected='spectral_unit_selected')
+        # TODO: can choices be defined in the object init?
+        self.spectral_unit.choices = create_spectral_equivalencies_list(u.Hz)
+
 
         self.has_flux = self.config in ('specviz', 'cubeviz', 'specviz2d', 'mosviz')
         self.flux_unit = UnitSelectPluginComponent(self,
                                                    items='flux_unit_items',
                                                    selected='flux_unit_selected')
+        # TODO: double check if using Hz to grab the equivalencies will cause a problem for wavelength
+        # NOTE: will switch to count only if first data loaded into viewer in in counts
+        self.flux_unit.choices = create_flux_equivalencies_list(u.Jy, u.Hz)
 
         self.has_angle = self.config in ('cubeviz', 'specviz', 'mosviz')
         self.angle_unit = UnitSelectPluginComponent(self,
                                                     items='angle_unit_items',
                                                     selected='angle_unit_selected')
+        # NOTE: will switch to pix2 only if first data loaded into viewer is in pix2 units
+        self.angle_unit.choices = create_angle_equivalencies_list(u.sr)
 
         self.has_sb = self.has_angle or self.config in ('imviz',)
         # NOTE: always read_only, exposed through sb_unit property
@@ -133,6 +150,7 @@ def __init__(self, *args, **kwargs):
                                                      selected='spectral_y_type_selected',
                                                      manual_options=['Surface Brightness', 'Flux'])
 
+
     @property
     def user_api(self):
         expose = []
@@ -153,266 +171,164 @@ def user_api(self):
 
     @property
     def sb_unit(self):
+        # expose selected surface-brightness unit as read-only (rather than exposing a select object)
         return self.sb_unit_selected
 
-    def _on_glue_x_display_unit_changed(self, x_unit_str):
-        if x_unit_str is None:
-            return
-        self.spectrum_viewer.set_plot_axes()
-        if x_unit_str != self.spectral_unit.selected:
-            x_unit_str = _valid_glue_display_unit(x_unit_str, self.spectrum_viewer, 'x')
-            x_unit = u.Unit(x_unit_str)
-            choices = create_spectral_equivalencies_list(x_unit)
-            # ensure that original entry is in the list of choices
-            if not np.any([x_unit == u.Unit(choice) for choice in choices]):
-                choices = [x_unit_str] + choices
-            self.spectral_unit.choices = choices
-            # in addition to the jdaviz options, allow the user to set any glue-valid unit
-            # which would then be appended on to the list of choices going forward
-            self.spectral_unit._addl_unit_strings = self.spectrum_viewer.state.__class__.x_display_unit.get_choices(self.spectrum_viewer.state)  # noqa
-            self.spectral_unit.selected = x_unit_str
-            if not len(self.flux_unit.choices) or not len(self.angle_unit.choices):
-                # in case flux_unit was triggered first (but could not be set because there
-                # was no spectral_unit to determine valid equivalencies)
-                self._on_glue_y_display_unit_changed(self.spectrum_viewer.state.y_display_unit)
-
-    def _on_glue_y_display_unit_changed(self, y_unit_str):
-        if y_unit_str is None:
-            return
-        if self.spectral_unit.selected == "":
-            # no spectral unit set yet, cannot determine equivalencies
-            # setting the spectral unit will check len(spectral_y_type_unit.choices)
-            # and call this manually in the case that that is triggered second.
-            return
-        self.spectrum_viewer.set_plot_axes()
-
-        x_unit = u.Unit(self.spectral_unit.selected)
-        y_unit_str = _valid_glue_display_unit(y_unit_str, self.spectrum_viewer, 'y')
-        y_unit = u.Unit(y_unit_str)
-        y_unit_solid_angle = check_if_unit_is_per_solid_angle(y_unit_str, return_unit=True)
-
-        if not check_if_unit_is_per_solid_angle(y_unit_str) and y_unit_str != self.flux_unit.selected:  # noqa
-            flux_choices = create_flux_equivalencies_list(y_unit, x_unit)
-            # ensure that original entry is in the list of choices
-            if not np.any([y_unit == u.Unit(choice) for choice in flux_choices]):
-                flux_choices = [y_unit_str] + flux_choices
-
-            self.flux_unit.choices = flux_choices
-            self.flux_unit.selected = y_unit_str
-
-        # if the y-axis is set to surface brightness,
-        # untranslatable units need to be removed from the flux choices
-        if y_unit_solid_angle:
-            flux_choices = [(y_unit * y_unit_solid_angle).to_string()]
-            flux_choices += create_flux_equivalencies_list(y_unit * y_unit_solid_angle, x_unit)
-            self.flux_unit.choices = flux_choices
-            flux_unit = str(y_unit * y_unit_solid_angle)
-            # We need to set the angle_unit before triggering _on_flux_unit_changed via
-            # setting self.flux_unit.selected below, or the lack of angle unit will make it think
-            # we're in Flux units.
-            self.angle_unit.choices = create_angle_equivalencies_list(y_unit)
-            self.angle_unit.selected = self.angle_unit.choices[0]
-            if flux_unit in self.flux_unit.choices and flux_unit != self.flux_unit.selected:
-                self.flux_unit.selected = flux_unit
-
-        # sets the angle unit drop down and the surface brightness read-only text
-        if self.app.data_collection[0]:
-            dc_unit = self.app.data_collection[0].get_component("flux").units
-
-            # angle choices will be angle equivalencies to the solid-angle component of the cube
-            dc_solid_angle_unit = check_if_unit_is_per_solid_angle(dc_unit, return_unit=True)
-
-            self.angle_unit.choices = create_angle_equivalencies_list(dc_solid_angle_unit)
-            self.angle_unit.selected = self.angle_unit.choices[0]
-            self.sb_unit_selected = self._append_angle_correctly(
-                self.flux_unit.selected,
-                self.angle_unit.selected
-            )
-            if self.angle_unit.selected == 'pix':
-                mouseover_unit = self.flux_unit.selected
-            else:
-                mouseover_unit = self.sb_unit_selected
-            self.hub.broadcast(GlobalDisplayUnitChanged("sb", mouseover_unit, sender=self))
-
-        else:
-            # if cube was loaded in flux units, we still need to broadcast
-            # a 'sb' message for mouseover info. this should be removed when
-            # unit change messaging is improved and is a temporary fix
-            self.hub.broadcast(GlobalDisplayUnitChanged('sb',
-                                                        self.flux_unit.selected,
-                                                        sender=self))
-
-            if not self.flux_unit.selected:
-                y_display_unit = self.spectrum_viewer.state.y_display_unit
-                flux_unit_str = str(u.Unit(y_display_unit * y_unit_solid_angle))
-                self.flux_unit.selected = flux_unit_str
-
-    @observe('spectral_unit_selected')
-    def _on_spectral_unit_changed(self, *args):
-        xunit = _valid_glue_display_unit(self.spectral_unit.selected, self.spectrum_viewer, 'x')
-        if self.spectrum_viewer.state.x_display_unit != xunit:
-            self.spectrum_viewer.state.x_display_unit = xunit
-            self.hub.broadcast(GlobalDisplayUnitChanged('spectral',
-                               self.spectral_unit.selected,
-                               sender=self))
-
-    @observe('spectral_y_type_selected')
-    def _on_spectral_y_type_selected(self, msg):
-        """
-        Observes toggle between surface brightness or flux selection for
-        spectrum viewer to trigger translation.
-        """
-
-        if msg.get('name') == 'spectral_y_type_selected':
-            self._translate(self.spectral_y_type_selected)
-
-    @observe('flux_unit_selected')
-    def _on_flux_unit_changed(self, msg):
-
+    def _on_add_data_to_viewer(self, msg):
+        # toggle warning message for cubes without PIXAR_SR defined
+        if self.config == 'cubeviz':
+            # NOTE: this assumes data_collection[0] is the science (flux/sb) cube
+            if (
+                len(self.app.data_collection) > 0
+                and not self.app.data_collection[0].meta.get('PIXAR_SR')
+            ):
+                self.pixar_sr_exists = False
+
+        viewer = msg.viewer
+ 
+        if isinstance(viewer, JdavizProfileView):  # TODO: think about rampviz
+            if viewer.state.x_display_unit == self.spectral_unit_selected and viewer.state.y_display_unit == self.app._get_display_unit('spectral_y'):
+                # data already existed in this viewer and display units were already set
+                return
+            if not (len(self.spectral_unit_selected)
+                    and len(self.flux_unit_selected)
+                    and len(self.angle_unit_selected)
+                    and (self.config != 'cubeviz' or len(self.spectral_y_type_selected))):
+
+                spec = msg.data.get_object()
+
+                self.spectral_unit._addl_unit_strings = self.spectrum_viewer.state.__class__.x_display_unit.get_choices(self.spectrum_viewer.state)
+                if not len(self.spectral_unit_selected):
+                    self.spectral_unit.selected = str(spec.spectral_axis.unit)
+
+                angle_unit = check_if_unit_is_per_solid_angle(spec.flux.unit, return_unit=True)
+                flux_unit = spec.flux.unit if angle_unit is None else spec.flux.unit * angle_unit
+
+                if not self.flux_unit_selected:
+                    if flux_unit in (u.count, u.DN):
+                        self.flux_unit.choices = [flux_unit]
+                    self.flux_unit.selected = str(flux_unit)
+
+                if not self.angle_unit_selected:
+                    if angle_unit == u.pix**2:
+                        self.angle_unit.choices = ['pix2']
+
+                    if angle_unit is None:
+                        # default to sr if input spectrum is not in surface brightness units
+                        # TODO: for cubeviz, should we check the cube itself?
+                        self.angle_unit.selected = 'sr'
+                    else:
+                        self.angle_unit.selected = str(angle_unit)
+
+                if not len(self.spectral_y_type_selected):
+                    # set spectral_y_type_selected to 'Flux' if the y-axis unit is not per solid angle
+                    if angle_unit is None:
+                        self.spectral_y_type_selected = 'Flux'
+                    else:
+                        self.spectral_y_type_selected = 'Surface Brightness'
+
+                # setting default values will trigger the observes to set the units in _on_unit_selected,
+                # so return here to avoid setting twice
+                return
+
+            # this spectral viewer was empty (did not have display units set yet), but global selections
+            # are available in the plugin, so we'll set them to the viewer here
+            viewer.state.x_display_unit = self.spectral_unit_selected
+            # _handle_spectral_y_unit will call viewer.set_plot_axes()
+            self._handle_spectral_y_unit()
+
+        elif isinstance(viewer, BqplotImageView):
+            # set the attribute display unit (contour and stretch units) for the new layer
+            # NOTE: this assumes that all image data is coerced to surface brightness units
+            layers = [lyr for lyr in msg.viewer.layers if lyr.layer.data.label == msg.data.label]
+            self._handle_attribute_display_unit(self.sb_unit_selected, layers=layers)
+
+
+    @observe('spectral_unit_selected', 'flux_unit_selected',
+             'angle_unit_selected', 'sb_unit_selected',
+             'time_unit_selected')
+    def _on_unit_selected(self, msg):
         """
-        Observes changes in selected flux unit.
-
-        When the selected flux unit changes, a GlobalDisplayUnitChange needs
-        to be broadcasted indicating that the flux unit has changed.
-
-        Note: The 'axis' of the broadcast should always be 'flux', even though a
-        change in flux unit indicates a change in surface brightness unit, because
-        SB is read only, so anything observing for changes in surface brightness
-        should be looking for a change in 'flux' (as well as angle).
+        When any user selection is made, update the relevant viewer(s) with the new unit,
+        and then emit a GlobalDisplayUnitChanged message to notify other plugins of the change.
         """
-
-        if msg.get('name') != 'flux_unit_selected':
-            # not sure when this would be encountered but keeping as a safeguard
-            return
-        if not hasattr(self, 'flux_unit'):
+        if not len(msg.get('new', '')):
+            # empty string, nothing to set yet
             return
-        if not self.flux_unit.choices and self.app.config == 'cubeviz':
-            return
-
-        # various plugins are listening for changes in either flux or sb and
-        # need to be able to filter messages accordingly, so broadcast both when
-        # flux unit is updated. if data was loaded in a flux unit (i.e MJy), it
-        # can be reperesented as a per-pixel surface brightness unit
-        flux_unit = self.flux_unit.selected
-        sb_unit = self._append_angle_correctly(flux_unit, self.angle_unit.selected)
 
-        self.hub.broadcast(GlobalDisplayUnitChanged("flux", flux_unit, sender=self))
-        self.hub.broadcast(GlobalDisplayUnitChanged("sb", sb_unit, sender=self))
+        axis = msg.get('name').split('_')[0]
 
-        spectral_y = sb_unit if self.spectral_y_type == 'Surface Brightness' else flux_unit
+        if axis == 'spectral':
+            xunit = _valid_glue_display_unit(self.spectral_unit.selected, self.spectrum_viewer, 'x')
+            self.spectrum_viewer.state.x_display_unit = xunit
+            self.spectrum_viewer.set_plot_axes()
 
-        yunit = _valid_glue_display_unit(spectral_y, self.spectrum_viewer, 'y')
+        elif axis == 'flux':
+            if len(self.angle_unit_selected):
+                # NOTE: setting sb_unit_selected will call this method again with axis=='sb',
+                # which in turn will call _handle_spectral_y_unit and send a second GlobalDisplayUnitChanged message for sb
+                self.sb_unit_selected = _flux_to_sb_unit(self.flux_unit.selected, self.angle_unit.selected)
 
-        # update spectrum viewer with new y display unit
-        if self.spectrum_viewer.state.y_display_unit != yunit:
-            self.spectrum_viewer.state.y_display_unit = yunit
-            self.spectrum_viewer.reset_limits()
+            if self.spectral_y_type_selected == 'Flux':
+                self._handle_spectral_y_unit()
 
-            # and broacast that there has been a change in the spectral axis y unit
-            # to either a flux or surface brightness unit, for plugins that specifically
-            # care about this toggle selection
-            self.hub.broadcast(GlobalDisplayUnitChanged("spectral_y", spectral_y, sender=self))
+        elif axis == 'angle':
+            if len(self.flux_unit_selected):
+                # NOTE: setting sb_unit_selected will call this method again and send a second GlobalDisplayUnitChanged message for sb
+                self.sb_unit_selected = _flux_to_sb_unit(self.flux_unit.selected, self.angle_unit.selected)
 
-        if not check_if_unit_is_per_solid_angle(self.spectrum_viewer.state.y_display_unit):
-            self.spectral_y_type_selected = 'Flux'
-        else:
-            self.spectral_y_type_selected = 'Surface Brightness'
-
-        # Always send a surface brightness unit to contours
-        if self.spectral_y_type_selected == 'Flux':
-            yunit = self._append_angle_correctly(yunit, self.angle_unit.selected)
-        self._find_and_convert_contour_units(yunit=yunit)
-
-        # for displaying message that PIXAR_SR = 1 if it is not found in the FITS header
-        if (
-            len(self.app.data_collection) > 0
-            and not self.app.data_collection[0].meta.get('PIXAR_SR')
-        ):
-            self.pixar_sr_exists = False
-
-    def _find_and_convert_contour_units(self, msg=None, yunit=None):
-        if not yunit:
-            yunit = self.sb_unit_selected
-
-        if msg is not None:
-            viewers = [self.app.get_viewer(msg.viewer_reference)]
-        else:
-            viewers = self._app._viewer_store.values()
+        elif axis == 'sb':
+            self._handle_attribute_display_unit(self.sb_unit_selected)
 
-        if self.angle_unit_selected is None or self.angle_unit_selected == '':
-            # Can't do this before the plugin is initialized completely
-            return
+            if self.spectral_y_type_selected == 'Surface Brightness':
+                self._handle_spectral_y_unit()
 
-        for viewer in viewers:
-            if not isinstance(viewer, BqplotImageView):
-                continue
-            for layer in viewer.state.layers:
+        elif axis == 'time':
+            pass
 
-                # DQ layer doesn't play nicely with this attribute
-                if "DQ" in layer.layer.label or isinstance(layer.layer, GroupedSubset):
-                    continue
-                elif u.Unit(layer.layer.get_component("flux").units).physical_type != 'surface brightness':  # noqa
-                    continue
-                if hasattr(layer, 'attribute_display_unit'):
-                    layer.attribute_display_unit = yunit
+        # axis (first) argument will be one of: spectral, flux, angle, sb, time
+        self.hub.broadcast(GlobalDisplayUnitChanged(msg.name.split('_')[0],
+                           msg.new, sender=self))
 
-    def _translate(self, spectral_y_type=None):
-
-        # currently unsupported, can be supported with a scale factor
-        if self.app.config == 'specviz':
-            return
-
-        if self.spectrum_viewer.state.y_display_unit:
-            spec_units = u.Unit(self.spectrum_viewer.state.y_display_unit)
-        else:
-            return
-
-        # on instantiation, we set determine flux choices and selection
-        # after surface brightness
-        if not self.flux_unit.choices:
-            return
 
-        selected_display_solid_angle_unit = u.Unit(self.angle_unit_selected)
-        spec_axis_ang_unit = check_if_unit_is_per_solid_angle(spec_units)
-
-        # Surface Brightness -> Flux
-        if spec_axis_ang_unit and spectral_y_type == 'Flux':
-            spec_units *= selected_display_solid_angle_unit
-            # update display units
-            self.spectrum_viewer.state.y_display_unit = str(spec_units)
-
-        # Flux -> Surface Brightness
-        elif (not spec_axis_ang_unit and spectral_y_type == 'Surface Brightness'):
-            spec_units /= selected_display_solid_angle_unit
-            # update display units
-            self.spectrum_viewer.state.y_display_unit = str(spec_units)
-
-        # entered the translator when we shouldn't translate
-        else:
+    @observe('spectral_y_type_selected')
+    def _handle_spectral_y_unit(self, *args):
+        """
+        When the spectral_y_type is changed, or the unit corresponding to the currently selected spectral_y_type is changed,
+        update the y-axis of the spectrum viewer with the new unit, and then emit a GlobalDisplayUnitChanged message to notify
+        """
+        yunit_selected = self.sb_unit_selected if self.spectral_y_type_selected == 'Surface Brightness' else self.flux_unit_selected
+        yunit = _valid_glue_display_unit(yunit_selected, self.spectrum_viewer, 'y')
+        if self.spectrum_viewer.state.y_display_unit == yunit:
             return
-
-        # broadcast that there has been a change in the spectrum viewer y axis,
-        # if translation was completed
-        self.hub.broadcast(GlobalDisplayUnitChanged('spectral_y',
-                                                    spec_units,
-                                                    sender=self))
-        self.spectrum_viewer.reset_limits()
-
-    def _append_angle_correctly(self, flux_unit, angle_unit):
-        if angle_unit not in ['pix2', 'sr']:
-            self.sb_unit_selected = flux_unit
-            return flux_unit
-        if '(' in flux_unit:
-            pos = flux_unit.rfind(')')
-            sb_unit_selected = flux_unit[:pos] + ' ' + angle_unit + flux_unit[pos:]
+        try:
+            self.spectrum_viewer.state.y_display_unit = yunit
+        except ValueError:
+            # may not be data in the viewer, or unit may be incompatible
+            pass
         else:
-            # append angle if there are no parentheses
-            sb_unit_selected = flux_unit + ' / ' + angle_unit
+            self.spectrum_viewer.set_plot_axes()
 
-        if sb_unit_selected:
-            # convert string to and from u.Unit to get rid of any
-            # formatting inconstancies, order of units in string
-            # for a composite unit matters
-            sb_unit_selected = u.Unit(sb_unit_selected).to_string()
+            # broadcast that there has been a change in the spectrum viewer y axis,
+            self.hub.broadcast(GlobalDisplayUnitChanged('spectral_y',
+                                                        yunit,
+                                                        sender=self))
 
-        return sb_unit_selected
+ 
+    def _handle_attribute_display_unit(self, attr_unit, layers=None):
+        """
+        Update the per-layer attribute display unit in glue for image viewers (updating stretch and contour units).
+        """
+        if layers is None:
+            layers = [layer
+                      for viewer in self._app._viewer_store.values() if isinstance(viewer, BqplotImageView)
+                      for layer in viewer.layers]
+ 
+        for layer in layers:
+            # DQ layer doesn't play nicely with this attribute
+            if "DQ" in layer.layer.label or isinstance(layer.layer, GroupedSubset):
+                continue
+            elif u.Unit(layer.layer.get_component("flux").units).physical_type != 'surface brightness':  # noqa
+                continue
+            if hasattr(layer, 'attribute_display_unit'):
+                layer.attribute_display_unit = attr_unit
diff --git a/jdaviz/core/events.py b/jdaviz/core/events.py
index cb12c416b2..b9fdfd5ed5 100644
--- a/jdaviz/core/events.py
+++ b/jdaviz/core/events.py
@@ -92,6 +92,9 @@ def path(self):
 
 
 class AddDataMessage(Message):
+    """
+    Emitted AFTER data is added to a viewer
+    """
     def __init__(self, data, viewer, viewer_id=None, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
@@ -198,6 +201,9 @@ def config(self):
 
 
 class AddDataToViewerMessage(Message):
+    """
+    Emitted to request data is added to a viewer (BEFORE the data is actually added)
+    """
     def __init__(self, viewer_reference, data_label, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
diff --git a/jdaviz/core/template_mixin.py b/jdaviz/core/template_mixin.py
index 92feb3bd70..78e13aef32 100644
--- a/jdaviz/core/template_mixin.py
+++ b/jdaviz/core/template_mixin.py
@@ -1452,8 +1452,6 @@ def __init__(self, plugin, items, selected, viewer,
                            handler=self._on_data_added)
         self.hub.subscribe(self, RemoveDataMessage,
                            handler=lambda _: self._update_layer_items())
-        self.hub.subscribe(self, AddDataToViewerMessage,
-                           handler=self._on_data_added)
         self.hub.subscribe(self, SubsetCreateMessage,
                            handler=lambda _: self._on_subset_created())
         # will need SubsetUpdateMessage for name only (style shouldn't force a full refresh)
diff --git a/jdaviz/core/validunits.py b/jdaviz/core/validunits.py
index cfafb0b08a..394a81194b 100644
--- a/jdaviz/core/validunits.py
+++ b/jdaviz/core/validunits.py
@@ -31,7 +31,7 @@ def create_spectral_equivalencies_list(spectral_axis_unit,
                                                 u.lyr, u.AU, u.pc, u.Bq, u.micron, u.lsec]):
     """Get all possible conversions from current spectral_axis_unit."""
     if spectral_axis_unit in (u.pix, u.dimensionless_unscaled):
-        return []
+        return [spectral_axis_unit]
 
     # Get unit equivalencies.
     try:

From 452eb5e98b0617009d1bd41337a293ffd90ea7d9 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Tue, 10 Sep 2024 13:13:53 -0400
Subject: [PATCH 02/22] generalized solution in place of #3185

---
 jdaviz/configs/specviz/plugins/parsers.py | 10 ----------
 1 file changed, 10 deletions(-)

diff --git a/jdaviz/configs/specviz/plugins/parsers.py b/jdaviz/configs/specviz/plugins/parsers.py
index 0ae7546d95..51fbfc1ee6 100644
--- a/jdaviz/configs/specviz/plugins/parsers.py
+++ b/jdaviz/configs/specviz/plugins/parsers.py
@@ -11,7 +11,6 @@
 from jdaviz.core.events import SnackbarMessage
 from jdaviz.core.registries import data_parser_registry
 from jdaviz.utils import standardize_metadata, download_uri_to_path
-from jdaviz.core.validunits import check_if_unit_is_per_solid_angle
 
 
 __all__ = ["specviz_spectrum1d_parser"]
@@ -160,15 +159,6 @@ def specviz_spectrum1d_parser(app, data, data_label=None, format=None, show_in_v
             # Make metadata layout conform with other viz.
             spec.meta = standardize_metadata(spec.meta)
 
-            # If this is the first loaded data, we want to set spectral y unit type to Flux or
-            # Surface Brightness as appropriate
-            if len(app.data_collection) == 0 and "Unit Conversion" in app._jdaviz_helper.plugins:
-                uc = app._jdaviz_helper.plugins["Unit Conversion"]
-                if check_if_unit_is_per_solid_angle(flux_units):
-                    uc._obj.spectral_y_type = "Surface Brightness"
-                else:
-                    uc._obj.spectral_y_type = "Flux"
-
             app.add_data(spec, data_label[i])
 
             # handle display, with the SpectrumList special case in mind.

From a1b5baefb1045ea747399457f02528ff9ab98517 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Tue, 10 Sep 2024 14:33:17 -0400
Subject: [PATCH 03/22] code cleanup

---
 .../specviz/plugins/unit_conversion/unit_conversion.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
index 30d8f63307..e858e1b537 100644
--- a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
+++ b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
@@ -118,7 +118,6 @@ def __init__(self, *args, **kwargs):
         self.spectral_unit = UnitSelectPluginComponent(self,
                                                        items='spectral_unit_items',
                                                        selected='spectral_unit_selected')
-        # TODO: can choices be defined in the object init?
         self.spectral_unit.choices = create_spectral_equivalencies_list(u.Hz)
 
 
@@ -126,7 +125,6 @@ def __init__(self, *args, **kwargs):
         self.flux_unit = UnitSelectPluginComponent(self,
                                                    items='flux_unit_items',
                                                    selected='flux_unit_selected')
-        # TODO: double check if using Hz to grab the equivalencies will cause a problem for wavelength
         # NOTE: will switch to count only if first data loaded into viewer in in counts
         self.flux_unit.choices = create_flux_equivalencies_list(u.Jy, u.Hz)
 
@@ -138,7 +136,7 @@ def __init__(self, *args, **kwargs):
         self.angle_unit.choices = create_angle_equivalencies_list(u.sr)
 
         self.has_sb = self.has_angle or self.config in ('imviz',)
-        # NOTE: always read_only, exposed through sb_unit property
+        # NOTE: sb_unit is read_only, exposed through sb_unit property
 
         self.has_time = False
         self.time_unit = UnitSelectPluginComponent(self,
@@ -150,7 +148,6 @@ def __init__(self, *args, **kwargs):
                                                      selected='spectral_y_type_selected',
                                                      manual_options=['Surface Brightness', 'Flux'])
 
-
     @property
     def user_api(self):
         expose = []
@@ -185,8 +182,9 @@ def _on_add_data_to_viewer(self, msg):
                 self.pixar_sr_exists = False
 
         viewer = msg.viewer
- 
-        if isinstance(viewer, JdavizProfileView):  # TODO: think about rampviz
+
+        # TODO: when enabling unit-conversion in rampviz, this may need to be more specific or handle other cases for ramp profile viewers
+        if isinstance(viewer, JdavizProfileView):
             if viewer.state.x_display_unit == self.spectral_unit_selected and viewer.state.y_display_unit == self.app._get_display_unit('spectral_y'):
                 # data already existed in this viewer and display units were already set
                 return

From 5bd8241afc3b173d70be7f716deec6ad2ce480d4 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Tue, 10 Sep 2024 15:54:25 -0400
Subject: [PATCH 04/22] generalized solution in place of #3177

---
 jdaviz/configs/specviz/tests/test_viewers.py | 1 -
 jdaviz/core/template_mixin.py                | 9 ---------
 2 files changed, 10 deletions(-)

diff --git a/jdaviz/configs/specviz/tests/test_viewers.py b/jdaviz/configs/specviz/tests/test_viewers.py
index cdbd6ea5ee..e8fab5f95b 100644
--- a/jdaviz/configs/specviz/tests/test_viewers.py
+++ b/jdaviz/configs/specviz/tests/test_viewers.py
@@ -27,7 +27,6 @@ def test_spectrum_viewer_axis_labels(specviz_helper, input_unit, y_axis_label):
     assert (y_axis_label in label)
 
 
-@pytest.mark.xfail(reason="FIXME: Some callback magic needs to happen somewhere.")
 def test_spectrum_viewer_keep_unit_when_removed(specviz_helper, spectrum1d):
     specviz_helper.load_data(spectrum1d, data_label="Test")
     uc = specviz_helper.plugins["Unit Conversion"]
diff --git a/jdaviz/core/template_mixin.py b/jdaviz/core/template_mixin.py
index 78e13aef32..0deba57075 100644
--- a/jdaviz/core/template_mixin.py
+++ b/jdaviz/core/template_mixin.py
@@ -3908,13 +3908,7 @@ def add_results_from_plugin(self, data_item, replace=None, label=None):
                 add_to_viewer_vis = [True]
                 preserved_attributes = [{}]
 
-        enforce_flux_unit = None
         if label in self.app.data_collection:
-            if self.app.config == "cubeviz":
-                sv = self.app.get_viewer(
-                    self.app._jdaviz_helper._default_spectrum_viewer_reference_name)
-                if len(sv.state.layers) == 1:
-                    enforce_flux_unit = self.app._get_display_unit('spectral_y')
             for viewer_ref in add_to_viewer_refs:
                 self.app.remove_data_from_viewer(viewer_ref, label)
             self.app.data_collection.remove(self.app.data_collection[label])
@@ -3947,9 +3941,6 @@ def add_results_from_plugin(self, data_item, replace=None, label=None):
                                         label,
                                         visible=visible, clear_other_data=this_replace)
 
-            if enforce_flux_unit:
-                sv.state.y_display_unit = enforce_flux_unit
-
             if preserved != {}:
                 layer_state = [layer.state for layer in this_viewer.layers if
                                layer.layer.label == label][0]

From 186868af2fa04a08dcd2bdf830e9f0f08ed6c8e5 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Thu, 12 Sep 2024 09:10:49 -0400
Subject: [PATCH 05/22] update slice plugin to handle new messaging order

---
 jdaviz/configs/cubeviz/plugins/mixins.py      | 11 ++++----
 jdaviz/configs/cubeviz/plugins/slice/slice.py | 26 ++++++++++++-------
 2 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/jdaviz/configs/cubeviz/plugins/mixins.py b/jdaviz/configs/cubeviz/plugins/mixins.py
index ea57ab01cb..41ec392cf7 100644
--- a/jdaviz/configs/cubeviz/plugins/mixins.py
+++ b/jdaviz/configs/cubeviz/plugins/mixins.py
@@ -89,6 +89,12 @@ def slice_values(self):
         take_inds = [2, 1, 0]
         take_inds.remove(self.slice_index)
         converted_axis = np.array([])
+
+        # Retrieve display units
+        slice_display_units = self.jdaviz_app._get_display_unit(
+            self.slice_display_unit_name
+        )
+
         for layer in self.layers:
             world_comp_ids = layer.layer.data.world_component_ids
 
@@ -100,11 +106,6 @@ def slice_values(self):
                 # Case where 2D image is loaded in image viewer
                 continue
 
-            # Retrieve display units
-            slice_display_units = self.jdaviz_app._get_display_unit(
-                self.slice_display_unit_name
-            )
-
             try:
                 # Retrieve layer data and units using the slice index of the world components ids
                 data_comp = layer.layer.data.get_component(world_comp_ids[self.slice_index])
diff --git a/jdaviz/configs/cubeviz/plugins/slice/slice.py b/jdaviz/configs/cubeviz/plugins/slice/slice.py
index 0bd4380ca0..a9b1e310a0 100644
--- a/jdaviz/configs/cubeviz/plugins/slice/slice.py
+++ b/jdaviz/configs/cubeviz/plugins/slice/slice.py
@@ -139,12 +139,13 @@ def _initialize_location(self, *args):
             if str(viewer.state.x_att) not in self.valid_slice_att_names:
                 # avoid setting value to degs, before x_att is changed to wavelength, for example
                 continue
-            # ensure the cache is reset (if previous attempts to initialize failed resulting in an
-            # empty list as the cache)
-            viewer._clear_cache('slice_values')
+            if self.app._get_display_unit(viewer.slice_display_unit_name) == '':
+                # viewer is not ready to retrieve slice_values in display units
+                continue
             slice_values = viewer.slice_values
             if len(slice_values):
-                self.value = slice_values[int(len(slice_values)/2)]
+                new_value = slice_values[int(len(slice_values)/2)]
+                self.value = new_value
                 self._indicator_initialized = True
                 return
 
@@ -229,22 +230,29 @@ def _on_select_slice_message(self, msg):
             self.value = msg.value
 
     def _on_global_display_unit_changed(self, msg):
-        if not self.app.config == 'cubeviz':
-            return
-
         if msg.axis != self.slice_display_unit_name:
             return
+        self._clear_cache()
         if not self.value_unit:
             self.value_unit = str(msg.unit)
             return
+        if not self._indicator_initialized:
+            self._initialize_location()
+            return
         prev_unit = u.Unit(self.value_unit)
         self.value_unit = str(msg.unit)
-        self._clear_cache()
-        self.value = (self.value * prev_unit).to_value(msg.unit, equivalencies=u.spectral())
+        self.value = self._convert_value_to_unit(self.value, prev_unit, msg.unit)
+
+    def _convert_value_to_unit(self, value, prev_unit, new_unit):
+        return (value * prev_unit).to_value(new_unit, equivalencies=u.spectral())
 
     def _clear_cache(self, *attrs):
         if not len(attrs):
             attrs = self._cached_properties
+        if len(attrs):
+            # most internally cached properties rely on viewer slice_values, so let's also clear those caches
+            for viewer in self.slice_selection_viewers:
+                viewer._clear_cache('slice_values')
         for attr in attrs:
             if attr in self.__dict__:
                 del self.__dict__[attr]

From 58e9d638ced674f0c7b0739130b2ee76324ad3c6 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Thu, 12 Sep 2024 11:18:45 -0400
Subject: [PATCH 06/22] generalize initializing options to also act on first
 valid data

---
 .../unit_conversion/unit_conversion.py        | 60 +++++++++++--------
 1 file changed, 36 insertions(+), 24 deletions(-)

diff --git a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
index e858e1b537..babf694623 100644
--- a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
+++ b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
@@ -2,6 +2,7 @@
 from glue.core.subset_group import GroupedSubset
 from glue_jupyter.bqplot.image import BqplotImageView
 import numpy as np
+from specutils import Spectrum1D
 from traitlets import List, Unicode, observe, Bool
 
 from jdaviz.configs.default.plugins.viewers import JdavizProfileView
@@ -145,8 +146,7 @@ def __init__(self, *args, **kwargs):
 
         self.spectral_y_type = SelectPluginComponent(self,
                                                      items='spectral_y_type_items',
-                                                     selected='spectral_y_type_selected',
-                                                     manual_options=['Surface Brightness', 'Flux'])
+                                                     selected='spectral_y_type_selected')
 
     @property
     def user_api(self):
@@ -183,43 +183,49 @@ def _on_add_data_to_viewer(self, msg):
 
         viewer = msg.viewer
 
-        # TODO: when enabling unit-conversion in rampviz, this may need to be more specific or handle other cases for ramp profile viewers
-        if isinstance(viewer, JdavizProfileView):
-            if viewer.state.x_display_unit == self.spectral_unit_selected and viewer.state.y_display_unit == self.app._get_display_unit('spectral_y'):
-                # data already existed in this viewer and display units were already set
-                return
-            if not (len(self.spectral_unit_selected)
-                    and len(self.flux_unit_selected)
-                    and len(self.angle_unit_selected)
-                    and (self.config != 'cubeviz' or len(self.spectral_y_type_selected))):
+        if not (len(self.spectral_unit_selected)
+                and len(self.flux_unit_selected)
+                and len(self.angle_unit_selected)
+                and (self.config == 'cubeviz' and not len(self.spectral_y_type_selected))):
 
-                spec = msg.data.get_object()
+            data_obj = msg.data.get_object()
+            if isinstance(data_obj, Spectrum1D):
 
                 self.spectral_unit._addl_unit_strings = self.spectrum_viewer.state.__class__.x_display_unit.get_choices(self.spectrum_viewer.state)
                 if not len(self.spectral_unit_selected):
-                    self.spectral_unit.selected = str(spec.spectral_axis.unit)
+                    try:
+                        self.spectral_unit.selected = str(data_obj.spectral_axis.unit)
+                    except ValueError:
+                        self.spectral_unit.selected = ''
 
-                angle_unit = check_if_unit_is_per_solid_angle(spec.flux.unit, return_unit=True)
-                flux_unit = spec.flux.unit if angle_unit is None else spec.flux.unit * angle_unit
+                angle_unit = check_if_unit_is_per_solid_angle(data_obj.flux.unit, return_unit=True)
+                flux_unit = data_obj.flux.unit if angle_unit is None else data_obj.flux.unit * angle_unit
 
                 if not self.flux_unit_selected:
                     if flux_unit in (u.count, u.DN):
                         self.flux_unit.choices = [flux_unit]
-                    self.flux_unit.selected = str(flux_unit)
+                    try:
+                        self.flux_unit.selected = str(flux_unit)
+                    except ValueError:
+                        self.flux_unit.selected = ''
 
                 if not self.angle_unit_selected:
                     if angle_unit == u.pix**2:
                         self.angle_unit.choices = ['pix2']
 
-                    if angle_unit is None:
-                        # default to sr if input spectrum is not in surface brightness units
-                        # TODO: for cubeviz, should we check the cube itself?
-                        self.angle_unit.selected = 'sr'
-                    else:
-                        self.angle_unit.selected = str(angle_unit)
-
-                if not len(self.spectral_y_type_selected):
+                    try:
+                        if angle_unit is None:
+                            # default to sr if input spectrum is not in surface brightness units
+                            # TODO: for cubeviz, should we check the cube itself?
+                            self.angle_unit.selected = 'sr'
+                        else:
+                            self.angle_unit.selected = str(angle_unit)
+                    except ValueError:
+                        self.angle_unit.selected = ''
+
+                if not len(self.spectral_y_type_selected) and isinstance(viewer, JdavizProfileView):
                     # set spectral_y_type_selected to 'Flux' if the y-axis unit is not per solid angle
+                    self.spectral_y_type.choices = ['Surface Brightness', 'Flux']
                     if angle_unit is None:
                         self.spectral_y_type_selected = 'Flux'
                     else:
@@ -229,6 +235,12 @@ def _on_add_data_to_viewer(self, msg):
                 # so return here to avoid setting twice
                 return
 
+        # TODO: when enabling unit-conversion in rampviz, this may need to be more specific or handle other cases for ramp profile viewers
+        if isinstance(viewer, JdavizProfileView):
+            if viewer.state.x_display_unit == self.spectral_unit_selected and viewer.state.y_display_unit == self.app._get_display_unit('spectral_y'):
+                # data already existed in this viewer and display units were already set
+                return
+
             # this spectral viewer was empty (did not have display units set yet), but global selections
             # are available in the plugin, so we'll set them to the viewer here
             viewer.state.x_display_unit = self.spectral_unit_selected

From 9607eb4256e833e46b1b0d023fc824775f44ef55 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Thu, 12 Sep 2024 11:40:55 -0400
Subject: [PATCH 07/22] code cleanup

---
 jdaviz/configs/cubeviz/plugins/slice/slice.py |  3 +-
 .../unit_conversion/unit_conversion.py        | 60 +++++++++++--------
 jdaviz/core/template_mixin.py                 |  2 +-
 3 files changed, 37 insertions(+), 28 deletions(-)

diff --git a/jdaviz/configs/cubeviz/plugins/slice/slice.py b/jdaviz/configs/cubeviz/plugins/slice/slice.py
index a9b1e310a0..92b9e73ac5 100644
--- a/jdaviz/configs/cubeviz/plugins/slice/slice.py
+++ b/jdaviz/configs/cubeviz/plugins/slice/slice.py
@@ -250,7 +250,8 @@ def _clear_cache(self, *attrs):
         if not len(attrs):
             attrs = self._cached_properties
         if len(attrs):
-            # most internally cached properties rely on viewer slice_values, so let's also clear those caches
+            # most internally cached properties rely on
+            # viewer slice_values, so let's also clear those caches
             for viewer in self.slice_selection_viewers:
                 viewer._clear_cache('slice_values')
         for attr in attrs:
diff --git a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
index babf694623..4a287d4123 100644
--- a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
+++ b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
@@ -1,7 +1,6 @@
 from astropy import units as u
 from glue.core.subset_group import GroupedSubset
 from glue_jupyter.bqplot.image import BqplotImageView
-import numpy as np
 from specutils import Spectrum1D
 from traitlets import List, Unicode, observe, Bool
 
@@ -121,7 +120,6 @@ def __init__(self, *args, **kwargs):
                                                        selected='spectral_unit_selected')
         self.spectral_unit.choices = create_spectral_equivalencies_list(u.Hz)
 
-
         self.has_flux = self.config in ('specviz', 'cubeviz', 'specviz2d', 'mosviz')
         self.flux_unit = UnitSelectPluginComponent(self,
                                                    items='flux_unit_items',
@@ -168,7 +166,8 @@ def user_api(self):
 
     @property
     def sb_unit(self):
-        # expose selected surface-brightness unit as read-only (rather than exposing a select object)
+        # expose selected surface-brightness unit as read-only
+        # (rather than exposing a select object)
         return self.sb_unit_selected
 
     def _on_add_data_to_viewer(self, msg):
@@ -191,7 +190,7 @@ def _on_add_data_to_viewer(self, msg):
             data_obj = msg.data.get_object()
             if isinstance(data_obj, Spectrum1D):
 
-                self.spectral_unit._addl_unit_strings = self.spectrum_viewer.state.__class__.x_display_unit.get_choices(self.spectrum_viewer.state)
+                self.spectral_unit._addl_unit_strings = self.spectrum_viewer.state.__class__.x_display_unit.get_choices(self.spectrum_viewer.state)  # noqa
                 if not len(self.spectral_unit_selected):
                     try:
                         self.spectral_unit.selected = str(data_obj.spectral_axis.unit)
@@ -199,7 +198,7 @@ def _on_add_data_to_viewer(self, msg):
                         self.spectral_unit.selected = ''
 
                 angle_unit = check_if_unit_is_per_solid_angle(data_obj.flux.unit, return_unit=True)
-                flux_unit = data_obj.flux.unit if angle_unit is None else data_obj.flux.unit * angle_unit
+                flux_unit = data_obj.flux.unit if angle_unit is None else data_obj.flux.unit * angle_unit  # noqa
 
                 if not self.flux_unit_selected:
                     if flux_unit in (u.count, u.DN):
@@ -223,26 +222,31 @@ def _on_add_data_to_viewer(self, msg):
                     except ValueError:
                         self.angle_unit.selected = ''
 
-                if not len(self.spectral_y_type_selected) and isinstance(viewer, JdavizProfileView):
-                    # set spectral_y_type_selected to 'Flux' if the y-axis unit is not per solid angle
+                if (not len(self.spectral_y_type_selected)
+                        and isinstance(viewer, JdavizProfileView)):
+                    # set spectral_y_type_selected to 'Flux'
+                    # if the y-axis unit is not per solid angle
                     self.spectral_y_type.choices = ['Surface Brightness', 'Flux']
                     if angle_unit is None:
                         self.spectral_y_type_selected = 'Flux'
                     else:
                         self.spectral_y_type_selected = 'Surface Brightness'
 
-                # setting default values will trigger the observes to set the units in _on_unit_selected,
-                # so return here to avoid setting twice
+                # setting default values will trigger the observes to set the units
+                # in _on_unit_selected, so return here to avoid setting twice
                 return
 
-        # TODO: when enabling unit-conversion in rampviz, this may need to be more specific or handle other cases for ramp profile viewers
+        # TODO: when enabling unit-conversion in rampviz, this may need to be more specific
+        # or handle other cases for ramp profile viewers
         if isinstance(viewer, JdavizProfileView):
-            if viewer.state.x_display_unit == self.spectral_unit_selected and viewer.state.y_display_unit == self.app._get_display_unit('spectral_y'):
+            if (viewer.state.x_display_unit == self.spectral_unit_selected
+                    and viewer.state.y_display_unit == self.app._get_display_unit('spectral_y')):
                 # data already existed in this viewer and display units were already set
                 return
 
-            # this spectral viewer was empty (did not have display units set yet), but global selections
-            # are available in the plugin, so we'll set them to the viewer here
+            # this spectral viewer was empty (did not have display units set yet),
+            # but global selections are available in the plugin,
+            # so we'll set them to the viewer here
             viewer.state.x_display_unit = self.spectral_unit_selected
             # _handle_spectral_y_unit will call viewer.set_plot_axes()
             self._handle_spectral_y_unit()
@@ -253,7 +257,6 @@ def _on_add_data_to_viewer(self, msg):
             layers = [lyr for lyr in msg.viewer.layers if lyr.layer.data.label == msg.data.label]
             self._handle_attribute_display_unit(self.sb_unit_selected, layers=layers)
 
-
     @observe('spectral_unit_selected', 'flux_unit_selected',
              'angle_unit_selected', 'sb_unit_selected',
              'time_unit_selected')
@@ -276,16 +279,20 @@ def _on_unit_selected(self, msg):
         elif axis == 'flux':
             if len(self.angle_unit_selected):
                 # NOTE: setting sb_unit_selected will call this method again with axis=='sb',
-                # which in turn will call _handle_spectral_y_unit and send a second GlobalDisplayUnitChanged message for sb
-                self.sb_unit_selected = _flux_to_sb_unit(self.flux_unit.selected, self.angle_unit.selected)
+                # which in turn will call _handle_spectral_y_unit and
+                # send a second GlobalDisplayUnitChanged message for sb
+                self.sb_unit_selected = _flux_to_sb_unit(self.flux_unit.selected,
+                                                         self.angle_unit.selected)
 
             if self.spectral_y_type_selected == 'Flux':
                 self._handle_spectral_y_unit()
 
         elif axis == 'angle':
             if len(self.flux_unit_selected):
-                # NOTE: setting sb_unit_selected will call this method again and send a second GlobalDisplayUnitChanged message for sb
-                self.sb_unit_selected = _flux_to_sb_unit(self.flux_unit.selected, self.angle_unit.selected)
+                # NOTE: setting sb_unit_selected will call this method again and
+                # send a second GlobalDisplayUnitChanged message for sb
+                self.sb_unit_selected = _flux_to_sb_unit(self.flux_unit.selected,
+                                                         self.angle_unit.selected)
 
         elif axis == 'sb':
             self._handle_attribute_display_unit(self.sb_unit_selected)
@@ -300,14 +307,15 @@ def _on_unit_selected(self, msg):
         self.hub.broadcast(GlobalDisplayUnitChanged(msg.name.split('_')[0],
                            msg.new, sender=self))
 
-
     @observe('spectral_y_type_selected')
     def _handle_spectral_y_unit(self, *args):
         """
-        When the spectral_y_type is changed, or the unit corresponding to the currently selected spectral_y_type is changed,
-        update the y-axis of the spectrum viewer with the new unit, and then emit a GlobalDisplayUnitChanged message to notify
+        When the spectral_y_type is changed, or the unit corresponding to the
+        currently selected spectral_y_type is changed, update the y-axis of
+        the spectrum viewer with the new unit, and then emit a
+        GlobalDisplayUnitChanged message to notify
         """
-        yunit_selected = self.sb_unit_selected if self.spectral_y_type_selected == 'Surface Brightness' else self.flux_unit_selected
+        yunit_selected = self.sb_unit_selected if self.spectral_y_type_selected == 'Surface Brightness' else self.flux_unit_selected  # noqa
         yunit = _valid_glue_display_unit(yunit_selected, self.spectrum_viewer, 'y')
         if self.spectrum_viewer.state.y_display_unit == yunit:
             return
@@ -324,16 +332,16 @@ def _handle_spectral_y_unit(self, *args):
                                                         yunit,
                                                         sender=self))
 
- 
     def _handle_attribute_display_unit(self, attr_unit, layers=None):
         """
-        Update the per-layer attribute display unit in glue for image viewers (updating stretch and contour units).
+        Update the per-layer attribute display unit in glue for image viewers
+        (updating stretch and contour units).
         """
         if layers is None:
             layers = [layer
-                      for viewer in self._app._viewer_store.values() if isinstance(viewer, BqplotImageView)
+                      for viewer in self._app._viewer_store.values() if isinstance(viewer, BqplotImageView)  # noqa
                       for layer in viewer.layers]
- 
+
         for layer in layers:
             # DQ layer doesn't play nicely with this attribute
             if "DQ" in layer.layer.label or isinstance(layer.layer, GroupedSubset):
diff --git a/jdaviz/core/template_mixin.py b/jdaviz/core/template_mixin.py
index 0deba57075..0eb455691c 100644
--- a/jdaviz/core/template_mixin.py
+++ b/jdaviz/core/template_mixin.py
@@ -45,7 +45,7 @@
 from jdaviz.core.events import (AddDataMessage, RemoveDataMessage,
                                 ViewerAddedMessage, ViewerRemovedMessage,
                                 ViewerRenamedMessage, SnackbarMessage,
-                                AddDataToViewerMessage, ChangeRefDataMessage,
+                                ChangeRefDataMessage,
                                 PluginTableAddedMessage, PluginTableModifiedMessage,
                                 PluginPlotAddedMessage, PluginPlotModifiedMessage,
                                 GlobalDisplayUnitChanged)

From cc43f459abd81b85d12b757730d5cf2aaf8cf5bf Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Fri, 20 Sep 2024 15:51:56 -0400
Subject: [PATCH 08/22] fix test failures

---
 jdaviz/app.py                                 |  3 +-
 .../tests/test_spectral_extraction.py         |  1 +
 jdaviz/configs/default/plugins/viewers.py     |  7 ++-
 .../tests/test_unit_conversion.py             | 10 ++++-
 .../unit_conversion/unit_conversion.py        | 45 ++++++++++++-------
 5 files changed, 44 insertions(+), 22 deletions(-)

diff --git a/jdaviz/app.py b/jdaviz/app.py
index 17c0981c2d..895a0e90b6 100644
--- a/jdaviz/app.py
+++ b/jdaviz/app.py
@@ -1345,8 +1345,7 @@ def _get_display_unit(self, axis):
                 raise ValueError(f"could not find units for axis='{axis}'")
         uc = self._jdaviz_helper.plugins.get('Unit Conversion')._obj
         if axis == 'spectral_y':
-            # translate options from uc.spectral_y_type to the prefix used in uc.??_unit_selected
-            axis = {'Surface Brightness': 'sb', 'Flux': 'flux'}[uc.spectral_y_type_selected]
+            return uc.spectral_y_unit
         try:
             return getattr(uc, f'{axis}_unit_selected')
         except AttributeError:
diff --git a/jdaviz/configs/cubeviz/plugins/spectral_extraction/tests/test_spectral_extraction.py b/jdaviz/configs/cubeviz/plugins/spectral_extraction/tests/test_spectral_extraction.py
index 07e75f884a..6e855decfd 100644
--- a/jdaviz/configs/cubeviz/plugins/spectral_extraction/tests/test_spectral_extraction.py
+++ b/jdaviz/configs/cubeviz/plugins/spectral_extraction/tests/test_spectral_extraction.py
@@ -584,6 +584,7 @@ def test_spectral_extraction_unit_conv_one_spec(
     uc = cubeviz_helper.plugins["Unit Conversion"]
     assert uc.flux_unit == "Jy"
     uc.flux_unit.selected = "MJy"
+    assert spectrum_viewer.state.y_display_unit == "MJy"
     spec_extr_plugin = cubeviz_helper.plugins['Spectral Extraction']
     # Overwrite the one and only default extraction.
     collapsed = spec_extr_plugin.extract()
diff --git a/jdaviz/configs/default/plugins/viewers.py b/jdaviz/configs/default/plugins/viewers.py
index ff9077542d..48e55d9a7f 100644
--- a/jdaviz/configs/default/plugins/viewers.py
+++ b/jdaviz/configs/default/plugins/viewers.py
@@ -770,8 +770,11 @@ def set_plot_axes(self):
 
         if solid_angle_unit is not None:
 
-            for un in locally_defined_flux_units:
-                locally_defined_sb_unit = un / solid_angle_unit
+            # create an equivalency for each flux unit for flux <> flux/pix2.
+            # for similar reasons to the 'untranslatable units' issue, custom
+            # equivs. can't be combined, so a workaround is creating an equiv
+            # for each flux that may need an additional equiv.
+            angle_to_pixel_equivs = [_eqv_sb_per_pixel_to_per_angle(un) for un in locally_defined_flux_units]  # noqa
 
                 # create an equivalency for each flux unit for flux <> flux/pix2.
                 # for similar reasons to the 'untranslatable units' issue, custom
diff --git a/jdaviz/configs/specviz/plugins/unit_conversion/tests/test_unit_conversion.py b/jdaviz/configs/specviz/plugins/unit_conversion/tests/test_unit_conversion.py
index b967e90ffe..51fc06537b 100644
--- a/jdaviz/configs/specviz/plugins/unit_conversion/tests/test_unit_conversion.py
+++ b/jdaviz/configs/specviz/plugins/unit_conversion/tests/test_unit_conversion.py
@@ -100,7 +100,7 @@ def test_conv_no_data(specviz_helper, spectrum1d):
     # spectrum not load is in Flux units, sb_unit and flux_unit
     # should be enabled, spectral_y_type should not be
     plg = specviz_helper.plugins["Unit Conversion"]
-    with pytest.raises(ValueError, match="no valid unit choices"):
+    with pytest.raises(ValueError, match="could not find match in valid x display units"):
         plg.spectral_unit = "micron"
     assert len(specviz_helper.app.data_collection) == 0
 
@@ -290,11 +290,17 @@ def test_contour_unit_conversion(cubeviz_helper, spectrum1d_cube_fluxunit_jy_per
     # Make sure that the contour values get updated
     po_plg.contour_visible = True
 
+    assert uc_plg.spectral_y_type == 'Flux'
+    assert uc_plg.flux_unit == 'Jy'
+    assert uc_plg.sb_unit == "Jy / sr"
+    assert cubeviz_helper.viewers['flux-viewer']._obj.layers[0].state.attribute_display_unit == "Jy / sr"  # noqa
     assert np.allclose(po_plg.contour_max.value, 199)
 
-    uc_plg._obj.spectral_y_type_selected = 'Surface Brightness'
+    uc_plg.spectral_y_type = 'Surface Brightness'
     uc_plg.flux_unit = 'MJy'
 
+    assert uc_plg.sb_unit == "MJy / sr"
+    assert cubeviz_helper.viewers['flux-viewer']._obj.layers[0].state.attribute_display_unit == "MJy / sr"  # noqa
     assert np.allclose(po_plg.contour_max.value, 1.99e-4)
 
 
diff --git a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
index 4a287d4123..562dab80c3 100644
--- a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
+++ b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
@@ -69,6 +69,9 @@ class UnitConversion(PluginTemplateMixin):
       derived from the set values of ``flux_unit`` and ``angle_unit``.
     * ``spectral_y_type`` (:class:`~jdaviz.core.template_mixin.SelectPluginComponent`):
       Select the y-axis physical type for the spectrum-viewer (applicable only to Cubeviz).
+    * ``spectral_y_unit``: Read-only property for the current y-axis unit in the spectrum-viewer,
+      either ``flux_unit`` or ``sb_unit`` depending on the selected ``spectral_y_type``
+      (applicable only to Cubeviz).
     """
     template_file = __file__, "unit_conversion.vue"
 
@@ -161,7 +164,7 @@ def user_api(self):
         if self.has_time:
             expose += ['time_unit']
         if self.config == 'cubeviz':
-            expose += ['spectral_y_type']
+            expose += ['spectral_y_type', 'spectral_y_unit']
         return PluginUserApi(self, expose=expose, readonly=readonly)
 
     @property
@@ -170,6 +173,10 @@ def sb_unit(self):
         # (rather than exposing a select object)
         return self.sb_unit_selected
 
+    @property
+    def spectral_y_unit(self):
+        return self.sb_unit_selected if self.spectral_y_type_selected == 'Surface Brightness' else self.flux_unit_selected  # noqa
+
     def _on_add_data_to_viewer(self, msg):
         # toggle warning message for cubes without PIXAR_SR defined
         if self.config == 'cubeviz':
@@ -182,10 +189,10 @@ def _on_add_data_to_viewer(self, msg):
 
         viewer = msg.viewer
 
-        if not (len(self.spectral_unit_selected)
-                and len(self.flux_unit_selected)
-                and len(self.angle_unit_selected)
-                and (self.config == 'cubeviz' and not len(self.spectral_y_type_selected))):
+        if (not len(self.spectral_unit_selected)
+                or not len(self.flux_unit_selected)
+                or not len(self.angle_unit_selected)
+                or (self.config == 'cubeviz' and not len(self.spectral_y_type_selected))):
 
             data_obj = msg.data.get_object()
             if isinstance(data_obj, Spectrum1D):
@@ -240,11 +247,11 @@ def _on_add_data_to_viewer(self, msg):
         # or handle other cases for ramp profile viewers
         if isinstance(viewer, JdavizProfileView):
             if (viewer.state.x_display_unit == self.spectral_unit_selected
-                    and viewer.state.y_display_unit == self.app._get_display_unit('spectral_y')):
+                    and viewer.state.y_display_unit == self.spectral_y_unit):
                 # data already existed in this viewer and display units were already set
                 return
 
-            # this spectral viewer was empty (did not have display units set yet),
+            # this spectral viewer was empty (did not have display units set yet),˜
             # but global selections are available in the plugin,
             # so we'll set them to the viewer here
             viewer.state.x_display_unit = self.spectral_unit_selected
@@ -279,8 +286,9 @@ def _on_unit_selected(self, msg):
         elif axis == 'flux':
             if len(self.angle_unit_selected):
                 # NOTE: setting sb_unit_selected will call this method again with axis=='sb',
-                # which in turn will call _handle_spectral_y_unit and
-                # send a second GlobalDisplayUnitChanged message for sb
+                # which in turn will call _handle_attribute_display_unit,
+                # _handle_spectral_y_unit (if spectral_y_type_selected == 'Surface Brightness'),
+                #  and send a second GlobalDisplayUnitChanged message for sb
                 self.sb_unit_selected = _flux_to_sb_unit(self.flux_unit.selected,
                                                          self.angle_unit.selected)
 
@@ -289,8 +297,10 @@ def _on_unit_selected(self, msg):
 
         elif axis == 'angle':
             if len(self.flux_unit_selected):
-                # NOTE: setting sb_unit_selected will call this method again and
-                # send a second GlobalDisplayUnitChanged message for sb
+                # NOTE: setting sb_unit_selected will call this method again with axis=='sb',
+                # which in turn will call _handle_attribute_display_unit,
+                # _handle_spectral_y_unit (if spectral_y_type_selected == 'Surface Brightness'),
+                #  and send a second GlobalDisplayUnitChanged message for sb
                 self.sb_unit_selected = _flux_to_sb_unit(self.flux_unit.selected,
                                                          self.angle_unit.selected)
 
@@ -315,9 +325,9 @@ def _handle_spectral_y_unit(self, *args):
         the spectrum viewer with the new unit, and then emit a
         GlobalDisplayUnitChanged message to notify
         """
-        yunit_selected = self.sb_unit_selected if self.spectral_y_type_selected == 'Surface Brightness' else self.flux_unit_selected  # noqa
-        yunit = _valid_glue_display_unit(yunit_selected, self.spectrum_viewer, 'y')
+        yunit = _valid_glue_display_unit(self.spectral_y_unit, self.spectrum_viewer, 'y')
         if self.spectrum_viewer.state.y_display_unit == yunit:
+            self.spectrum_viewer.set_plot_axes()
             return
         try:
             self.spectrum_viewer.state.y_display_unit = yunit
@@ -346,7 +356,10 @@ def _handle_attribute_display_unit(self, attr_unit, layers=None):
             # DQ layer doesn't play nicely with this attribute
             if "DQ" in layer.layer.label or isinstance(layer.layer, GroupedSubset):
                 continue
-            elif u.Unit(layer.layer.get_component("flux").units).physical_type != 'surface brightness':  # noqa
+            elif ("flux" not in [str(c) for c in layer.layer.components]
+                    or u.Unit(layer.layer.get_component("flux").units).physical_type != 'surface brightness'):  # noqa
                 continue
-            if hasattr(layer, 'attribute_display_unit'):
-                layer.attribute_display_unit = attr_unit
+            if hasattr(layer.state, 'attribute_display_unit'):
+                layer.state.attribute_display_unit = _valid_glue_display_unit(attr_unit,
+                                                                              layer,
+                                                                              'attribute')

From 72b4677157d1794336d13fd57b0ce82f7bb8cd30 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Wed, 18 Sep 2024 11:28:22 -0400
Subject: [PATCH 09/22] Revert "skip test in parsers"

This reverts commit aa0d3848672bffe0ebaffa53ee436b16f67c21ce.
---
 jdaviz/configs/cubeviz/plugins/tests/test_parsers.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/jdaviz/configs/cubeviz/plugins/tests/test_parsers.py b/jdaviz/configs/cubeviz/plugins/tests/test_parsers.py
index 912bf47ae8..720959202a 100644
--- a/jdaviz/configs/cubeviz/plugins/tests/test_parsers.py
+++ b/jdaviz/configs/cubeviz/plugins/tests/test_parsers.py
@@ -160,7 +160,6 @@ def test_spectrum3d_no_wcs_parse(cubeviz_helper):
     assert flux.units == 'nJy / pix2'
 
 
-@pytest.mark.skip(reason="unskip after 3192 merged")
 def test_spectrum1d_parse(spectrum1d, cubeviz_helper):
     cubeviz_helper.load_data(spectrum1d)
 

From 08cb09c01515d1b35a47eb8f57eed6ef29000ad1 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Mon, 23 Sep 2024 07:37:42 -0400
Subject: [PATCH 10/22] fix rebase failure

---
 jdaviz/configs/default/plugins/viewers.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/jdaviz/configs/default/plugins/viewers.py b/jdaviz/configs/default/plugins/viewers.py
index 48e55d9a7f..ff9077542d 100644
--- a/jdaviz/configs/default/plugins/viewers.py
+++ b/jdaviz/configs/default/plugins/viewers.py
@@ -770,11 +770,8 @@ def set_plot_axes(self):
 
         if solid_angle_unit is not None:
 
-            # create an equivalency for each flux unit for flux <> flux/pix2.
-            # for similar reasons to the 'untranslatable units' issue, custom
-            # equivs. can't be combined, so a workaround is creating an equiv
-            # for each flux that may need an additional equiv.
-            angle_to_pixel_equivs = [_eqv_sb_per_pixel_to_per_angle(un) for un in locally_defined_flux_units]  # noqa
+            for un in locally_defined_flux_units:
+                locally_defined_sb_unit = un / solid_angle_unit
 
                 # create an equivalency for each flux unit for flux <> flux/pix2.
                 # for similar reasons to the 'untranslatable units' issue, custom

From 029955c75820beec7b2d8ad6ba68b520bf5b7c92 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Mon, 23 Sep 2024 07:38:24 -0400
Subject: [PATCH 11/22] unskip test

---
 jdaviz/configs/specviz/tests/test_viewers.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/jdaviz/configs/specviz/tests/test_viewers.py b/jdaviz/configs/specviz/tests/test_viewers.py
index e8fab5f95b..a64ba989a1 100644
--- a/jdaviz/configs/specviz/tests/test_viewers.py
+++ b/jdaviz/configs/specviz/tests/test_viewers.py
@@ -4,7 +4,6 @@
 from specutils import Spectrum1D
 
 
-@pytest.mark.skip(reason="unskip after 3192 merged")
 @pytest.mark.parametrize(
     ('input_unit', 'y_axis_label'),
     [(u.MJy, 'Flux'),

From 8fb3d5d6913d79c4050559a9cd7fb74ea35a6925 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Mon, 23 Sep 2024 07:47:19 -0400
Subject: [PATCH 12/22] replace time axis elif with comment

---
 .../specviz/plugins/unit_conversion/unit_conversion.py      | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
index 562dab80c3..f684c00c18 100644
--- a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
+++ b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
@@ -310,11 +310,11 @@ def _on_unit_selected(self, msg):
             if self.spectral_y_type_selected == 'Surface Brightness':
                 self._handle_spectral_y_unit()
 
-        elif axis == 'time':
-            pass
+        # custom axes downstream can override _on_unit_selected if anything needs to be
+        # processed before the GlobalDisplayUnitChanged message is broadcast
 
         # axis (first) argument will be one of: spectral, flux, angle, sb, time
-        self.hub.broadcast(GlobalDisplayUnitChanged(msg.name.split('_')[0],
+        self.hub.broadcast(GlobalDisplayUnitChanged(axis,
                            msg.new, sender=self))
 
     @observe('spectral_y_type_selected')

From 7a2fef7b539d62d4fc58c62a39dd1506a02976b2 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Mon, 23 Sep 2024 07:51:26 -0400
Subject: [PATCH 13/22] add e/s

---
 .../configs/specviz/plugins/unit_conversion/unit_conversion.py  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
index f684c00c18..5fb7bfd48b 100644
--- a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
+++ b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
@@ -208,7 +208,7 @@ def _on_add_data_to_viewer(self, msg):
                 flux_unit = data_obj.flux.unit if angle_unit is None else data_obj.flux.unit * angle_unit  # noqa
 
                 if not self.flux_unit_selected:
-                    if flux_unit in (u.count, u.DN):
+                    if flux_unit in (u.count, u.DN, u.electron / u.s):
                         self.flux_unit.choices = [flux_unit]
                     try:
                         self.flux_unit.selected = str(flux_unit)

From 51e2f702d2320e8e2c16b702336e3ff8aad9fb95 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Mon, 23 Sep 2024 07:56:05 -0400
Subject: [PATCH 14/22] changelog entry

---
 CHANGES.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGES.rst b/CHANGES.rst
index 74ad53eab5..4f8438f6ed 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -6,7 +6,7 @@ New Features
 
 - Added flux/surface brightness translation and surface brightness
   unit conversion in Cubeviz and Specviz. [#2781, #2940, #3088, #3111, #3113, #3129,
-  #3139, #3149, #3155, #3178, #3185, #3187, #3190, #3156, #3200]
+  #3139, #3149, #3155, #3178, #3185, #3187, #3190, #3156, #3200, #3192]
 
 - Plugin tray is now open by default. [#2892]
 

From 08611aa2ac779ddcb2a93dc0ecfc71d06ec8810d Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Mon, 23 Sep 2024 07:57:53 -0400
Subject: [PATCH 15/22] revert debugging separation

---
 jdaviz/configs/cubeviz/plugins/slice/slice.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/jdaviz/configs/cubeviz/plugins/slice/slice.py b/jdaviz/configs/cubeviz/plugins/slice/slice.py
index 92b9e73ac5..6e50e3bdd3 100644
--- a/jdaviz/configs/cubeviz/plugins/slice/slice.py
+++ b/jdaviz/configs/cubeviz/plugins/slice/slice.py
@@ -144,8 +144,7 @@ def _initialize_location(self, *args):
                 continue
             slice_values = viewer.slice_values
             if len(slice_values):
-                new_value = slice_values[int(len(slice_values)/2)]
-                self.value = new_value
+                self.value = slice_values[int(len(slice_values)/2)]
                 self._indicator_initialized = True
                 return
 

From 4fc54598c4a804d130f50d916b557c5c79fbc682 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Tue, 24 Sep 2024 11:17:11 -0400
Subject: [PATCH 16/22] fix ylabel test failures

---
 jdaviz/configs/default/plugins/viewers.py | 48 +++++++++++------------
 1 file changed, 23 insertions(+), 25 deletions(-)

diff --git a/jdaviz/configs/default/plugins/viewers.py b/jdaviz/configs/default/plugins/viewers.py
index ff9077542d..aab17b8e5d 100644
--- a/jdaviz/configs/default/plugins/viewers.py
+++ b/jdaviz/configs/default/plugins/viewers.py
@@ -768,31 +768,29 @@ def set_plot_axes(self):
         # default to the catchall 'flux density' label
         flux_unit_type = None
 
-        if solid_angle_unit is not None:
-
-            for un in locally_defined_flux_units:
-                locally_defined_sb_unit = un / solid_angle_unit
-
-                # create an equivalency for each flux unit for flux <> flux/pix2.
-                # for similar reasons to the 'untranslatable units' issue, custom
-                # equivs. can't be combined, so a workaround is creating an eqiv
-                # for each flux that may need an additional equiv.
-                angle_to_pixel_equiv = _eqv_sb_per_pixel_to_per_angle(un)
-
-                if y_unit.is_equivalent(locally_defined_sb_unit, angle_to_pixel_equiv):
-                    flux_unit_type = "Surface Brightness"
-                elif y_unit.is_equivalent(un):
-                    flux_unit_type = 'Flux'
-                elif y_unit.is_equivalent(u.electron / u.s) or y_unit.physical_type == 'dimensionless':  # noqa
-                    # electron / s or 'dimensionless_unscaled' should be labeled counts
-                    flux_unit_type = "Counts"
-                elif y_unit.is_equivalent(u.W):
-                    flux_unit_type = "Luminosity"
-                if flux_unit_type is not None:
-                    # if we determined a label, stop checking
-                    break
-
-        if flux_unit_type is None:
+        for un in locally_defined_flux_units:
+            locally_defined_sb_unit = un / solid_angle_unit if solid_angle_unit is not None else None  # noqa
+
+            # create an equivalency for each flux unit for flux <> flux/pix2.
+            # for similar reasons to the 'untranslatable units' issue, custom
+            # equivs. can't be combined, so a workaround is creating an eqiv
+            # for each flux that may need an additional equiv.
+            angle_to_pixel_equiv = _eqv_sb_per_pixel_to_per_angle(un)
+
+            if (locally_defined_sb_unit is not None
+                    and y_unit.is_equivalent(locally_defined_sb_unit, angle_to_pixel_equiv)):
+                flux_unit_type = "Surface Brightness"
+            elif y_unit.is_equivalent(un):
+                flux_unit_type = 'Flux'
+            elif y_unit.is_equivalent(u.electron / u.s) or y_unit.physical_type == 'dimensionless':  # noqa
+                # electron / s or 'dimensionless_unscaled' should be labeled counts
+                flux_unit_type = "Counts"
+            elif y_unit.is_equivalent(u.W):
+                flux_unit_type = "Luminosity"
+            if flux_unit_type is not None:
+                # if we determined a label, stop checking
+                break
+        else:
             # default to Flux Density for flux density or uncaught types
             flux_unit_type = "Flux density"
 

From 5ea0327e825b8222328ce23ee1a663ea84530c20 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Tue, 24 Sep 2024 11:33:22 -0400
Subject: [PATCH 17/22] code simplification

---
 jdaviz/utils.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/jdaviz/utils.py b/jdaviz/utils.py
index a1e7e93309..fb0997872b 100644
--- a/jdaviz/utils.py
+++ b/jdaviz/utils.py
@@ -550,11 +550,8 @@ def _eqv_flux_to_sb_pixel():
     flux_units = [u.MJy, u.erg / (u.s * u.cm**2 * u.Angstrom),
                   u.ph / (u.Angstrom * u.s * u.cm**2),
                   u.ph / (u.Hz * u.s * u.cm**2)]
-    equiv = []
-    for flux_unit in flux_units:
-        equiv.append((flux_unit, flux_unit / pix2, lambda x: x, lambda x: x))
-
-    return equiv
+    return [(flux_unit, flux_unit / pix2, lambda x: x, lambda x: x)
+            for flux_unit in flux_units]
 
 
 def _eqv_sb_per_pixel_to_per_angle(flux_unit, scale_factor=1):

From 54c65dfb5a3264d0e296faeb348aea642da30512 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Tue, 24 Sep 2024 13:33:09 -0400
Subject: [PATCH 18/22] fix model_fitting test failures

---
 .../configs/default/plugins/model_fitting/model_fitting.py  | 6 +++---
 .../default/plugins/model_fitting/tests/test_fitting.py     | 4 ++++
 .../specviz/plugins/unit_conversion/unit_conversion.py      | 3 +++
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/jdaviz/configs/default/plugins/model_fitting/model_fitting.py b/jdaviz/configs/default/plugins/model_fitting/model_fitting.py
index 4213baca47..aae4127ca4 100644
--- a/jdaviz/configs/default/plugins/model_fitting/model_fitting.py
+++ b/jdaviz/configs/default/plugins/model_fitting/model_fitting.py
@@ -472,9 +472,9 @@ def _initialize_model_component(self, model_comp, comp_label, poly_order=None):
 
         # Need to set the units the first time we initialize a model component, after this
         # we listen for display unit changes
-        if (self._units is None or self._units == {} or 'x' not in self._units or
-                'y' not in self._units):
+        if self._units.get('x', '') == '':
             self._units['x'] = self.app._get_display_unit('spectral')
+        if self._units.get('y', '') == '':
             if self.cube_fit:
                 self._units['y'] = self.app._get_display_unit('sb')
             else:
@@ -537,7 +537,7 @@ def _initialize_model_component(self, model_comp, comp_label, poly_order=None):
 
         # equivs for spectral density and flux<>flux/pix2. revisit
         # when generalizing plugin UC equivs.
-        equivs = _eqv_flux_to_sb_pixel() + [u.spectral_density(init_x)]
+        equivs = _eqv_flux_to_sb_pixel() + u.spectral_density(init_x)
         init_y = init_y.to(self._units['y'], equivs)
 
         initialized_model = initialize(
diff --git a/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py b/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py
index 9e701412ae..69ed0251a5 100644
--- a/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py
+++ b/jdaviz/configs/default/plugins/model_fitting/tests/test_fitting.py
@@ -93,6 +93,10 @@ def test_parameter_retrieval(cubeviz_helper, spectral_cube_wcs):
     # even though the spectral y axis is in 'flux' by default
     plugin.cube_fit = True
 
+    assert cubeviz_helper.app._get_display_unit('spectral') == wav_unit
+    assert cubeviz_helper.app._get_display_unit('spectral_y') == flux_unit
+    assert cubeviz_helper.app._get_display_unit('sb') == sb_unit
+
     plugin.create_model_component("Linear1D", "L")
     with warnings.catch_warnings():
         warnings.filterwarnings('ignore', message='Model is linear in parameters.*')
diff --git a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
index 5fb7bfd48b..58abc018b4 100644
--- a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
+++ b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
@@ -210,6 +210,9 @@ def _on_add_data_to_viewer(self, msg):
                 if not self.flux_unit_selected:
                     if flux_unit in (u.count, u.DN, u.electron / u.s):
                         self.flux_unit.choices = [flux_unit]
+                    elif flux_unit not in self.flux_unit.choices:
+                        # ensure that the native units are in the list of choices
+                        self.flux_unit.choices += [flux_unit]
                     try:
                         self.flux_unit.selected = str(flux_unit)
                     except ValueError:

From 721ac054612dcee3545fde36e37db6550c2afe15 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Tue, 24 Sep 2024 13:56:15 -0400
Subject: [PATCH 19/22] ignore multiple slashes warning

---
 jdaviz/configs/specviz/tests/test_viewers.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/jdaviz/configs/specviz/tests/test_viewers.py b/jdaviz/configs/specviz/tests/test_viewers.py
index a64ba989a1..28981fbaea 100644
--- a/jdaviz/configs/specviz/tests/test_viewers.py
+++ b/jdaviz/configs/specviz/tests/test_viewers.py
@@ -1,6 +1,7 @@
 import astropy.units as u
 import numpy as np
 import pytest
+import warnings
 from specutils import Spectrum1D
 
 
@@ -19,7 +20,9 @@ def test_spectrum_viewer_axis_labels(specviz_helper, input_unit, y_axis_label):
 
     spec = Spectrum1D(flux, spectral_axis)
 
-    specviz_helper.load_data(spec)
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", message=".*contains multiple slashes, which is discouraged by the FITS standard.*")  # noqa
+        specviz_helper.load_data(spec)
 
     label = specviz_helper.app.get_viewer_by_id('specviz-0').figure.axes[1].label
 

From b356b594381ec5498969a4c003377c5c48fd2cf3 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Wed, 25 Sep 2024 11:04:41 -0400
Subject: [PATCH 20/22] fix uncertainty unit conversion

---
 jdaviz/configs/cubeviz/plugins/parsers.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/jdaviz/configs/cubeviz/plugins/parsers.py b/jdaviz/configs/cubeviz/plugins/parsers.py
index ae2113ab63..40b0236c2f 100644
--- a/jdaviz/configs/cubeviz/plugins/parsers.py
+++ b/jdaviz/configs/cubeviz/plugins/parsers.py
@@ -611,8 +611,9 @@ def convert_spectrum1d_from_flux_to_flux_per_pixel(spectrum):
     # and uncerts, if present
     uncerts = getattr(spectrum, 'uncertainty')
     if uncerts is not None:
-        old_uncerts = uncerts.represent_as(StdDevUncertainty)  # enforce common uncert type.
-        uncerts = old_uncerts.quantity / (u.pix * u.pix)
+        # enforce common uncert type.
+        uncerts = uncerts.represent_as(StdDevUncertainty)
+        uncerts = StdDevUncertainty(uncerts.quantity / (u.pix * u.pix))
 
     # create a new spectrum 1d with all the info from the input spectrum 1d,
     # and the flux / uncerts converted from flux to SB per square pixel

From 59d1192bd104236dddeab9144da14568ad8f40e3 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Wed, 25 Sep 2024 14:47:03 -0400
Subject: [PATCH 21/22] re-introduce resetting y-limits (until can be fixed)

---
 .../configs/specviz/plugins/unit_conversion/unit_conversion.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
index 58abc018b4..a8f0f3828b 100644
--- a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
+++ b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
@@ -339,6 +339,9 @@ def _handle_spectral_y_unit(self, *args):
             pass
         else:
             self.spectrum_viewer.set_plot_axes()
+            # until we can have upstream automatic limit updating on change
+            # in display units with equivalencies, we'll reset the limits
+            self.spectrum_viewer.reset_limits()
 
             # broadcast that there has been a change in the spectrum viewer y axis,
             self.hub.broadcast(GlobalDisplayUnitChanged('spectral_y',

From d467ccb20f6f6e35aab9359cc152f59c269f5eab Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Thu, 26 Sep 2024 09:27:53 -0400
Subject: [PATCH 22/22] optimize changing spectral y-unit

* apply handle_spectral_y_unit before attribute_display_unit
* optimize handle_attribute_display_unit by caching image layers
---
 .../unit_conversion/unit_conversion.py        | 28 +++++++++++++------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
index a8f0f3828b..e7474f14b0 100644
--- a/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
+++ b/jdaviz/configs/specviz/plugins/unit_conversion/unit_conversion.py
@@ -1,4 +1,5 @@
 from astropy import units as u
+from functools import cached_property
 from glue.core.subset_group import GroupedSubset
 from glue_jupyter.bqplot.image import BqplotImageView
 from specutils import Spectrum1D
@@ -105,6 +106,8 @@ class UnitConversion(PluginTemplateMixin):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
+        self._cached_properties = ['image_layers']
+
         if self.config not in ['specviz', 'cubeviz']:
             # TODO [specviz2d, mosviz] x_display_unit is not implemented in glue for image viewer
             # used by spectrum-2d-viewer
@@ -177,6 +180,12 @@ def sb_unit(self):
     def spectral_y_unit(self):
         return self.sb_unit_selected if self.spectral_y_type_selected == 'Surface Brightness' else self.flux_unit_selected  # noqa
 
+    @cached_property
+    def image_layers(self):
+        return [layer
+                for viewer in self._app._viewer_store.values() if isinstance(viewer, BqplotImageView)  # noqa
+                for layer in viewer.layers]
+
     def _on_add_data_to_viewer(self, msg):
         # toggle warning message for cubes without PIXAR_SR defined
         if self.config == 'cubeviz':
@@ -266,6 +275,7 @@ def _on_add_data_to_viewer(self, msg):
             # NOTE: this assumes that all image data is coerced to surface brightness units
             layers = [lyr for lyr in msg.viewer.layers if lyr.layer.data.label == msg.data.label]
             self._handle_attribute_display_unit(self.sb_unit_selected, layers=layers)
+            self._clear_cache('image_layers')
 
     @observe('spectral_unit_selected', 'flux_unit_selected',
              'angle_unit_selected', 'sb_unit_selected',
@@ -287,6 +297,11 @@ def _on_unit_selected(self, msg):
             self.spectrum_viewer.set_plot_axes()
 
         elif axis == 'flux':
+            # handle spectral y-unit first since that is a more apparent change to the user
+            # and feels laggy if it is done later
+            if self.spectral_y_type_selected == 'Flux':
+                self._handle_spectral_y_unit()
+
             if len(self.angle_unit_selected):
                 # NOTE: setting sb_unit_selected will call this method again with axis=='sb',
                 # which in turn will call _handle_attribute_display_unit,
@@ -295,9 +310,6 @@ def _on_unit_selected(self, msg):
                 self.sb_unit_selected = _flux_to_sb_unit(self.flux_unit.selected,
                                                          self.angle_unit.selected)
 
-            if self.spectral_y_type_selected == 'Flux':
-                self._handle_spectral_y_unit()
-
         elif axis == 'angle':
             if len(self.flux_unit_selected):
                 # NOTE: setting sb_unit_selected will call this method again with axis=='sb',
@@ -308,11 +320,13 @@ def _on_unit_selected(self, msg):
                                                          self.angle_unit.selected)
 
         elif axis == 'sb':
-            self._handle_attribute_display_unit(self.sb_unit_selected)
-
+            # handle spectral y-unit first since that is a more apparent change to the user
+            # and feels laggy if it is done later
             if self.spectral_y_type_selected == 'Surface Brightness':
                 self._handle_spectral_y_unit()
 
+            self._handle_attribute_display_unit(self.sb_unit_selected)
+
         # custom axes downstream can override _on_unit_selected if anything needs to be
         # processed before the GlobalDisplayUnitChanged message is broadcast
 
@@ -354,9 +368,7 @@ def _handle_attribute_display_unit(self, attr_unit, layers=None):
         (updating stretch and contour units).
         """
         if layers is None:
-            layers = [layer
-                      for viewer in self._app._viewer_store.values() if isinstance(viewer, BqplotImageView)  # noqa
-                      for layer in viewer.layers]
+            layers = self.image_layers
 
         for layer in layers:
             # DQ layer doesn't play nicely with this attribute