Refactor get subsets to work with composite subset state

CHANGES.rst

@@ -22,6 +22,8 @@ Specviz2d
 API Changes
 -----------
 
+- Add ``get_subsets()`` method to app level to centralize subset information retrieval. [#2087]
+
 Cubeviz
 ^^^^^^^
 

docs/cubeviz/export_data.rst

@@ -30,12 +30,12 @@ An example without accessing Specviz:
 
     subset1_spec1d = cubeviz.get_data(data_label=flux_data_label, 
                                       subset_to_apply="Subset 1",
-                                      statistic="mean")
+                                      function="mean")
 
-Note that in the above example, the ``statistic`` keyword is used to tell Cubeviz
+Note that in the above example, the ``function`` keyword is used to tell Cubeviz
 how to collapse the flux cube down to a one dimensional spectrum - this is not 
 necessarily equivalent to the collapsed spectrum in the spectrum viewer, which 
-may have used a different collapse statistic.
+may have used a different collapse function.
 
 To get all subsets from the spectrum viewer:
 

jdaviz/app.py

@@ -4,12 +4,16 @@
 import uuid
 import warnings
 from inspect import isclass
+import operator
 
 from ipywidgets import widget_serialization
 import ipyvue
 
 from astropy.nddata import CCDData, NDData
 from astropy.io import fits
+from astropy import units as u
+from astropy.coordinates import Angle
+from regions import PixCoord, CirclePixelRegion, RectanglePixelRegion, EllipsePixelRegion
 
 from echo import CallbackProperty, DictCallbackProperty, ListCallbackProperty
 from ipygoldenlayout import GoldenLayout
@@ -31,7 +35,9 @@
                                SubsetUpdateMessage,
                                SubsetDeleteMessage)
 from glue.core.state_objects import State
-from glue.core.subset import Subset, RangeSubsetState, RoiSubsetState
+from glue.core.subset import (Subset, RangeSubsetState, RoiSubsetState,
+                              CompositeSubsetState, InvertState)
+from glue.core.roi import CircularROI, EllipticalROI, RectangularROI
 from glue_astronomy.spectral_coordinates import SpectralCoordinates
 from glue_jupyter.app import JupyterApplication
 from glue_jupyter.common.toolbar_vuetify import read_icon
@@ -729,54 +735,6 @@ def get_subsets_from_viewer(self, viewer_reference, data_label=None, subset_type
                                          cls=None)
         regions = {}
 
-        def _get_all_subregions(mask, spec_axis_data):
-            """
-            Return all subregions within a subset.
-
-            Parameters
-            ----------
-            mask : list
-                List of indices in spec_axis_data that are part of the subset.
-            spec_axis_data : list
-                List of spectral axis values.
-            Returns
-            -------
-            combined_spec_region : `~specutils.SpectralRegion`
-                SpectralRegion object containing all subregions of the subset.
-            """
-            if len(mask) == 0:
-                # Mask should only be 0 if ApplyROI is used to incorrectly
-                # create subsets via the API
-                raise ValueError("Mask has length 0, ApplyROI may have been used incorrectly")
-
-            current_edge = 0
-            combined_spec_region = None
-            for index in range(1, len(mask)):
-                # Find spot where mask == True is for a different region of the subset
-                # i.e. mask = [0, 1, 4, 5]
-                # mask[2] != mask[1] + 1
-                if mask[index] != mask[index - 1] + 1:
-                    subset_region = spec_axis_data[mask[current_edge]: mask[index - 1] + 1]
-                    if not combined_spec_region:
-                        combined_spec_region = SpectralRegion(min(subset_region),
-                                                              max(subset_region))
-                    else:
-                        combined_spec_region += SpectralRegion(min(subset_region),
-                                                               max(subset_region))
-                    current_edge = index
-
-            # Get last region within the subset
-            if current_edge != index:
-                subset_region = spec_axis_data[mask[current_edge]: mask[index] + 1]
-                # No if check here because len(mask) must be greater than 1
-                # so combined_spec_region will have been instantiated in the for loop
-                if combined_spec_region is None:
-                    combined_spec_region = SpectralRegion(min(subset_region), max(subset_region))
-                else:
-                    combined_spec_region += SpectralRegion(min(subset_region), max(subset_region))
-
-            return combined_spec_region
-
         if data_label is not None:
             data = {data_label: data}
 
@@ -811,30 +769,12 @@ def _get_all_subregions(mask, spec_axis_data):
                 #  translation machinery entirely and construct the astropy
                 #  region ourselves.
                 elif value.data.ndim == 1:
-                    # Grab the data units from the glue-astronomy spectral axis
-                    # TODO: this needs to be much simpler; i.e. data units in
-                    #  the glue component objects
-                    # Cases where there is a single subset
-                    if '_orig_spec' in value.data.meta:  # Hack for Cubeviz WCS propagation
-                        data_wcs = value.data.meta['_orig_spec']
-                    else:
-                        data_wcs = value.data.coords
-
-                    subregions_in_subset = _get_all_subregions(
-                            np.where(value.to_mask() == True)[0], # noqa
-                            data_wcs.spectral_axis)
-
-                    regions[key] = subregions_in_subset
+                    regions[key] = self.get_subsets(key)
                     continue
 
                 temp_data = self.get_data_from_viewer(viewer_reference, value.label)
                 if isinstance(temp_data, Spectrum1D):
-                    # Note that we look for mask == False here, rather than True above,
-                    # because specutils masks are the reverse of Glue (of course)
-                    subregions_in_subset = _get_all_subregions(
-                             np.where(~temp_data.mask)[0], # noqa
-                             temp_data.spectral_axis)
-                    regions[key] = subregions_in_subset
+                    regions[key] = self.get_subsets(key)
                     continue
 
                 # Get the pixel coordinate [z] of the 3D data, repeating the
@@ -887,6 +827,203 @@ def _get_all_subregions(mask, spec_axis_data):
 
         return regions
 
+    def get_subsets(self, subset_name=None, spectral_only=False,
+                    spatial_only=False, object_only=False):
+        """
+        Returns all branches of glue subset tree in the form that subset plugin can recognize.
+
+        Parameters
+        ----------
+        subset_name : str
+            The subset name.
+        spectral_only : bool
+            Return only spectral subsets.
+        spatial_only : bool
+            Return only spatial subsets.
+        object_only : bool
+            Return only object relevant information and
+            leave out the region class name and glue_state.
+
+        Returns
+        -------
+        data : dict
+            A dict with keys representing the subset name and values as astropy regions
+            objects
+        """
+
+        dc = self.data_collection
+        subsets = dc.subset_groups
+
+        all_subsets = {}
+
+        for subset in subsets:
+            label = subset.label
+            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)
+            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)
+            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):
+                subset_region = self._remove_duplicate_bounds(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 object_only:
+                    all_subsets[label] = [reg['region'] for reg in subset_region]
+                else:
+                    all_subsets[label] = subset_region
+            elif not spectral_only and not spatial_only:
+                if object_only and not isinstance(subset_region, SpectralRegion):
+                    all_subsets[label] = [reg['region'] for reg in subset_region]
+                else:
+                    all_subsets[label] = subset_region
+
+        all_subset_names = [subset.label for subset in dc.subset_groups]
+        if subset_name and subset_name in all_subset_names:
+            return all_subsets[subset_name]
+        elif subset_name:
+            raise ValueError(f"{subset_name} not in {all_subset_names}")
+        else:
+            return all_subsets
+
+    def _remove_duplicate_bounds(self, spec_regions):
+        regions_no_dups = None
+
+        for region in spec_regions:
+            if not regions_no_dups:
+                regions_no_dups = region
+            elif region.bounds not in regions_no_dups.subregions:
+                regions_no_dups += region
+        return regions_no_dups
+
+    def _get_range_subset_bounds(self, subset_state):
+        # TODO: Use global display units
+        # units = dc[0].data.coords.spectral_axis.unit
+        viewer = self.get_viewer(self._jdaviz_helper. _default_spectrum_viewer_reference_name)
+        data = viewer.data()
+        if viewer:
+            units = u.Unit(viewer.state.x_display_unit)
+        elif data and len(data) > 0 and isinstance(data[0], Spectrum1D):
+            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)
+
+    def _get_roi_subset_definition(self, subset_state):
+        _around_decimals = 6
+        roi = subset_state.roi
+        roi_as_region = None
+        if isinstance(roi, CircularROI):
+            x, y = roi.get_center()
+            r = roi.radius
+            roi_as_region = CirclePixelRegion(PixCoord(x, y), r)
+
+        elif isinstance(roi, RectangularROI):
+            theta = np.around(np.degrees(roi.theta), decimals=_around_decimals)
+            roi_as_region = RectanglePixelRegion(PixCoord(roi.center()[0], roi.center()[1]),
+                                                 roi.width(), roi.height(), Angle(theta, "deg"))
+
+        elif isinstance(roi, EllipticalROI):
+            xc = roi.xc
+            yc = roi.yc
+            rx = roi.radius_x
+            ry = roi.radius_y
+            theta = np.around(np.degrees(roi.theta), decimals=_around_decimals)
+            roi_as_region = EllipsePixelRegion(PixCoord(xc, yc), rx, ry, Angle(theta, "deg"))
+
+        return [{"name": subset_state.roi.__class__.__name__,
+                 "glue_state": subset_state.__class__.__name__,
+                 "region": roi_as_region}]
+
+    def get_sub_regions(self, subset_state):
+
+        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)
+
+                if (isinstance(one, list) and "glue_state" in one[0] and
+                        one[0]["glue_state"] == "RoiSubsetState"):
+                    one[0]["glue_state"] = subset_state.__class__.__name__
+
+                if isinstance(subset_state.state2, InvertState):
+                    # This covers the REMOVE subset mode
+
+                    # As an example for how this works:
+                    # a = SpectralRegion(4 * u.um, 7 * u.um) + SpectralRegion(9 * u.um, 11 * u.um)
+                    # b = SpectralRegion(5 * u.um, 6 * u.um)
+                    # After running the following code with a as one and b as two:
+                    # Spectral Region, 3 sub-regions:
+                    #   (4.0 um, 5.0 um)    (6.0 um, 7.0 um)    (9.0 um, 11.0 um)
+                    if isinstance(two, SpectralRegion):
+                        new_spec = None
+                        for sub in one:
+                            if not new_spec:
+                                new_spec = two.invert(sub.lower, sub.upper)
+                            else:
+                                new_spec += two.invert(sub.lower, sub.upper)
+                        return new_spec
+                    else:
+                        if isinstance(two, list):
+                            # two[0]['glue_state'] = subset_state.state2.__class__.__name__
+                            two[0]['glue_state'] = "AndNotState"
+                        # Return two first so that we preserve the chronology of how
+                        # subset regions are applied.
+                        return two + one
+                elif subset_state.op is operator.and_:
+                    # This covers the AND subset mode
+
+                    # Example of how this works:
+                    # a = SpectralRegion(4 * u.um, 7 * u.um)
+                    # b = SpectralRegion(5 * u.um, 6 * u.um)
+                    #
+                    # b.invert(a.lower, a.upper)
+                    # Spectral Region, 2 sub-regions:
+                    #   (4.0 um, 5.0 um)   (6.0 um, 7.0 um)
+                    if isinstance(two, SpectralRegion):
+                        return two.invert(one.lower, one.upper)
+                    else:
+                        return one + two
+                elif subset_state.op is operator.or_:
+                    # This covers the ADD subset mode
+                    # one + two works for both Range and ROI subsets
+                    if one and two:
+                        return one + two
+                    elif one:
+                        return one
+                    elif two:
+                        return two
+                elif subset_state.op is operator.xor:
+                    # This covers the XOR case which is currently not working
+                    return None
+                else:
+                    return None
+            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)
+        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)
+
     def add_data(self, data, data_label=None, notify_done=True):
         """
         Add data to the Glue ``DataCollection``.

jdaviz/configs/cubeviz/plugins/tests/test_parsers.py

@@ -77,8 +77,8 @@ def test_spectrum1d_with_fake_fixed_units(spectrum1d, cubeviz_helper):
     reg = subsets.get('Subset 1')
 
     assert len(subsets) == 1
-    assert_allclose(reg.lower.value, 6666.666666666667)
-    assert_allclose(reg.upper.value, 7333.333333333334)
+    assert_allclose(reg.lower.value, 6600.)
+    assert_allclose(reg.upper.value, 7400.)
     assert reg.lower.unit == 'Angstrom'
     assert reg.upper.unit == 'Angstrom'
 

jdaviz/configs/default/plugins/model_fitting/tests/test_plugin.py

@@ -350,7 +350,7 @@ def test_invalid_subset(specviz_helper, spectrum1d):
     plugin.spectral_subset = 'Subset 1'
     assert not plugin._obj.spectral_subset_valid
 
-    with pytest.raises(ValueError, match=r"spectral subset 'Subset 1' \(5000.0, 5888.888888888889\) is outside data range of 'right_spectrum' \(6000.0, 8000.0\)"):  # noqa
+    with pytest.raises(ValueError, match=r"spectral subset 'Subset 1' \(5000.0, 6000.0\) is outside data range of 'right_spectrum' \(6000.0, 8000.0\)"):  # noqa
         plugin.calculate_fit()
 
     plugin.dataset = 'left_spectrum'

jdaviz/configs/specviz/helper.py

@@ -139,9 +139,7 @@ def get_spectral_regions(self):
             Mapping from the names of the subsets to the subsets expressed
             as `specutils.SpectralRegion` objects.
         """
-        return self.app.get_subsets_from_viewer(
-            self._default_spectrum_viewer_reference_name, subset_type="spectral"
-        )
+        return self.app.get_subsets(spectral_only=True)
 
     def x_limits(self, x_min=None, x_max=None):
         """Sets the limits of the x-axis