First pass at handling mask cube in extraction

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,10 @@ 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:
+            mask = np.logical_or(mask, mask_cube.get_component('flux').data.astype('bool'))
+
         nddata_reshaped = NDDataArray(
             flux, mask=mask, uncertainty=uncertainties, wcs=wcs, meta=nddata.meta
         )
@@ -588,7 +602,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 +656,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)