From 1934515d91817965e6e7ca848827ab1ab2281ad7 Mon Sep 17 00:00:00 2001 From: Ricky O'Steen Date: Mon, 25 Nov 2024 14:49:19 -0500 Subject: [PATCH] First pass at handling mask cube in extraction --- .../spectral_extraction.py | 20 ++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.py b/jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.py index 913faf818b..90db9ff4d9 100644 --- a/jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.py +++ b/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)