Add Sky Regions to get_subsets()

CHANGES.rst

@@ -12,11 +12,16 @@ Cubeviz
 
 - Add circular annulus subset to toolbar. [#2438]
 
+- Expose sky regions in get_subsets. If 'include_sky_region' is True, a sky Region will be returned (in addition to a pixel Region) for spatial subsets with parent data that was a WCS. [#2496]
+
 Imviz
 ^^^^^
 
 - Aperture photometry (previously "Imviz Simple Aperture Photometry") now supports batch mode. [#2465]
 
+
+- Expose sky regions in get_subsets. If 'include_sky_region' is True, a sky Region will be returned (in addition to a pixel Region) for spatial subsets with parent data that was a WCS. [#2496]
+
 Mosviz
 ^^^^^^
 

jdaviz/app.py

@@ -922,10 +922,12 @@ 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, simplify_spectral=True,
-                    use_display_units=False):
+                    spatial_only=False, object_only=False,
+                    simplify_spectral=True, use_display_units=False,
+                    include_sky_region=False):
         """
-        Returns all branches of glue subset tree in the form that subset plugin can recognize.
+        Returns all branches of glue subset tree in the form that subset plugin
+        can recognize.
 
         Parameters
         ----------
@@ -939,16 +941,23 @@ def get_subsets(self, subset_name=None, spectral_only=False,
             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.
-        use_display_units: bool, optional
+            Return a composite spectral subset collapsed to a simplified
+            SpectralRegion.
+        use_display_units : bool, optional
             Whether to convert to the display units defined in the
             :ref:`Unit Conversion <unit-conversion>` plugin.
+        include_sky_region : bool
+            If True, for spatial subsets that have a WCS associated with their
+            parent data, return a sky region in addition to pixel region. If
+            subset is composite, a sky region for each constituent subset will
+            be returned.
 
         Returns
         -------
         data : dict
-            A dict with keys representing the subset name and values as astropy regions
-            objects.
+            Returns a nested dictionary with, for each subset, 'name', 'glue_state',
+            'region', 'sky_region' (set to None if not applicable), and 'subset_state'.
+            If subset is composite, each constituant subset will be included individually.
         """
 
         dc = self.data_collection
@@ -957,29 +966,37 @@ def get_subsets(self, subset_name=None, spectral_only=False,
         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, use_display_units)
+                                                     simplify_spectral, use_display_units,
+                                                     get_sky_regions=include_sky_region)
+
             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)
+
+                subset_region = self._get_roi_subset_definition(subset.subset_state,
+                                                                to_sky=include_sky_region)
+
             elif isinstance(subset.subset_state, RangeSubsetState):
                 # 2D regions represented as SpectralRegion objects
                 subset_region = self._get_range_subset_bounds(subset.subset_state,
                                                               simplify_spectral,
                                                               use_display_units)
+
             elif isinstance(subset.subset_state, MultiMaskSubsetState):
                 subset_region = self._get_multi_mask_subset_definition(subset.subset_state)
+
             else:
                 # subset.subset_state can be an instance of something else
                 # we do not know how to handle yet
                 all_subsets[label] = [{"name": subset.subset_state.__class__.__name__,
                                        "glue_state": subset.subset_state.__class__.__name__,
                                        "region": None,
+                                       "sky_region": None,
                                        "subset_state": subset.subset_state}]
                 continue
 
@@ -1014,6 +1031,8 @@ def get_subsets(self, subset_name=None, spectral_only=False,
                 else:
                     all_subsets[label] = subset_region
 
+        # can this be done at the top to avoid traversing all subsets if only
+        # one is requested?
         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]
@@ -1071,31 +1090,54 @@ def _get_range_subset_bounds(self, subset_state,
             return [{"name": subset_state.__class__.__name__,
                      "glue_state": subset_state.__class__.__name__,
                      "region": spec_region,
+                     "sky_region": None,  # not implemented for spectral, include for consistancy
                      "subset_state": subset_state}]
         return spec_region
 
     def _get_multi_mask_subset_definition(self, subset_state):
+
         return [{"name": subset_state.__class__.__name__,
                  "glue_state": subset_state.__class__.__name__,
                  "region": subset_state.total_masked_first_data(),
+                 "sky_region": None,
                  "subset_state": subset_state}]
 
-    def _get_roi_subset_definition(self, subset_state):
-        # TODO: Imviz: Return sky region if link type is WCS.
+    def _get_roi_subset_definition(self, subset_state, to_sky=False):
+
+        # pixel region
         roi_as_region = roi_subset_state_to_region(subset_state)
+
+        wcs = None
+        if to_sky:
+            if self.config == 'cubeviz':
+                parent_data = subset_state.attributes[0].parent
+                wcs = parent_data.meta.get("_orig_spatial_wcs", None)
+            else:
+                wcs = subset_state.xatt.parent.coords  # imviz, try getting WCS from subset data
+
+        # if no spatial wcs on subset, we have to skip computing sky region for this subset
+        # but want to do so without raising an error (since many subsets could be requested)
+        roi_as_sky_region = None
+        if wcs is not None:
+            roi_as_sky_region = roi_as_region.to_sky(wcs)
+
         return [{"name": subset_state.roi.__class__.__name__,
                  "glue_state": subset_state.__class__.__name__,
                  "region": roi_as_region,
+                 "sky_region": roi_as_sky_region,
                  "subset_state": subset_state}]
 
-    def get_sub_regions(self, subset_state, simplify_spectral=True, use_display_units=False):
+    def get_sub_regions(self, subset_state, simplify_spectral=True,
+                        use_display_units=False, get_sky_regions=False):
 
         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, use_display_units)
+                                           simplify_spectral, use_display_units,
+                                           get_sky_regions=get_sky_regions)
                 two = self.get_sub_regions(subset_state.state2,
-                                           simplify_spectral, use_display_units)
+                                           simplify_spectral, use_display_units,
+                                           get_sky_regions=get_sky_regions)
 
                 if isinstance(one, list) and "glue_state" in one[0]:
                     one[0]["glue_state"] = subset_state.__class__.__name__
@@ -1229,7 +1271,7 @@ def get_sub_regions(self, subset_state, simplify_spectral=True, use_display_unit
             # This is the leaf node of the glue subset state tree where
             # a subset_state is either ROI, Range, or MultiMask.
             if isinstance(subset_state, RoiSubsetState):
-                return self._get_roi_subset_definition(subset_state)
+                return self._get_roi_subset_definition(subset_state, to_sky=get_sky_regions)
 
             elif isinstance(subset_state, RangeSubsetState):
                 return self._get_range_subset_bounds(subset_state,

jdaviz/configs/cubeviz/plugins/parsers.py

@@ -130,6 +130,11 @@ def parse_data(app, file_obj, data_type=None, data_label=None):
         raise NotImplementedError(f'Unsupported data format: {file_obj}')
 
 
+def _get_celestial_wcs(wcs):
+    """ If `wcs` has a celestial component return that, otherwise return None """
+    return wcs.celestial if hasattr(wcs, 'celestial') else None
+
+
 def _return_spectrum_with_correct_units(flux, wcs, metadata, data_type, target_wave_unit=None,
                                         hdulist=None, uncertainty=None, mask=None):
     """Upstream issue of WCS not using the correct units for data must be fixed here.
@@ -232,6 +237,11 @@ def _parse_hdulist(app, hdulist, file_name=None,
             metadata[PRIHDR_KEY] = standardize_metadata(hdulist['PRIMARY'].header)
 
         sc = _return_spectrum_with_correct_units(flux, wcs, metadata, data_type, hdulist=hdulist)
+
+        # store original WCS in metadata. this is a hacky workaround for converting subsets
+        # to sky regions, where the parent data of the subset might have dropped spatial WCS info
+        metadata['_orig_spatial_wcs'] = _get_celestial_wcs(wcs)
+
         app.add_data(sc, data_label)
         if data_type == 'flux':  # Forced wave unit conversion made it lose stuff, so re-add
             app.data_collection[-1].get_component("flux").units = flux_unit
@@ -278,6 +288,11 @@ def _parse_jwst_s3d(app, hdulist, data_label, ext='SCI',
     wcs = WCS(hdulist['SCI'].header, hdulist)  # Everything uses SCI WCS
 
     metadata = standardize_metadata(hdu.header)
+
+    # store original WCS in metadata. this is a hacky workaround for converting subsets
+    # to sky regions, where the parent data of the subset might have dropped spatial WCS info
+    metadata['_orig_spatial_wcs'] = _get_celestial_wcs(wcs)
+
     if hdu.name != 'PRIMARY' and 'PRIMARY' in hdulist:
         metadata[PRIHDR_KEY] = standardize_metadata(hdulist['PRIMARY'].header)
 
@@ -324,6 +339,10 @@ def _parse_esa_s3d(app, hdulist, data_label, ext='DATA', flux_viewer_reference_n
     if hdu.name != 'PRIMARY' and 'PRIMARY' in hdulist:
         metadata[PRIHDR_KEY] = standardize_metadata(hdulist['PRIMARY'].header)
 
+    # store original WCS in metadata. this is a hacky workaround for converting subsets
+    # to sky regions, where the parent data of the subset might have dropped spatial WCS info
+    metadata['_orig_spatial_wcs'] = _get_celestial_wcs(wcs)
+
     data = _return_spectrum_with_correct_units(flux, wcs, metadata, data_type, hdulist=hdulist)
 
     app.add_data(data, data_label)
@@ -364,8 +383,16 @@ def _parse_spectrum1d_3d(app, file_obj, data_label=None,
             warnings.filterwarnings(
                 'ignore', message='Input WCS indicates that the spectral axis is not last',
                 category=UserWarning)
-            s1d = Spectrum1D(flux=flux, wcs=file_obj.wcs,
-                             meta=standardize_metadata(file_obj.meta))
+            meta = standardize_metadata(file_obj.meta)
+
+            # store original WCS in metadata. this is a hacky workaround for
+            # converting subsets to sky regions, where the parent data of the
+            # subset might have dropped spatial WCS info
+            meta['_orig_spatial_wcs'] = None
+            if hasattr(file_obj, 'wcs'):
+                meta['_orig_spatial_wcs'] = _get_celestial_wcs(file_obj.wcs)
+
+            s1d = Spectrum1D(flux=flux, wcs=file_obj.wcs, meta=meta)
 
         cur_data_label = app.return_data_label(data_label, attr.upper())
         app.add_data(s1d, cur_data_label)
@@ -379,9 +406,16 @@ def _parse_spectrum1d_3d(app, file_obj, data_label=None,
 
 
 def _parse_spectrum1d(app, file_obj, data_label=None, spectrum_viewer_reference_name=None):
+
+    # Here 'file_obj' is a Spectrum1D
+
     if data_label is None:
         data_label = app.return_data_label(file_obj)
 
+    # store original WCS in metadata. this is a hacky workaround for converting subsets
+    # to sky regions, where the parent data of the subset might have dropped spatial WCS info
+    file_obj.meta['_orig_spatial_wcs'] = _get_celestial_wcs(file_obj.wcs) if hasattr(file_obj, 'wcs') else None  # noqa: E501
+
     # TODO: glue-astronomy translators only look at the flux property of
     #  specutils Spectrum1D objects. Fix to support uncertainties and masks.
 
@@ -404,7 +438,9 @@ def _parse_ndarray(app, file_obj, data_label=None, data_type=None,
 
     if not hasattr(flux, 'unit'):
         flux = flux << u.count
-    s3d = Spectrum1D(flux=flux)
+
+    meta = standardize_metadata({'_orig_spatial_wcs': None})
+    s3d = Spectrum1D(flux=flux, meta=meta)
     app.add_data(s3d, data_label)
 
     if data_type == 'flux':
@@ -427,7 +463,9 @@ def _parse_gif(app, file_obj, data_label=None, flux_viewer_reference_name=None,
 
     flux = imageio.v3.imread(file_obj, mode='P')  # All frames as gray scale
     flux = np.rot90(np.moveaxis(flux, 0, 2), k=-1, axes=(0, 1))
-    s3d = Spectrum1D(flux=flux * u.count, meta={'filename': file_name})
+
+    meta = {'filename': file_name, '_orig_spatial_wcs': None}
+    s3d = Spectrum1D(flux=flux * u.count, meta=standardize_metadata(meta))
 
     app.add_data(s3d, data_label)
     app.add_data_to_viewer(flux_viewer_reference_name, data_label)