Respect loaded mask cube properly in spectral extraction. (#3319)

* Track loaded mask cube in cubeviz

* First pass at handling mask cube in extraction

* Fix type

* Add changelog

* Handle mask cubes that have NaNs instead of 0s

* Act on copy of this mask array

* Don't save DQ array as loaded_mask for JWST

* Add snackbar warning and note in docs

* codestyle

* Add extraction test on file with mask

* Codestyle

(cherry picked from commit bf8cb26)

CHANGES.rst

@@ -29,6 +29,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

@@ -306,8 +306,9 @@ def _parse_hdulist(app, hdulist, file_name=None,
         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)
@@ -429,8 +430,10 @@ def _parse_esa_s3d(app, hdulist, data_label, ext='DATA', flux_viewer_reference_n
 
     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,
@@ -482,7 +485,8 @@ def _parse_spectrum1d_3d(app, file_obj, data_label=None,
         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):
@@ -540,6 +544,8 @@ def _parse_ndarray(app, file_obj, data_label=None, data_type=None,
     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

@@ -370,6 +370,16 @@ def uncert_cube(self):
             # 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))
@@ -430,7 +440,7 @@ def aperture_area_along_spectral(self):
     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,
@@ -486,6 +496,18 @@ def _extract_from_aperture(self, cube, uncert_cube, aperture,
         # 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
         )
@@ -588,7 +610,7 @@ def extract(self, return_bg=False, add_data=True, **kwargs):
             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)
@@ -642,7 +664,7 @@ def extract_bg_spectrum(self, add_data=False, **kwargs):
         # 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
@@ -205,6 +207,25 @@ def test_numpy_cube(cubeviz_helper):
     assert flux.units == 'ct'
 
 
+@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')