Backport PR spacetelescope#2168: adding support for temporal subsets

jdaviz/app.py

@@ -10,6 +10,10 @@
 
 from astropy.nddata import CCDData, NDData
 from astropy.io import fits
+from astropy import units as u
+from astropy.coordinates import Angle
+from astropy.time import Time
+from regions import PixCoord, CirclePixelRegion, RectanglePixelRegion, EllipsePixelRegion
 
 from echo import CallbackProperty, DictCallbackProperty, ListCallbackProperty
 from ipygoldenlayout import GoldenLayout
@@ -680,8 +684,13 @@ def get_data_from_viewer(self, viewer_reference, data_label=None,
                         if cls is not None:
                             handler, _ = data_translator.get_handler_for(cls)
                             try:
-                                layer_data = handler.to_object(layer_data,
-                                                               statistic=statistic)
+                                if cls == Spectrum1D:
+                                    # if this is a spectrum, apply the `statistic`:
+                                    layer_data = handler.to_object(layer_data,
+                                                                   statistic=statistic)
+                                else:
+                                    # otherwise simply translate to an object:
+                                    layer_data = handler.to_object(layer_data)
                             except IncompatibleAttribute:
                                 continue
 
@@ -887,6 +896,261 @@ 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, simplify_spectral=True):
+        """
+        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.
+        simplify_spectral : bool
+            Return a composite spectral subset collapsed to a simplified SpectralRegion.
+
+        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, 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,
+                                                              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
+
+            # Is the subset spectral, spatial, temporal?
+            is_spectral = self._is_subset_spectral(subset_region)
+            is_temporal = self._is_subset_temporal(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 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:
+                    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
+
+            if not (spectral_only or spatial_only) and is_temporal:
+                if object_only:
+                    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_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 _is_subset_temporal(self, subset_region):
+        if isinstance(subset_region, Time):
+            return True
+        elif isinstance(subset_region, list) and len(subset_region) > 0:
+            if isinstance(subset_region[0]['region'], Time):
+                return True
+        return False
+
+    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, 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)
+        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")
+
+        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
+        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,
+                 "subset_state": 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, simplify_spectral)
+                two = self.get_sub_regions(subset_state.state2, simplify_spectral)
+
+                if isinstance(one, list) and "glue_state" in one[0]:
+                    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 one + two
+                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 two + one
+                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 two + one
+                    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, 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, simplify_spectral)
+
     def add_data(self, data, data_label=None, notify_done=True):
         """
         Add data to the Glue ``DataCollection``.