Add option to get_data to apply spectral subset to collapse spatial (#…

…2199)

CHANGES.rst

@@ -28,6 +28,9 @@ Cubeviz
 - ``get_data`` now supports ``function=True`` to adopt the collapse-function from the spectrum viewer.
   [#2117]
 
+- ``get_data`` now supports applying a spectral mask to a collapse spatial subset. [#2199]
+
+
 Imviz
 ^^^^^
 

jdaviz/configs/cubeviz/helper.py

@@ -120,7 +120,8 @@ def specviz(self):
             self._specviz = Specviz(app=self.app)
         return self._specviz
 
-    def get_data(self, data_label=None, cls=None, subset_to_apply=None, function=None):
+    def get_data(self, data_label=None, spatial_subset=None, spectral_subset=None, function=None,
+                 cls=None):
         """
         Returns data with name equal to data_label of type cls with subsets applied from
         subset_to_apply.
@@ -129,29 +130,39 @@ def get_data(self, data_label=None, cls=None, subset_to_apply=None, function=Non
         ----------
         data_label : str, optional
             Provide a label to retrieve a specific data set from data_collection.
-        cls : `~specutils.Spectrum1D`, `~astropy.nddata.CCDData`, optional
-            The type that data will be returned as.
-        subset_to_apply : str, optional
-            Subset that is to be applied to data before it is returned.
+        spatial_subset : str, optional
+            Spatial subset applied to data.
+        spectral_subset : str, optional
+            Spectral subset applied to data.
         function : {True, False, 'minimum', 'maximum', 'mean', 'median', 'sum'}, optional
             Ignored if ``data_label`` does not point to cube-like data.
             If True, will collapse according to the current collapse function defined in the
             spectrum viewer.  If provided as a string, the cube will be collapsed with the provided
-            function.  If False, None, or not passed, the entire cube will be returned.
+            function.  If False, None, or not passed, the entire cube will be returned (unless there
+            are values for ``spatial_subset`` and ``spectral_subset``).
+        cls : `~specutils.Spectrum1D`, `~astropy.nddata.CCDData`, optional
+            The type that data will be returned as.
 
         Returns
         -------
         data : cls
             Data is returned as type cls with subsets applied.
 
         """
-        if function is True:
-            return self.specviz.get_data(data_label=data_label, cls=cls,
-                                         subset_to_apply=subset_to_apply)
+        # If function is a value ('sum' or 'minimum') or True and spatial and spectral
+        # are set, then we collapse the cube along the spatial subset using the function, then
+        # we apply the mask from the spectral subset.
+        # If function is any value other than False, we use specviz
+        if (function is not False and spectral_subset and spatial_subset) or function:
+            return self.specviz.get_data(data_label=data_label, spectral_subset=spectral_subset,
+                                         cls=cls, spatial_subset=spatial_subset, function=function)
+        elif function is False and spectral_subset:
+            raise ValueError("function cannot be False if spectral_subset"
+                             " is set")
         elif function is False:
             function = None
-        return self._get_data(data_label=data_label, cls=cls, subset_to_apply=subset_to_apply,
-                              function=function)
+        return self._get_data(data_label=data_label, spatial_subset=spatial_subset,
+                              spectral_subset=spectral_subset, function=function, cls=cls)
 
 
 def layer_is_cube_image_data(layer):

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

@@ -4,6 +4,8 @@
 from astropy.tests.helper import assert_quantity_allclose
 from specutils import Spectrum1D
 
+from glue.core.roi import XRangeROI
+
 
 def test_nested_helper(cubeviz_helper):
     '''Ensures the Cubeviz helper is always returned, even after the Specviz helper is called'''
@@ -29,7 +31,7 @@ def test_invalid_function(cubeviz_helper, spectrum1d_cube):
     cubeviz_helper._apply_interactive_region('bqplot:ellipse', (0, 0), (9, 8))
 
     with pytest.raises(ValueError, match='function 42 not in list of valid '):
-        cubeviz_helper.get_data(data_label="test[FLUX]", subset_to_apply='Subset 1', function=42)
+        cubeviz_helper.get_data(data_label="test[FLUX]", spatial_subset='Subset 1', function=42)
 
     # Also make sure specviz redshift slider warning does not show up.
     # https://github.com/spacetelescope/jdaviz/issues/2029
@@ -41,20 +43,20 @@ def test_valid_function(cubeviz_helper, spectrum1d_cube):
     cubeviz_helper._apply_interactive_region('bqplot:ellipse', (0, 0), (9, 8))
 
     results_cube = cubeviz_helper.get_data(data_label="test[FLUX]",
-                                           subset_to_apply='Subset 1')
+                                           spatial_subset='Subset 1')
     assert results_cube.flux.ndim == 3
     results_false = cubeviz_helper.get_data(data_label="test[FLUX]",
-                                            subset_to_apply='Subset 1', function=False)
+                                            spatial_subset='Subset 1', function=False)
     assert results_false.flux.ndim == 3
 
     results_def = cubeviz_helper.get_data(data_label="test[FLUX]",
-                                          subset_to_apply='Subset 1', function=True)
+                                          spatial_subset='Subset 1', function=True)
     assert results_def.flux.ndim == 1
 
     results_min = cubeviz_helper.get_data(data_label="test[FLUX]",
-                                          subset_to_apply='Subset 1', function="minimum")
+                                          spatial_subset='Subset 1', function="minimum")
     results_max = cubeviz_helper.get_data(data_label="test[FLUX]",
-                                          subset_to_apply='Subset 1', function="maximum")
+                                          spatial_subset='Subset 1', function="maximum")
     assert isinstance(results_min, Spectrum1D)
     assert_quantity_allclose(results_min.flux,
                              [6., 14.] * u.Jy, atol=1e-5 * u.Jy)
@@ -65,3 +67,50 @@ def test_valid_function(cubeviz_helper, spectrum1d_cube):
     assert cubeviz_helper.get_data(data_label="test[FLUX]").flux.ndim == 3
     # but calling through specviz automatically collapses
     assert cubeviz_helper.specviz.get_data(data_label="test[FLUX]").flux.ndim == 1
+
+
+def test_get_data_spatial_and_spectral(cubeviz_helper, spectrum1d_cube_larger):
+    data_label = "test"
+    spatial_subset = "Subset 1"
+    spectral_subset = "Subset 2"
+    cubeviz_helper.load_data(spectrum1d_cube_larger, data_label)
+    cubeviz_helper._apply_interactive_region('bqplot:ellipse', (0, 0), (9, 8))
+
+    spec_viewer = cubeviz_helper.app.get_viewer(cubeviz_helper._default_spectrum_viewer_reference_name) # noqa
+    spec_viewer.apply_roi(XRangeROI(4.62440061e-07, 4.62520112e-07))
+
+    data_label = data_label + "[FLUX]"
+    # This will be the same if function is None or True
+    spatial_with_spec = cubeviz_helper.get_data(data_label=data_label,
+                                                spatial_subset=spatial_subset,
+                                                spectral_subset=spectral_subset)
+    assert spatial_with_spec.flux.ndim == 1
+    assert list(spatial_with_spec.mask) == [True, True, False, False, True,
+                                            True, True, True, True, True]
+    assert max(list(spatial_with_spec.flux.value)) == 157.
+    assert min(list(spatial_with_spec.flux.value)) == 13.
+
+    spatial_with_spec = cubeviz_helper.get_data(data_label=data_label,
+                                                spatial_subset=spatial_subset,
+                                                spectral_subset=spectral_subset,
+                                                function='minimum')
+    assert max(list(spatial_with_spec.flux.value)) == 78.
+    assert min(list(spatial_with_spec.flux.value)) == 6.
+
+    collapse_with_spectral = cubeviz_helper.get_data(data_label=data_label,
+                                                     spectral_subset=spectral_subset,
+                                                     function=True)
+    collapse_with_spectral2 = cubeviz_helper.get_data(data_label=data_label,
+                                                      function=True)
+
+    assert list(collapse_with_spectral.flux) == list(collapse_with_spectral2.flux)
+
+    with pytest.raises(ValueError, match=f'{spectral_subset} is not a spatial subset.'):
+        cubeviz_helper.get_data(data_label=data_label, spatial_subset=spectral_subset,
+                                function=True)
+    with pytest.raises(ValueError, match=f'{spatial_subset} is not a spectral subset.'):
+        cubeviz_helper.get_data(data_label=data_label, spectral_subset=spatial_subset,
+                                function=True)
+    with pytest.raises(ValueError, match='function cannot be False if spectral_subset'):
+        cubeviz_helper.get_data(data_label=data_label, spectral_subset=spectral_subset,
+                                function=False)

jdaviz/configs/imviz/helper.py

@@ -239,6 +239,28 @@ def get_catalog_source_results(self):
         """
         return getattr(self.app, '_catalog_source_table', None)
 
+    def get_data(self, data_label=None, spatial_subset=None, cls=None):
+        """
+        Returns data with name equal to data_label of type cls with subsets applied from
+        spatial_subset.
+
+        Parameters
+        ----------
+        data_label : str, optional
+            Provide a label to retrieve a specific data set from data_collection.
+        spatial_subset : str, optional
+            Spatial subset applied to data.
+        cls : `~specutils.Spectrum1D`, `~astropy.nddata.CCDData`, optional
+            The type that data will be returned as.
+
+        Returns
+        -------
+        data : cls
+            Data is returned as type cls with subsets applied.
+
+        """
+        return self._get_data(data_label=data_label, spatial_subset=spatial_subset, cls=cls)
+
 
 def split_filename_with_fits_ext(filename):
     """Split a ``filename[ext]`` input into filename and FITS extension.

jdaviz/configs/mosviz/helper.py

@@ -1076,3 +1076,25 @@ def get_spectrum_2d(self, row=None, apply_slider_redshift="Warn"):
         `~specutils.Spectrum1D`
         """
         return self._get_spectrum('2D Spectra', row, apply_slider_redshift)
+
+    def get_data(self, data_label=None, spectral_subset=None, cls=None):
+        """
+        Returns data with name equal to data_label of type cls with subsets applied from
+        spectral_subset.
+
+        Parameters
+        ----------
+        data_label : str, optional
+            Provide a label to retrieve a specific data set from data_collection.
+        spectral_subset : str, optional
+            Spectral subset applied to data.
+        cls : `~specutils.Spectrum1D`, `~astropy.nddata.CCDData`, optional
+            The type that data will be returned as.
+
+        Returns
+        -------
+        data : cls
+            Data is returned as type cls with subsets applied.
+
+        """
+        return self._get_data(data_label=data_label, spectral_subset=spectral_subset, cls=cls)

jdaviz/configs/specviz/helper.py

@@ -1,6 +1,7 @@
 import warnings
 
 from astropy import units as u
+from regions.core.core import Region
 from glue.core.subset_group import GroupedSubset
 from specutils import SpectralRegion, Spectrum1D
 
@@ -76,10 +77,11 @@ def get_spectra(self, data_label=None, subset_to_apply=None, apply_slider_redshi
         get_data_method = self.app._jdaviz_helper.get_data
         viewer = self.app.get_viewer(self._default_spectrum_viewer_reference_name)
         function_kwargs = {'function': getattr(viewer.state, "function")} if self.app.config == 'cubeviz' else {}  # noqa
+        all_subsets = self.app.get_subsets(object_only=True)
 
         if data_label is not None:
             spectrum = get_data_method(data_label=data_label,
-                                       subset_to_apply=subset_to_apply,
+                                       spectral_subset=subset_to_apply,
                                        cls=Spectrum1D)
             spectra[data_label] = spectrum
         else:
@@ -88,16 +90,24 @@ def get_spectra(self, data_label=None, subset_to_apply=None, apply_slider_redshi
                 if subset_to_apply is not None:
                     if lyr.label == subset_to_apply:
                         spectrum = get_data_method(data_label=lyr.data.label,
-                                                   subset_to_apply=subset_to_apply,
+                                                   spectral_subset=subset_to_apply,
                                                    cls=Spectrum1D,
                                                    **function_kwargs)
                         spectra[lyr.data.label] = spectrum
                     else:
                         continue
                 else:
-                    if isinstance(lyr, GroupedSubset):
+                    if (isinstance(lyr, GroupedSubset) and lyr.label in all_subsets.keys() and
+                            isinstance(all_subsets[lyr.label][0], Region)):
                         spectrum = get_data_method(data_label=lyr.data.label,
-                                                   subset_to_apply=lyr.label,
+                                                   spatial_subset=lyr.label,
+                                                   cls=Spectrum1D,
+                                                   **function_kwargs)
+                        spectra[f'{lyr.data.label} ({lyr.label})'] = spectrum
+                    elif (isinstance(lyr, GroupedSubset) and lyr.label in all_subsets.keys() and
+                            isinstance(all_subsets[lyr.label], SpectralRegion)):
+                        spectrum = get_data_method(data_label=lyr.data.label,
+                                                   spectral_subset=lyr.label,
                                                    cls=Spectrum1D,
                                                    **function_kwargs)
                         spectra[f'{lyr.data.label} ({lyr.label})'] = spectrum
@@ -271,7 +281,7 @@ def set_spectrum_tick_format(self, fmt, axis=None):
             self._default_spectrum_viewer_reference_name
         ).figure.axes[axis].tick_format = fmt
 
-    def get_data(self, data_label=None, cls=None, subset_to_apply=None):
+    def get_data(self, data_label=None, spectral_subset=None, cls=None, **kwargs):
         """
         Returns data with name equal to data_label of type cls with subsets applied from
         subset_to_apply.
@@ -280,26 +290,36 @@ def get_data(self, data_label=None, cls=None, subset_to_apply=None):
         ----------
         data_label : str, optional
             Provide a label to retrieve a specific data set from data_collection.
+        spectral_subset : str, optional
+            Spectral subset applied to data.
         cls : `~specutils.Spectrum1D`, optional
             The type that data will be returned as.
-        subset_to_apply : str, optional
-            Subset that is to be applied to data before it is returned.
 
         Returns
         -------
         data : cls
             Data is returned as type cls with subsets applied.
 
         """
+        spatial_subset = kwargs.pop("spatial_subset", None)
+        function = kwargs.pop("function", None)
+        if len(kwargs) > 0:
+            raise ValueError(f'kwargs {[x for x in kwargs.keys()]} are not valid')
+
         if self.app.config == 'cubeviz':
             # then this is a specviz instance inside cubeviz and we want to default to the
             # viewer's collapse function
             default_sp_viewer = self.app.get_viewer(self._default_spectrum_viewer_reference_name)
-            function = getattr(default_sp_viewer.state, 'function', None)
+            if function is True or function is None:
+                function = getattr(default_sp_viewer.state, 'function', None)
+
             if cls is None:
                 cls = Spectrum1D
+        elif spatial_subset or function:
+            raise ValueError('kwargs spatial subset and function are not valid in specviz')
         else:
+            spatial_subset = None
             function = None
 
-        return self._get_data(data_label=data_label, cls=cls, subset_to_apply=subset_to_apply,
-                              function=function)
+        return self._get_data(data_label=data_label, spatial_subset=spatial_subset,
+                              spectral_subset=spectral_subset, function=function, cls=cls)

jdaviz/configs/specviz/plugins/line_analysis/line_analysis.py

@@ -383,7 +383,7 @@ def _calculate_statistics(self, *args, **kwargs):
             # cube, but still apply that to the spatially-collapsed spectrum.
             continuum_mask = ~self._specviz_helper.get_data(
                 self.dataset.selected,
-                subset_to_apply=self.continuum_subset_selected).mask
+                spectral_subset=self.continuum_subset_selected).mask
             spectral_axis_nanmasked = spectral_axis.value.copy()
             spectral_axis_nanmasked[~continuum_mask] = np.nan
             if self.spectral_subset_selected == "Entire Spectrum":

jdaviz/configs/specviz2d/helper.py

@@ -238,3 +238,25 @@ def load_trace(self, trace, data_label, show_in_viewer=True):
             self.app.add_data_to_viewer(
                 self._default_spectrum_2d_viewer_reference_name, data_label
             )
+
+    def get_data(self, data_label=None, spectral_subset=None, cls=None):
+        """
+        Returns data with name equal to data_label of type cls with subsets applied from
+        spectral_subset.
+
+        Parameters
+        ----------
+        data_label : str, optional
+            Provide a label to retrieve a specific data set from data_collection.
+        spectral_subset : str, optional
+            Spectral subset applied to data.
+        cls : `~specutils.Spectrum1D`, `~astropy.nddata.CCDData`, optional
+            The type that data will be returned as.
+
+        Returns
+        -------
+        data : cls
+            Data is returned as type cls with subsets applied.
+
+        """
+        return self._get_data(data_label=data_label, spectral_subset=spectral_subset, cls=cls)