Add option to see full glue state tree for spectral subsets (#2138)

* Add option to see full glue state tree for spectral subsets

* Fix codestyle checks

* Add changelog

CHANGES.rst

@@ -34,7 +34,7 @@ API Changes
 -----------
 
 - Add ``get_subsets()`` method to app level to centralize subset information
-  retrieval. [#2087, #2116]
+  retrieval. [#2087, #2116, #2138]
 
 Cubeviz
 ^^^^^^^

jdaviz/app.py

@@ -828,7 +828,7 @@ def get_subsets_from_viewer(self, viewer_reference, data_label=None, subset_type
         return regions
 
     def get_subsets(self, subset_name=None, spectral_only=False,
-                    spatial_only=False, object_only=False):
+                    spatial_only=False, object_only=False, simplify_spectral=True):
         """
         Returns all branches of glue subset tree in the form that subset plugin can recognize.
 
@@ -843,6 +843,8 @@ def get_subsets(self, subset_name=None, spectral_only=False,
         object_only : bool
             Return only object relevant information and
             leave out the region class name and glue_state.
+        simplify_spectral : bool
+            Return a composite spectral subset collapsed to a simplified SpectralRegion.
 
         Returns
         -------
@@ -861,26 +863,36 @@ def get_subsets(self, subset_name=None, spectral_only=False,
             if isinstance(subset.subset_state, CompositeSubsetState):
                 # Region composed of multiple ROI or Range subset
                 # objects that must be traversed
-                subset_region = self.get_sub_regions(subset.subset_state)
+                subset_region = self.get_sub_regions(subset.subset_state, simplify_spectral)
             elif isinstance(subset.subset_state, RoiSubsetState):
                 # 3D regions represented as a dict including an
                 # AstropyRegion object if possible
                 subset_region = self._get_roi_subset_definition(subset.subset_state)
             elif isinstance(subset.subset_state, RangeSubsetState):
                 # 2D regions represented as SpectralRegion objects
-                subset_region = self._get_range_subset_bounds(subset.subset_state)
+                subset_region = self._get_range_subset_bounds(subset.subset_state,
+                                                              simplify_spectral)
             else:
                 # subset.subset_state can be an instance of MaskSubsetState
                 # or something else we do not know how to handle
                 all_subsets[label] = None
                 continue
 
-            if isinstance(subset_region, SpectralRegion):
+            # Is the subset spectral or spatial?
+            is_spectral = self._is_subset_spectral(subset_region)
+
+            # Remove duplicate spectral regions
+            if is_spectral and isinstance(subset_region, SpectralRegion):
                 subset_region = self._remove_duplicate_bounds(subset_region)
+            elif is_spectral:
+                subset_region = self._remove_duplicate_bounds_in_dict(subset_region)
 
-            if spectral_only and isinstance(subset_region, SpectralRegion):
-                all_subsets[label] = subset_region
-            elif spatial_only and not isinstance(subset_region, SpectralRegion):
+            if spectral_only and is_spectral:
+                if object_only and not simplify_spectral:
+                    all_subsets[label] = [reg['region'] for reg in subset_region]
+                else:
+                    all_subsets[label] = subset_region
+            elif spatial_only and not is_spectral:
                 if object_only:
                     all_subsets[label] = [reg['region'] for reg in subset_region]
                 else:
@@ -899,6 +911,30 @@ def get_subsets(self, subset_name=None, spectral_only=False,
         else:
             return all_subsets
 
+    def _remove_duplicate_bounds_in_dict(self, subset_region):
+        new_subset_region = []
+        for elem in subset_region:
+            if not new_subset_region:
+                new_subset_region.append(elem)
+                continue
+            unique = True
+            for elem2 in new_subset_region:
+                if (elem['region'].lower == elem2['region'].lower and
+                        elem['region'].upper == elem2['region'].upper and
+                        elem['glue_state'] == elem2['glue_state']):
+                    unique = False
+            if unique:
+                new_subset_region.append(elem)
+        return new_subset_region
+
+    def _is_subset_spectral(self, subset_region):
+        if isinstance(subset_region, SpectralRegion):
+            return True
+        elif isinstance(subset_region, list) and len(subset_region) > 0:
+            if isinstance(subset_region[0]['region'], SpectralRegion):
+                return True
+        return False
+
     def _remove_duplicate_bounds(self, spec_regions):
         regions_no_dups = None
 
@@ -909,7 +945,7 @@ def _remove_duplicate_bounds(self, spec_regions):
                 regions_no_dups += region
         return regions_no_dups
 
-    def _get_range_subset_bounds(self, subset_state):
+    def _get_range_subset_bounds(self, subset_state, simplify_spectral=True):
         # 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)
@@ -920,7 +956,14 @@ def _get_range_subset_bounds(self, subset_state):
             units = data[0].spectral_axis.unit
         else:
             raise ValueError("Unable to find spectral axis units")
-        return SpectralRegion(subset_state.lo * units, subset_state.hi * units)
+
+        spec_region = SpectralRegion(subset_state.lo * units, subset_state.hi * units)
+        if not simplify_spectral:
+            return [{"name": subset_state.__class__.__name__,
+                     "glue_state": subset_state.__class__.__name__,
+                     "region": spec_region,
+                     "subset_state": subset_state}]
+        return spec_region
 
     def _get_roi_subset_definition(self, subset_state):
         _around_decimals = 6
@@ -946,17 +989,17 @@ def _get_roi_subset_definition(self, subset_state):
 
         return [{"name": subset_state.roi.__class__.__name__,
                  "glue_state": subset_state.__class__.__name__,
-                 "region": roi_as_region}]
+                 "region": roi_as_region,
+                 "subset_state": subset_state}]
 
-    def get_sub_regions(self, subset_state):
+    def get_sub_regions(self, subset_state, simplify_spectral=True):
 
         if isinstance(subset_state, CompositeSubsetState):
             if subset_state and hasattr(subset_state, "state2") and subset_state.state2:
-                one = self.get_sub_regions(subset_state.state1)
-                two = self.get_sub_regions(subset_state.state2)
+                one = self.get_sub_regions(subset_state.state1, simplify_spectral)
+                two = self.get_sub_regions(subset_state.state2, simplify_spectral)
 
-                if (isinstance(one, list) and "glue_state" in one[0] and
-                        one[0]["glue_state"] == "RoiSubsetState"):
+                if isinstance(one, list) and "glue_state" in one[0]:
                     one[0]["glue_state"] = subset_state.__class__.__name__
 
                 if isinstance(subset_state.state2, InvertState):
@@ -1014,15 +1057,15 @@ def get_sub_regions(self, subset_state):
             else:
                 # This gets triggered in the InvertState case where state1
                 # is an object and state2 is None
-                return self.get_sub_regions(subset_state.state1)
+                return self.get_sub_regions(subset_state.state1, simplify_spectral)
         elif subset_state is not None:
             # This is the leaf node of the glue subset state tree where
             # a subset_state is either ROI or Range.
             if isinstance(subset_state, RoiSubsetState):
                 return self._get_roi_subset_definition(subset_state)
 
             elif isinstance(subset_state, RangeSubsetState):
-                return self._get_range_subset_bounds(subset_state)
+                return self._get_range_subset_bounds(subset_state, simplify_spectral)
 
     def add_data(self, data, data_label=None, notify_done=True):
         """

jdaviz/tests/test_subsets.py

@@ -322,34 +322,39 @@ def test_composite_region_from_subset_3d(cubeviz_helper):
     viewer.apply_roi(CircularROI(xc=25, yc=25, radius=5))
     reg = cubeviz_helper.app.get_subsets("Subset 1")
     circle1 = CirclePixelRegion(center=PixCoord(x=25, y=25), radius=5)
-    assert reg[-1] == {'name': 'CircularROI', 'glue_state': 'RoiSubsetState', 'region': circle1}
+    assert reg[-1] == {'name': 'CircularROI', 'glue_state': 'RoiSubsetState', 'region': circle1,
+                       'subset_state': reg[-1]['subset_state']}
 
     cubeviz_helper.app.session.edit_subset_mode.mode = AndNotMode
     viewer.apply_roi(RectangularROI(25, 30, 25, 30))
     reg = cubeviz_helper.app.get_subsets("Subset 1")
     rectangle1 = RectanglePixelRegion(center=PixCoord(x=27.5, y=27.5),
                                       width=5, height=5, angle=0.0 * u.deg)
-    assert reg[-1] == {'name': 'RectangularROI', 'glue_state': 'AndNotState', 'region': rectangle1}
+    assert reg[-1] == {'name': 'RectangularROI', 'glue_state': 'AndNotState', 'region': rectangle1,
+                       'subset_state': reg[-1]['subset_state']}
 
     cubeviz_helper.app.session.edit_subset_mode.mode = OrMode
     viewer.apply_roi(EllipticalROI(30, 30, 3, 6))
     reg = cubeviz_helper.app.get_subsets("Subset 1")
     ellipse1 = EllipsePixelRegion(center=PixCoord(x=30, y=30),
                                   width=3, height=6, angle=0.0 * u.deg)
-    assert reg[-1] == {'name': 'EllipticalROI', 'glue_state': 'OrState', 'region': ellipse1}
+    assert reg[-1] == {'name': 'EllipticalROI', 'glue_state': 'OrState', 'region': ellipse1,
+                       'subset_state': reg[-1]['subset_state']}
 
     cubeviz_helper.app.session.edit_subset_mode.mode = AndMode
     viewer.apply_roi(RectangularROI(20, 25, 20, 25))
     reg = cubeviz_helper.app.get_subsets("Subset 1")
     rectangle2 = RectanglePixelRegion(center=PixCoord(x=22.5, y=22.5),
                                       width=5, height=5, angle=0.0 * u.deg)
-    assert reg[-1] == {'name': 'RectangularROI', 'glue_state': 'AndState', 'region': rectangle2}
+    assert reg[-1] == {'name': 'RectangularROI', 'glue_state': 'AndState', 'region': rectangle2,
+                       'subset_state': reg[-1]['subset_state']}
 
     cubeviz_helper.app.session.edit_subset_mode.mode = AndNotMode
     viewer.apply_roi(CircularROI(xc=21, yc=24, radius=1))
     reg = cubeviz_helper.app.get_subsets("Subset 1")
     circle2 = CirclePixelRegion(center=PixCoord(x=21, y=24), radius=1)
-    assert reg[-1] == {'name': 'CircularROI', 'glue_state': 'AndNotState', 'region': circle2}
+    assert reg[-1] == {'name': 'CircularROI', 'glue_state': 'AndNotState', 'region': circle2,
+                       'subset_state': reg[-1]['subset_state']}
 
 
 def test_composite_region_with_consecutive_and_not_states(cubeviz_helper):
@@ -362,21 +367,24 @@ def test_composite_region_with_consecutive_and_not_states(cubeviz_helper):
     viewer.apply_roi(CircularROI(xc=25, yc=25, radius=5))
     reg = cubeviz_helper.app.get_subsets("Subset 1")
     circle1 = CirclePixelRegion(center=PixCoord(x=25, y=25), radius=5)
-    assert reg[-1] == {'name': 'CircularROI', 'glue_state': 'RoiSubsetState', 'region': circle1}
+    assert reg[-1] == {'name': 'CircularROI', 'glue_state': 'RoiSubsetState', 'region': circle1,
+                       'subset_state': reg[-1]['subset_state']}
 
     cubeviz_helper.app.session.edit_subset_mode.mode = AndNotMode
     viewer.apply_roi(RectangularROI(25, 30, 25, 30))
     reg = cubeviz_helper.app.get_subsets("Subset 1")
     rectangle1 = RectanglePixelRegion(center=PixCoord(x=27.5, y=27.5),
                                       width=5, height=5, angle=0.0 * u.deg)
-    assert reg[-1] == {'name': 'RectangularROI', 'glue_state': 'AndNotState', 'region': rectangle1}
+    assert reg[-1] == {'name': 'RectangularROI', 'glue_state': 'AndNotState', 'region': rectangle1,
+                       'subset_state': reg[-1]['subset_state']}
 
     cubeviz_helper.app.session.edit_subset_mode.mode = AndNotMode
     viewer.apply_roi(EllipticalROI(30, 30, 3, 6))
     reg = cubeviz_helper.app.get_subsets("Subset 1")
     ellipse1 = EllipsePixelRegion(center=PixCoord(x=30, y=30),
                                   width=3, height=6, angle=0.0 * u.deg)
-    assert reg[-1] == {'name': 'EllipticalROI', 'glue_state': 'AndNotState', 'region': ellipse1}
+    assert reg[-1] == {'name': 'EllipticalROI', 'glue_state': 'AndNotState', 'region': ellipse1,
+                       'subset_state': reg[-1]['subset_state']}
 
     regions_list = cubeviz_helper.app.get_subsets("Subset 1", object_only=True)
     assert len(regions_list) == 3
@@ -400,24 +408,65 @@ def test_composite_region_with_imviz(imviz_helper, image_2d_wcs):
     viewer.apply_roi(CircularROI(xc=5, yc=5, radius=2))
     reg = imviz_helper.app.get_subsets("Subset 1")
     circle1 = CirclePixelRegion(center=PixCoord(x=5, y=5), radius=2)
-    assert reg[-1] == {'name': 'CircularROI', 'glue_state': 'RoiSubsetState', 'region': circle1}
+    assert reg[-1] == {'name': 'CircularROI', 'glue_state': 'RoiSubsetState', 'region': circle1,
+                       'subset_state': reg[-1]['subset_state']}
 
     imviz_helper.app.session.edit_subset_mode.mode = AndNotMode
     viewer.apply_roi(RectangularROI(2, 4, 2, 4))
     reg = imviz_helper.app.get_subsets("Subset 1")
     rectangle1 = RectanglePixelRegion(center=PixCoord(x=3, y=3),
                                       width=2, height=2, angle=0.0 * u.deg)
-    assert reg[-1] == {'name': 'RectangularROI', 'glue_state': 'AndNotState', 'region': rectangle1}
+    assert reg[-1] == {'name': 'RectangularROI', 'glue_state': 'AndNotState', 'region': rectangle1,
+                       'subset_state': reg[-1]['subset_state']}
 
     imviz_helper.app.session.edit_subset_mode.mode = AndNotMode
     viewer.apply_roi(EllipticalROI(3, 3, 3, 6))
     reg = imviz_helper.app.get_subsets("Subset 1")
     ellipse1 = EllipsePixelRegion(center=PixCoord(x=3, y=3),
                                   width=3, height=6, angle=0.0 * u.deg)
-    assert reg[-1] == {'name': 'EllipticalROI', 'glue_state': 'AndNotState', 'region': ellipse1}
+    assert reg[-1] == {'name': 'EllipticalROI', 'glue_state': 'AndNotState', 'region': ellipse1,
+                       'subset_state': reg[-1]['subset_state']}
 
 
 def test_with_invalid_subset_name(cubeviz_helper):
     subset_name = "Test"
     with pytest.raises(ValueError, match=f'{subset_name} not in '):
         cubeviz_helper.app.get_subsets(subset_name=subset_name)
+
+
+def test_composite_region_from_subset_2d(specviz_helper, spectrum1d):
+    specviz_helper.load_spectrum(spectrum1d)
+    viewer = specviz_helper.app.get_viewer(specviz_helper._default_spectrum_viewer_reference_name)
+    viewer.apply_roi(XRangeROI(6000, 7000))
+    reg = specviz_helper.app.get_subsets("Subset 1", simplify_spectral=False)
+    subset1 = SpectralRegion(6000 * spectrum1d.spectral_axis.unit,
+                             7000 * spectrum1d.spectral_axis.unit)
+    assert reg[-1]['region'].lower == subset1.lower and reg[-1]['region'].upper == subset1.upper
+    assert reg[-1]['glue_state'] == 'RangeSubsetState'
+
+    specviz_helper.app.session.edit_subset_mode.mode = AndNotMode
+
+    viewer.apply_roi(XRangeROI(6500, 6800))
+    reg = specviz_helper.app.get_subsets("Subset 1", simplify_spectral=False)
+    subset1 = SpectralRegion(6500 * spectrum1d.spectral_axis.unit,
+                             6800 * spectrum1d.spectral_axis.unit)
+    assert reg[-1]['region'].lower == subset1.lower and reg[-1]['region'].upper == subset1.upper
+    assert reg[-1]['glue_state'] == 'AndNotState'
+
+    specviz_helper.app.session.edit_subset_mode.mode = OrMode
+
+    viewer.apply_roi(XRangeROI(7200, 7800))
+    reg = specviz_helper.app.get_subsets("Subset 1", simplify_spectral=False)
+    subset1 = SpectralRegion(7200 * spectrum1d.spectral_axis.unit,
+                             7800 * spectrum1d.spectral_axis.unit)
+    assert reg[-1]['region'].lower == subset1.lower and reg[-1]['region'].upper == subset1.upper
+    assert reg[-1]['glue_state'] == 'OrState'
+
+    specviz_helper.app.session.edit_subset_mode.mode = AndMode
+
+    viewer.apply_roi(XRangeROI(6800, 7500))
+    reg = specviz_helper.app.get_subsets("Subset 1", simplify_spectral=False)
+    subset1 = SpectralRegion(6800 * spectrum1d.spectral_axis.unit,
+                             7500 * spectrum1d.spectral_axis.unit)
+    assert reg[-1]['region'].lower == subset1.lower and reg[-1]['region'].upper == subset1.upper
+    assert reg[-1]['glue_state'] == 'AndState'