Spectral extraction in Cubeviz with composite subsets

CHANGES.rst

@@ -13,6 +13,8 @@ Cubeviz
 
 - Enable spectral unit conversion in cubeviz. [#2758, #2803]
 
+- Enable spectral extraction for composite subsets. [#2837]
+
 Imviz
 ^^^^^
 

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

@@ -288,6 +288,11 @@ def collapse_to_spectrum(self, add_data=True, **kwargs):
             # Boolean cube which is True outside of the aperture
             # (i.e., the numpy boolean mask convention)
             mask = np.isclose(shape_mask, 0)
+
+            # composite subset masks are in `nddata.mask`:
+            if nddata.mask is not None and np.all(shape_mask == 0):
+                mask &= nddata.mask
+
         else:
             nddata = spectral_cube.get_object(cls=NDDataArray)
             if uncert_cube:
@@ -377,6 +382,14 @@ def get_aperture(self):
                                                                           statistic=None)
         display_unit = astropy.units.Unit(self.app._get_display_unit(self.slice_display_unit_name))
 
+        # if subset is a composite subset, skip the other logic:
+        if self.aperture.is_composite:
+            [subset_group] = [
+                subset_group for subset_group in self.app.data_collection.subset_groups
+                if subset_group.label == self.aperture_selected]
+            mask_weights = subset_group.subsets[0].to_mask().astype(np.float32)
+            return mask_weights
+
         # Center is reverse coordinates
         center = (self.aperture.selected_spatial_region.center.y,
                   self.aperture.selected_spatial_region.center.x)

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

@@ -6,8 +6,8 @@
 from astropy import units as u
 from astropy.nddata import NDDataArray, StdDevUncertainty
 from astropy.utils.exceptions import AstropyUserWarning
-from glue.core.roi import CircularROI
-from glue.core.edit_subset_mode import ReplaceMode
+from glue.core.roi import CircularROI, RectangularROI
+from glue.core.edit_subset_mode import ReplaceMode, AndNotMode, NewMode
 from numpy.testing import assert_allclose, assert_array_equal
 from regions import (CirclePixelRegion, CircleAnnulusPixelRegion, EllipsePixelRegion,
                      RectanglePixelRegion, PixCoord)
@@ -447,3 +447,75 @@ def test_autoupdate_results(cubeviz_helper, spectrum1d_cube_largest):
     # cover to make sure the logic does not crash
 #    new_med_flux = np.median(cubeviz_helper.get_data('extracted').flux)
 #    assert new_med_flux > orig_med_flux
+
+
+def test_aperture_composite_detection(cubeviz_helper, spectrum1d_cube):
+    cubeviz_helper.load_data(spectrum1d_cube)
+    flux_viewer = cubeviz_helper.app.get_viewer('flux-viewer')
+    subset_plugin = cubeviz_helper.plugins['Subset Tools']._obj
+    spec_extr_plugin = cubeviz_helper.plugins['Spectral Extraction']._obj
+
+    # create a rectangular subset with all spaxels:
+    rectangle = RectangularROI(-0.5, 1.5, -0.5, 3.5)
+    flux_viewer.toolbar.active_tool = flux_viewer.toolbar.tools['bqplot:rectangle']
+    flux_viewer.apply_roi(rectangle)
+
+    # select subset 1, ensure it's not a composite subset:
+    spec_extr_plugin.aperture_selected = 'Subset 1'
+    assert not spec_extr_plugin.aperture.is_composite
+
+    # now remove from this subset a circular region in the center:
+    flux_viewer.toolbar.active_tool = flux_viewer.toolbar.tools['bqplot:truecircle']
+    subset_plugin.subset_selected = 'Subset 1'
+    cubeviz_helper.app.session.edit_subset_mode.mode = AndNotMode
+
+    circle = CircularROI(0.5, 1.5, 1)
+    flux_viewer.apply_roi(circle)
+
+    # now the subset is composite:
+    assert spec_extr_plugin.aperture.is_composite
+
+
+def test_extraction_composite_subset(cubeviz_helper, spectrum1d_cube):
+    cubeviz_helper.load_data(spectrum1d_cube)
+
+    flux_viewer = cubeviz_helper.app.get_viewer('flux-viewer')
+    subset_plugin = cubeviz_helper.plugins['Subset Tools']._obj
+    spec_extr_plugin = cubeviz_helper.plugins['Spectral Extraction']._obj
+
+    lower_aperture = RectangularROI(-0.5, 1.5, -0.5, 0.5)
+    upper_aperture = RectangularROI(-0.5, 1.5, 2.5, 3.5)
+
+    flux_viewer.toolbar.active_tool = flux_viewer.toolbar.tools['bqplot:rectangle']
+    flux_viewer.apply_roi(lower_aperture)
+
+    cubeviz_helper.app.session.edit_subset_mode.mode = NewMode
+    flux_viewer.apply_roi(upper_aperture)
+
+    spec_extr_plugin.aperture_selected = 'Subset 1'
+    spectrum_1 = spec_extr_plugin.collapse_to_spectrum()
+
+    spec_extr_plugin.aperture_selected = 'Subset 2'
+    spectrum_2 = spec_extr_plugin.collapse_to_spectrum()
+
+    subset_plugin.subset_selected = 'Create New'
+    rectangle = RectangularROI(-0.5, 1.5, -0.5, 3.5)
+    flux_viewer.toolbar.active_tool = flux_viewer.toolbar.tools['bqplot:rectangle']
+    flux_viewer.apply_roi(rectangle)
+
+    flux_viewer.toolbar.active_tool = flux_viewer.toolbar.tools['bqplot:truecircle']
+    subset_plugin.subset_selected = 'Subset 3'
+    cubeviz_helper.app.session.edit_subset_mode.mode = AndNotMode
+    circle = CircularROI(0.5, 1.5, 1.1)
+    flux_viewer.apply_roi(circle)
+
+    spec_extr_plugin.aperture_selected = 'Subset 3'
+
+    assert spec_extr_plugin.aperture.is_composite
+
+    spectrum_3 = spec_extr_plugin.collapse_to_spectrum()
+
+    np.testing.assert_allclose(
+        (spectrum_1 + spectrum_2).flux.value,
+        (spectrum_3).flux.value
+    )

jdaviz/core/template_mixin.py

@@ -2257,6 +2257,10 @@ def _set_mark_visiblities(self, visible):
     def _plugin_active_changed(self, *args):
         self._set_mark_visiblities(self.plugin.is_active)
 
+    @property
+    def is_composite(self):
+        return hasattr(self.selected_obj, '__len__') and len(self.selected_obj) > 1
+
     @property
     def image_viewers(self):
         return [viewer for viewer in self.app._viewer_store.values()