Respect loaded mask cube properly in spectral extraction.

CHANGES.rst

@@ -32,7 +32,7 @@ Specviz2d
 API Changes
 -----------
 - Removed API access to plugins that have passed the deprecation period: Links Control, Canvas Rotation, Export Plot. [#3270]
-  
+
 - Subset Tools plugin now exposes the ``subset``, ``combination_mode``, ``recenter_dataset``,
   ``recenter``, ``get_center``, and ``set_center`` in the user API. [#3293, #3304, #3325]
 
@@ -114,6 +114,8 @@ Cubeviz
 
 - Fixed initializing a Gaussian1D model component when ``Cube Fit`` is toggled on. [#3295]
 
+- Spectral extraction now correctly respects the loaded mask cube. [#3319]
+
 Imviz
 ^^^^^
 

docs/cubeviz/plugins.rst

@@ -87,7 +87,7 @@ The slice plugin provides the ability to select the slice
 of the cube currently visible in the image viewers, with the
 corresponding wavelength highlighted in the spectrum viewer.
 
-To choose a specific slice, enter an approximate wavelength (in which case the nearest slice will 
+To choose a specific slice, enter an approximate wavelength (in which case the nearest slice will
 be selected and the wavelength entry will "span" to the exact value of that slice).  The snapping
 behavior can be disabled in the plugin settings to allow for smooth scrubbing, in which case the
 closest slice will still be displayed in the cube viewer.
@@ -301,6 +301,9 @@ optionally choose a :guilabel:`Spatial region`, if you have one.
 Click :guilabel:`EXTRACT` to produce a new 1D spectrum dataset
 from the spectral cube, which has uncertainties propagated by
 `astropy.nddata <https://docs.astropy.org/en/stable/nddata/nddata.html>`_.
+By default, if a mask was loaded with the cube, it will be applied to the
+cube when extracting in addition to any subsets chosen as an aperture. This
+is not currently done for Data Quality arrays, e.g. the DQ extension in JWST files.
 
 If using a simple subset (currently only works for a circular subset applied to data
 with spatial axis units in wavelength) for the spatial aperture, an option to

jdaviz/configs/cubeviz/helper.py

@@ -22,6 +22,7 @@ class Cubeviz(CubeConfigHelper, LineListMixin):
 
     _loaded_flux_cube = None
     _loaded_uncert_cube = None
+    _loaded_mask_cube = None
     _cube_viewer_cls = CubevizImageView
 
     def __init__(self, *args, **kwargs):

jdaviz/configs/cubeviz/plugins/parsers.py

@@ -307,8 +307,9 @@
         app.add_data(sc, data_label)
 
         if data_type == 'mask':
-            # We no longer auto-populate the mask cube into a viewer
-            pass
+            # We no longer auto-populate the mask cube into a viewer, but we still want
+            # to keep track of this cube for use in, e.g., spectral extraction.
+            app._jdaviz_helper._loaded_mask_cube = app.data_collection[data_label]
 
         elif data_type == 'uncert':
             app.add_data_to_viewer(uncert_viewer_reference_name, data_label)
@@ -432,8 +433,10 @@
 
     if data_type == 'flux':
         app._jdaviz_helper._loaded_flux_cube = app.data_collection[data_label]
-    if data_type == 'uncert':
+    elif data_type == 'uncert':
         app._jdaviz_helper._loaded_uncert_cube = app.data_collection[data_label]
+    elif data_type == 'mask':
+        app._jdaviz_helper._loaded_mask_cube = app.data_collection[data_label]
 
 
 def _parse_spectrum1d_3d(app, file_obj, data_label=None,
@@ -483,7 +486,8 @@
         elif attr == 'uncertainty':
             app.add_data_to_viewer(uncert_viewer_reference_name, cur_data_label)
             app._jdaviz_helper._loaded_uncert_cube = app.data_collection[cur_data_label]
-        # We no longer auto-populate the mask cube into a viewer
+        elif attr == 'mask':
+            app._jdaviz_helper._loaded_mask_cube = app.data_collection[cur_data_label]
 
 
 def _parse_spectrum1d(app, file_obj, data_label=None, spectrum_viewer_reference_name=None):
@@ -542,6 +546,8 @@
     elif data_type == 'uncert':
         app.add_data_to_viewer(uncert_viewer_reference_name, data_label)
         app._jdaviz_helper._loaded_uncert_cube = app.data_collection[data_label]
+    elif data_type == 'mask':
+        app._jdaviz_helper._loaded_mask_cube = app.data_collection[data_label]
 
 
 def _parse_gif(app, file_obj, data_label=None, flux_viewer_reference_name=None):  # pragma: no cover

jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.py

@@ -371,6 +371,16 @@
             # TODO: allow selecting or associating an uncertainty cube?
             return None
 
+    @property
+    def mask_cube(self):
+        if (hasattr(self._app._jdaviz_helper, '_loaded_flux_cube') and
+                hasattr(self.app._jdaviz_helper, '_loaded_mask_cube') and
+                self.dataset.selected == self._app._jdaviz_helper._loaded_flux_cube.label):
+            return self._app._jdaviz_helper._loaded_mask_cube
+        else:
+            # TODO: allow selecting or associating a mask/DQ cube?
+            return None
+
     @property
     def slice_display_unit(self):
         return astropy.units.Unit(self.app._get_display_unit(self.slice_display_unit_name))
@@ -431,7 +441,7 @@
     def bg_area_along_spectral(self):
         return np.sum(self.bg_weight_mask, axis=self.spatial_axes)
 
-    def _extract_from_aperture(self, cube, uncert_cube, aperture,
+    def _extract_from_aperture(self, cube, uncert_cube, mask_cube, aperture,
                                weight_mask, wavelength_dependent,
                                selected_func, **kwargs):
         # This plugin collapses over the *spatial axes* (optionally over a spatial subset,
@@ -487,6 +497,18 @@
         # Filter out NaNs (False = good)
         mask = np.logical_or(mask, np.isnan(flux))
 
+        # Also apply the cube's original mask array
+        if mask_cube:
+            snackbar_message = SnackbarMessage(
+                "Note: Applied loaded mask cube during extraction",
+                color="warning",
+                sender=self)
+            self.hub.broadcast(snackbar_message)
+            mask_from_cube = mask_cube.get_component('flux').data.copy()
+            # Some mask cubes have NaNs where they are not masked instead of 0
+            mask_from_cube[np.where(np.isnan(mask_from_cube))] = 0
+            mask = np.logical_or(mask, mask_from_cube.astype('bool'))
+
         nddata_reshaped = NDDataArray(
             flux, mask=mask, uncertainty=uncertainties, wcs=wcs, meta=nddata.meta
         )
@@ -600,7 +622,7 @@
             raise ValueError("aperture and background cannot be set to the same subset")
 
         selected_func = self.function_selected.lower()
-        spec = self._extract_from_aperture(self.cube, self.uncert_cube,
+        spec = self._extract_from_aperture(self.cube, self.uncert_cube, self.mask_cube,
                                            self.aperture, self.aperture_weight_mask,
                                            self.wavelength_dependent,
                                            selected_func, **kwargs)
@@ -654,7 +676,7 @@
         # allow internal calls to override the behavior of the bg_spec_per_spaxel traitlet
         bg_spec_per_spaxel = kwargs.pop('bg_spec_per_spaxel', self.bg_spec_per_spaxel)
         if self.background.selected != self.background.default_text:
-            bg_spec = self._extract_from_aperture(self.cube, self.uncert_cube,
+            bg_spec = self._extract_from_aperture(self.cube, self.uncert_cube, self.mask_cube,
                                                   self.background, self.bg_weight_mask,
                                                   self.bg_wavelength_dependent,
                                                   self.function_selected.lower(), **kwargs)

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

@@ -1,3 +1,5 @@
+import warnings
+
 import numpy as np
 import pytest
 from astropy import units as u
@@ -207,6 +209,25 @@ def test_numpy_cube(cubeviz_helper):
     assert flux.units == 'ct / pix2'
 
 
+@pytest.mark.remote_data
+def test_manga_cube(cubeviz_helper):
+    # Remote data test of loading and extracting an up-to-date (as of 11/19/2024) MaNGA cube
+    # This also tests that spaxel is converted to pix**2
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore")
+        cubeviz_helper.load_data("https://stsci.box.com/shared/static/gts87zqt5265msuwi4w5u003b6typ6h0.gz", cache=True)  # noqa
+
+    uc = cubeviz_helper.plugins['Unit Conversion']
+    uc.spectral_y_type = "Surface Brightness"
+
+    se = cubeviz_helper.plugins['Spectral Extraction']
+    se.function = "Mean"
+    se.extract()
+    extracted_max = cubeviz_helper.get_data("Spectrum (mean)").max()
+    assert_allclose(extracted_max.value, 2.836957E-18)
+    assert extracted_max.unit == u.Unit("erg / Angstrom s cm**2 pix**2")
+
+
 def test_invalid_data_types(cubeviz_helper):
     with pytest.raises(ValueError, match=r"The input file 'does_not_exist\.fits'"):
         cubeviz_helper.load_data('does_not_exist.fits')