Add masked subset support (spacetelescope#2462)

* Integrate MultiMaskSubsetState into get_subsets and subset_plugin

---------

Co-authored-by: Kyle Conroy <[email protected]>

CHANGES.rst

@@ -10,6 +10,9 @@ New Features
 
 - Data color cycler and marker color updates for increased accessibility. [#2453]
 
+- Add support for ``MultiMaskSubsetState`` in ``viz.app.get_subsets()`` and in
+  the Subset Plugin [#2462]
+
 Cubeviz
 ^^^^^^^
 

jdaviz/app.py

@@ -56,7 +56,7 @@
 from jdaviz.core.registries import (tool_registry, tray_registry, viewer_registry,
                                     data_parser_registry)
 from jdaviz.core.tools import ICON_DIR
-from jdaviz.utils import SnackbarQueue, alpha_index
+from jdaviz.utils import SnackbarQueue, alpha_index, MultiMaskSubsetState
 from ipypopout import PopoutButton
 
 __all__ = ['Application', 'ALL_JDAVIZ_CONFIGS']
@@ -1066,10 +1066,15 @@ def get_subsets(self, subset_name=None, spectral_only=False,
                 subset_region = self._get_range_subset_bounds(subset.subset_state,
                                                               simplify_spectral,
                                                               use_display_units)
+            elif isinstance(subset.subset_state, MultiMaskSubsetState):
+                subset_region = self._get_multi_mask_subset_definition(subset.subset_state)
             else:
-                # subset.subset_state can be an instance of MaskSubsetState
-                # or something else we do not know how to handle
-                all_subsets[label] = None
+                # subset.subset_state can be an instance of something else
+                # we do not know how to handle yet
+                all_subsets[label] = [{"name": subset.subset_state.__class__.__name__,
+                                       "glue_state": subset.subset_state.__class__.__name__,
+                                       "region": None,
+                                       "subset_state": subset.subset_state}]
                 continue
 
             # Is the subset spectral, spatial, temporal?
@@ -1091,6 +1096,7 @@ def get_subsets(self, subset_name=None, spectral_only=False,
                 else:
                     all_subsets[label] = subset_region
             elif not spectral_only and not spatial_only:
+                # General else statement if no type was specified
                 if object_only and not isinstance(subset_region, SpectralRegion):
                     all_subsets[label] = [reg['region'] for reg in subset_region]
                 else:
@@ -1162,6 +1168,12 @@ def _get_range_subset_bounds(self, subset_state,
                      "subset_state": subset_state}]
         return spec_region
 
+    def _get_multi_mask_subset_definition(self, subset_state):
+        return [{"name": subset_state.__class__.__name__,
+                 "glue_state": subset_state.__class__.__name__,
+                 "region": subset_state.total_masked_first_data(),
+                 "subset_state": subset_state}]
+
     def _get_roi_subset_definition(self, subset_state):
         # TODO: Imviz: Return sky region if link type is WCS.
         roi_as_region = roi_subset_state_to_region(subset_state)
@@ -1182,6 +1194,15 @@ def get_sub_regions(self, subset_state, simplify_spectral=True, use_display_unit
                 if isinstance(one, list) and "glue_state" in one[0]:
                     one[0]["glue_state"] = subset_state.__class__.__name__
 
+                if (isinstance(one, list)
+                        and isinstance(one[0]["subset_state"], MultiMaskSubsetState)
+                        and simplify_spectral):
+                    return two
+                elif (isinstance(two, list)
+                        and isinstance(two[0]["subset_state"], MultiMaskSubsetState)
+                        and simplify_spectral):
+                    return one
+
                 if isinstance(subset_state.state2, InvertState):
                     # This covers the REMOVE subset mode
 
@@ -1300,13 +1321,15 @@ def get_sub_regions(self, subset_state, simplify_spectral=True, use_display_unit
                                             simplify_spectral, use_display_units)
         elif subset_state is not None:
             # This is the leaf node of the glue subset state tree where
-            # a subset_state is either ROI or Range.
+            # a subset_state is either ROI, Range, or MultiMask.
             if isinstance(subset_state, RoiSubsetState):
                 return self._get_roi_subset_definition(subset_state)
 
             elif isinstance(subset_state, RangeSubsetState):
                 return self._get_range_subset_bounds(subset_state,
                                                      simplify_spectral, use_display_units)
+            elif isinstance(subset_state, MultiMaskSubsetState):
+                return self._get_multi_mask_subset_definition(subset_state)
 
     def _get_display_unit(self, axis):
         if self._jdaviz_helper is None or self._jdaviz_helper.plugins.get('Unit Conversion') is None:  # noqa

jdaviz/configs/default/plugins/subset_plugin/subset_plugin.py

@@ -18,6 +18,7 @@
 from jdaviz.core.registries import tray_registry
 from jdaviz.core.template_mixin import PluginTemplateMixin, DatasetSelectMixin, SubsetSelect
 from jdaviz.core.tools import ICON_DIR
+from jdaviz.utils import MultiMaskSubsetState
 
 __all__ = ['SubsetPlugin']
 
@@ -51,6 +52,7 @@ class SubsetPlugin(PluginTemplateMixin, DatasetSelectMixin):
     multiselect = Bool(False).tag(sync=True)
     is_centerable = Bool(False).tag(sync=True)
     can_simplify = Bool(False).tag(sync=True)
+    can_freeze = Bool(False).tag(sync=True)
 
     icon_replace = Unicode(read_icon(os.path.join(icon_path("glue_replace", icon_format="svg")), 'svg+xml')).tag(sync=True)  # noqa
     icon_or = Unicode(read_icon(os.path.join(icon_path("glue_or", icon_format="svg")), 'svg+xml')).tag(sync=True)  # noqa
@@ -227,6 +229,13 @@ def _unpack_get_subsets_for_ui(self):
                                          {"name": "Upper bound", "att": "hi", "value": hi.value,
                                           "orig": hi.value, "unit": str(hi.unit)}]
                 subset_type = "Range"
+
+            elif isinstance(subset_state, MultiMaskSubsetState):
+                total_masked = subset_state.total_masked_first_data()
+                subset_definition = [{"name": "Masked values", "att": "masked",
+                                      "value": total_masked,
+                                      "orig": total_masked}]
+                subset_type = "Mask"
             if len(subset_definition) > 0:
                 # Note: .append() does not work for List traitlet.
                 self.subset_definitions = self.subset_definitions + [subset_definition]
@@ -236,8 +245,11 @@ def _unpack_get_subsets_for_ui(self):
 
         simplifiable_states = set(['AndState', 'XorState', 'AndNotState'])
         # Check if the subset has more than one subregion, is a range subset type, and
-        # uses one of the states that can be simplified.
-        if (len(self.subset_states) > 1 and isinstance(self.subset_states[0], RangeSubsetState)
+        # uses one of the states that can be simplified. Mask subset types cannot be simplified
+        # so subsets contained that are skipped.
+        if 'Mask' in self.subset_types:
+            self.can_simplify = False
+        elif (len(self.subset_states) > 1 and isinstance(self.subset_states[0], RangeSubsetState)
                 and len(simplifiable_states - set(self.glue_state_types)) < 3):
             self.can_simplify = True
         elif (len(self.subset_states) > 1 and isinstance(self.subset_states[0], RangeSubsetState)
@@ -258,6 +270,16 @@ def _get_subset_definition(self, *args):
 
         self._unpack_get_subsets_for_ui()
 
+    def vue_freeze_subset(self, *args):
+        sgs = {sg.label: sg for sg in self.app.data_collection.subset_groups}
+        sg = sgs.get(self.subset_selected)
+
+        masks = {}
+        for data in self.app.data_collection:
+            masks[data.uuid] = sg.subset_state.to_mask(data)
+
+        sg.subset_state = MultiMaskSubsetState(masks)
+
     def vue_simplify_subset(self, *args):
         if self.multiselect:
             self.hub.broadcast(SnackbarMessage("Cannot simplify spectral subset "

jdaviz/configs/default/plugins/subset_plugin/subset_plugin.vue

@@ -95,14 +95,14 @@
       </v-row>
 
       <v-row v-for="(item, index2) in region" class="row-no-outside-padding">
-        <v-text-field v-if="item.name === 'Parent'"
+        <v-text-field v-if="item.name === 'Parent' || item.name === 'Masked values'"
           :label="item.name"
           :value="item.value"
           style="padding-top: 0px; margin-top: 0px"
           :readonly="true"
-          hint="Subset was defined with respect to this reference data (read-only)"
+          :hint="item.name === 'Parent' ? 'Subset was defined with respect to this reference data (read-only)' : 'Number of elements included by mask'"
         ></v-text-field>
-        <v-text-field v-if="item.name !== 'Parent'"
+        <v-text-field v-if="item.name !== 'Parent' && item.name !== 'Masked values'"
           :label="item.name"
           v-model.number="item.value"
           type="number"
@@ -112,11 +112,14 @@
       </v-row>
     </div>
 
-    <v-row v-if="!multiselect" justify="end">
-      <j-tooltip v-if="can_simplify" tooltipcontent="Convert composite subset to use only add mode to connect subregions">
-        <v-btn color="primary" text @click="simplify_subset">Simplify</v-btn>
-      </j-tooltip>
-      <v-btn color="primary" text @click="update_subset">Update</v-btn>
-    </v-row>
+      <v-row v-if="!multiselect" justify="end" no-gutters>
+        <j-tooltip v-if="can_freeze" tooltipcontent="Freeze subset to a mask on the underlying data entries">
+          <v-btn color="primary" text @click="freeze_subset">Freeze</v-btn>
+        </j-tooltip>
+        <j-tooltip tooltipcontent="Convert composite subset to use only add mode to connect subregions">
+          <v-btn :disabled="!can_simplify" color="primary" text @click="simplify_subset">Simplify</v-btn>
+        </j-tooltip>
+        <v-btn color="primary" text @click="update_subset">Update</v-btn>
+      </v-row>
   </j-tray-plugin>
 </template>

jdaviz/tests/test_subsets.py

@@ -14,7 +14,7 @@
 from specutils import SpectralRegion, Spectrum1D
 
 from jdaviz.core.marks import ShadowSpatialSpectral
-from jdaviz.utils import get_subset_type
+from jdaviz.utils import get_subset_type, MultiMaskSubsetState
 
 
 def test_region_from_subset_2d(cubeviz_helper):
@@ -766,3 +766,51 @@ def test_only_overlapping_in_specviz2d(specviz2d_helper, mos_spectrum2d):
     reg = specviz2d_helper.app.get_subsets("Subset 1")
     assert reg[0].lower.value == 6400 and reg[0].upper.value == 7400
     assert reg[1].lower.value == 7600 and reg[1].upper.value == 7800
+
+
+def test_multi_mask_subset(specviz_helper, spectrum1d):
+    specviz_helper.load_data(spectrum1d)
+    viewer = specviz_helper.app.get_viewer(specviz_helper._default_spectrum_viewer_reference_name)
+
+    viewer.apply_roi(XRangeROI(6200, 6800))
+
+    plugin = specviz_helper.app.get_tray_item_from_name("g-subset-plugin")
+    plugin.can_freeze = True
+    plugin.vue_freeze_subset()
+
+    reg = specviz_helper.app.get_subsets()
+    assert reg["Subset 1"][0]["region"] == 3
+    assert isinstance(reg["Subset 1"][0]["subset_state"], MultiMaskSubsetState)
+
+    specviz_helper.app.session.edit_subset_mode.mode = OrMode
+    viewer.apply_roi(XRangeROI(7200, 7600))
+
+    # Simplify subset ignores Mask subsets
+    reg = specviz_helper.app.get_subsets()
+    assert (reg["Subset 1"].lower.value == 7200
+            and reg["Subset 1"].upper.value == 7600)
+
+    # If we set simplify to False, we see all subregions
+    reg = specviz_helper.app.get_subsets(simplify_spectral=False)
+    assert (reg["Subset 1"][1]["region"].lower.value == 7200
+            and reg["Subset 1"][1]["region"].upper.value == 7600)
+    assert reg["Subset 1"][0]["region"] == 3
+    assert plugin.can_simplify is False
+
+    # If we freeze again, all subregions become a Mask subset object
+    plugin.vue_freeze_subset()
+    reg = specviz_helper.app.get_subsets()
+    assert reg["Subset 1"][0]["region"] == 5
+
+    # When freezing an AndNot state, the number of mask values should decrease
+    specviz_helper.app.session.edit_subset_mode.mode = AndNotMode
+    viewer.apply_roi(XRangeROI(6600, 7200))
+
+    reg = specviz_helper.app.get_subsets(simplify_spectral=False)
+    assert (reg["Subset 1"][1]["region"].lower.value == 6600
+            and reg["Subset 1"][1]["region"].upper.value == 7200)
+    assert reg["Subset 1"][0]["region"] == 5
+
+    plugin.vue_freeze_subset()
+    reg = specviz_helper.app.get_subsets()
+    assert reg["Subset 1"][0]["region"] == 4

jdaviz/utils.py

@@ -2,12 +2,14 @@
 import time
 import threading
 from collections import deque
+import numpy as np
 
 from astropy.io import fits
 from astropy.utils import minversion
 from ipyvue import watch
+from glue.core.exceptions import IncompatibleAttribute
 from glue.config import settings
-from glue.core.subset import RangeSubsetState, RoiSubsetState
+from glue.core.subset import SubsetState, RangeSubsetState, RoiSubsetState
 
 
 __all__ = ['SnackbarQueue', 'enable_hot_reloading', 'bqplot_clear_figure',
@@ -305,3 +307,45 @@ def get_subset_type(subset):
         return 'spectral'
     else:
         return None
+
+
+class MultiMaskSubsetState(SubsetState):
+    """
+    A subset state that can include a different mask for different datasets.
+    Adopted from https://github.com/glue-viz/glue/pull/2415
+
+    Parameters
+    ----------
+    masks : dict
+        A dictionary mapping data UUIDs to boolean arrays with the same
+        dimensions as the data arrays.
+    """
+
+    def __init__(self, masks=None):
+        super(MultiMaskSubsetState, self).__init__()
+        self._masks = masks
+
+    def to_mask(self, data, view=None):
+        if data.uuid in self._masks:
+            mask = self._masks[data.uuid]
+            if view is not None:
+                mask = mask[view]
+            return mask
+        else:
+            raise IncompatibleAttribute()
+
+    def copy(self):
+        return MultiMaskSubsetState(masks=self._masks)
+
+    def __gluestate__(self, context):
+        serialized = {key: context.do(value) for key, value in self._masks.items()}
+        return {'masks': serialized}
+
+    def total_masked_first_data(self):
+        first_data = next(iter(self._masks))
+        return len(np.where(self._masks[first_data])[0])
+
+    @classmethod
+    def __setgluestate__(cls, rec, context):
+        masks = {key: context.object(value) for key, value in rec['masks'].items()}
+        return cls(masks=masks)