Fix aperture weight mask with loaded mask cube (#3358)

* Track loaded mask cube in cubeviz

* Don't save DQ array as loaded_mask for JWST

* Add extraction test on file with mask

* Codestyle

* Add consideration of loaded mask to aperture weight mask

* Remove duplicate test, codestyle

* Changelog

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

Co-authored-by: P. L. Lim <[email protected]>

* Cast this as bool here

* Fix tests, address review comments

* adding regression tests to manga mask

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

Co-authored-by: Brett M. Morris <[email protected]>

---------

Co-authored-by: P. L. Lim <[email protected]>
Co-authored-by: Brett M. Morris <[email protected]>
Co-authored-by: Brett M. Morris <[email protected]>

CHANGES.rst

@@ -120,7 +120,7 @@ Cubeviz
 
 - Fixed initializing a Gaussian1D model component when ``Cube Fit`` is toggled on. [#3295]
 
-- Spectral extraction now correctly respects the loaded mask cube. [#3319]
+- Spectral extraction now correctly respects the loaded mask cube. [#3319, #3358]
 
 Imviz
 ^^^^^

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

@@ -386,14 +386,17 @@ def slice_display_unit(self):
         return astropy.units.Unit(self.app._get_display_unit(self.slice_display_unit_name))
 
     @property
-    def mask_non_science(self):
+    def inverted_mask_non_science(self):
         # Aperture masks begin by removing from consideration any pixel
         # set to NaN, which corresponds to a pixel on the "non-science" portions
         # of the detector. For JWST spectral cubes, these pixels are also marked in
-        # the DQ array with flag `513`.
-        return np.logical_not(
-            np.isnan(self.dataset.selected_obj.flux.value)
-        ).astype(float)
+        # the DQ array with flag `513`. Also respect the loaded mask, if it exists.
+        # This "inverted mask" is `True` where the data are included, `False` where excluded.
+        mask_non_science = np.isnan(self.dataset.selected_obj.flux.value)
+        if self.mask_cube is not None:
+            mask_non_science = np.logical_or(self.mask_cube.get_component('flux').data,
+                                             mask_non_science)
+        return np.logical_not(mask_non_science).astype(float)
 
     @property
     def aperture_weight_mask(self):
@@ -402,10 +405,10 @@ def aperture_weight_mask(self):
         # wavelength, on the range [0, 1].
         if self.aperture.selected == self.aperture.default_text:
             # Entire Cube
-            return self.mask_non_science
+            return self.inverted_mask_non_science
 
         return (
-            self.mask_non_science *
+            self.inverted_mask_non_science *
             self.aperture.get_mask(
                 self.dataset.selected_obj,
                 self.aperture_method_selected,
@@ -421,7 +424,7 @@ def bg_weight_mask(self):
             return np.zeros_like(self.dataset.selected_obj.flux.value)
 
         return (
-            self.mask_non_science *
+            self.inverted_mask_non_science *
             self.background.get_mask(
                 self.dataset.selected_obj,
                 self.aperture_method_selected,

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

@@ -209,8 +209,34 @@ def test_numpy_cube(cubeviz_helper):
     assert flux.units == 'ct / pix2'
 
 
+def test_loading_with_mask(cubeviz_helper):
+    # This also tests that spaxel is converted to pix**2
+    custom_spec = Spectrum1D(flux=[[[20, 1], [9, 1]], [[3, 1], [6, np.nan]]] * u.Unit("erg / Angstrom s cm**2 spaxel"),  # noqa
+                             spectral_axis=[1, 2]*u.AA,
+                             mask=[[[1, 0], [0, 0]], [[0, 0], [0, 0]]])
+    cubeviz_helper.load_data(custom_spec)
+
+    uc = cubeviz_helper.plugins['Unit Conversion']
+    uc.spectral_y_type = "Surface Brightness"
+
+    se = cubeviz_helper.plugins['Spectral Extraction']
+    se.function = "Mean"
+    se.extract()
+    extracted = cubeviz_helper.get_data("Spectrum (mean)")
+    assert_allclose(extracted.flux.value, [6, 1])
+    assert extracted.unit == u.Unit("erg / Angstrom s cm**2 pix**2")
+
+
 @pytest.mark.remote_data
-def test_manga_cube(cubeviz_helper):
+@pytest.mark.parametrize(
+    'function, expected_value,',
+    (
+        ('Mean', 5.566169e-18),
+        ('Sum', 1.553518e-14),
+        ('Max', 1e20),
+    )
+)
+def test_manga_with_mask(cubeviz_helper, function, expected_value):
     # 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():
@@ -221,11 +247,14 @@ def test_manga_cube(cubeviz_helper):
     uc.spectral_y_type = "Surface Brightness"
 
     se = cubeviz_helper.plugins['Spectral Extraction']
-    se.function = "Mean"
+    se.function = function
     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")
+    extracted_max = cubeviz_helper.get_data(f"Spectrum ({function.lower()})").max()
+    assert_allclose(extracted_max.value, expected_value, rtol=5e-7)
+    if function == "Sum":
+        assert extracted_max.unit == u.Unit("erg / Angstrom s cm**2")
+    else:
+        assert extracted_max.unit == u.Unit("erg / Angstrom s cm**2 pix**2")
 
 
 def test_invalid_data_types(cubeviz_helper):