BUG: Fix wrong Subset unit when created in 2D spectrum viewer (spacet…

…elescope#3201)

* BUG: Fix wrong Subset unit in Specviz2d

* Add test and change log

* cubeviz is so special
and fixed an ambiguous specviz2d test, maybe

* Address comment from kecnry

* Move change log to match new milestone

CHANGES.rst

@@ -162,6 +162,8 @@ Specviz
 Specviz2d
 ^^^^^^^^^
 
+- Fixed Subset unit when it is created in 2D spectrum viewer. [#3201]
+
 Other Changes and Additions
 ---------------------------
 

jdaviz/app.py

@@ -14,7 +14,7 @@
 from ipysplitpanes import SplitPanes
 import matplotlib.cm as cm
 import numpy as np
-from glue.config import colormaps, settings as glue_settings
+from glue.config import colormaps, data_translator, settings as glue_settings
 from glue.core import HubListener
 from glue.core.link_helpers import LinkSame, LinkSameWithUnits
 from glue.core.message import (DataCollectionAddMessage,
@@ -1052,13 +1052,24 @@ def _get_range_subset_bounds(self, subset_state,
                                  simplify_spectral=True, use_display_units=False):
         # TODO: Use global display units
         # units = dc[0].data.coords.spectral_axis.unit
-        viewer = self.get_viewer(self._jdaviz_helper. _default_spectrum_viewer_reference_name)
-        data = viewer.data()
-        if data and len(data) > 0 and isinstance(data[0], Spectrum1D):
-            units = data[0].spectral_axis.unit
+        if not hasattr(subset_state.att, "parent"):  # e.g., Cubeviz
+            viewer = self.get_viewer(self._jdaviz_helper._default_spectrum_viewer_reference_name)
+            data = viewer.data()
+            if data and len(data) > 0 and isinstance(data[0], Spectrum1D):
+                units = data[0].spectral_axis.unit
+            else:
+                raise ValueError("Unable to find spectral axis units")
         else:
-            raise ValueError("Unable to find spectral axis units")
-        if use_display_units:
+            data = subset_state.att.parent
+            ndim = data.get_component("flux").ndim
+            if ndim == 2:
+                units = u.pix
+            else:
+                handler, _ = data_translator.get_handler_for(Spectrum1D)
+                spec = handler.to_object(data)
+                units = spec.spectral_axis.unit
+
+        if use_display_units and units != u.pix:
             # converting may result in flipping order (wavelength <-> frequency)
             ret_units = self._get_display_unit('spectral')
             subset_bounds = [(subset_state.lo * units).to(ret_units, u.spectral()),

jdaviz/configs/specviz2d/tests/test_parsers.py

@@ -1,7 +1,10 @@
 import pytest
 import stdatamodels
+from astropy import units as u
 from astropy.io import fits
 from astropy.utils.data import download_file
+from glue.core.edit_subset_mode import NewMode
+from glue.core.roi import XRangeROI
 
 from jdaviz.utils import PRIHDR_KEY
 from jdaviz.configs.imviz.tests.utils import create_example_gwcs
@@ -134,6 +137,25 @@ def test_2d_1d_parser(specviz2d_helper, mos_spectrum2d, spectrum1d):
     specviz2d_helper.load_data(spectrum_2d=mos_spectrum2d, spectrum_1d=spectrum1d)
     assert specviz2d_helper.app.data_collection.labels == ['Spectrum 2D', 'Spectrum 1D']
 
+    spec2d_viewer = specviz2d_helper.app.get_viewer("spectrum-2d-viewer")
+    assert spec2d_viewer.figure.axes[0].label == "x: pixels"  # -0.5, 14.5
+    spec2d_viewer.apply_roi(XRangeROI(10, 12))
+
+    spec2d_viewer.session.edit_subset_mode._mode = NewMode
+
+    spec1d_viewer = specviz2d_helper.app.get_viewer("spectrum-viewer")
+    assert spec1d_viewer.figure.axes[0].label == "Wavelength [Angstrom]"  # 6000, 8000
+    spec1d_viewer.apply_roi(XRangeROI(7000, 7100))
+
+    # Subset 1 should follow Spectrum 2D viewer unit.
+    # Subset 2 should follow Spectrum 1D viewer unit.
+    subsets = specviz2d_helper.app.get_subsets()
+    assert len(subsets) == 2
+    s1 = subsets["Subset 1"]
+    s2 = subsets["Subset 2"]
+    assert s1.lower.unit == s1.upper.unit == u.pix
+    assert s2.lower.unit == s2.upper.unit == u.AA
+
 
 def test_parser_no_data(specviz2d_helper):
     with pytest.raises(ValueError, match='Must provide spectrum_2d or spectrum_1d'):