Add composite subset fixes (#2291)

* generalize handling of composite subsets when determining subset type

* Fix composite subset not applying coords info or shadow mark

* Add change

* Add code from #2287 that was accidentally removed during rebase

* Address review comments and add tests

* Update jdaviz/utils.py

Co-authored-by: Brett M. Morris <morrisbrettm@gmail.com>

---------

Co-authored-by: Kyle Conroy <kyleconroy@gmail.com>
Co-authored-by: Brett M. Morris <morrisbrettm@gmail.com>

CHANGES.rst

@@ -95,6 +95,8 @@ Cubeviz
 - Moment Map plugin now writes FITS file to working directory if no path provided
   in standalone mode. [#2264]
 
+- Fixes detection of spatial vs spectral subsets for composite subsets. [#2207, #2266, #2291]
+
 Imviz
 ^^^^^
 

jdaviz/configs/cubeviz/plugins/viewers.py

@@ -1,12 +1,12 @@
 from glue.core import BaseData
-from glue.core.subset import RoiSubsetState, RangeSubsetState
 from glue_jupyter.bqplot.image import BqplotImageView
 
 from jdaviz.core.registries import viewer_registry
 from jdaviz.core.marks import SliceIndicatorMarks, ShadowSpatialSpectral
 from jdaviz.configs.cubeviz.helper import layer_is_cube_image_data
 from jdaviz.configs.default.plugins.viewers import JdavizViewerMixin
 from jdaviz.configs.specviz.plugins.viewers import SpecvizProfileView
+from jdaviz.utils import get_subset_type
 
 __all__ = ['CubevizImageView', 'CubevizProfileView']
 
@@ -108,15 +108,23 @@ def __init__(self, *args, **kwargs):
     def _is_spatial_subset(self, layer):
         # spatial subset layers will have the same data-label as the collapsed flux cube
         ref_data_label = self.state.reference_data.label
-        return (isinstance(getattr(layer.layer, 'subset_state', None), RoiSubsetState)
+
+        subset_state = getattr(layer.layer, 'subset_state', None)
+        if subset_state is None:
+            return False
+        return (get_subset_type(subset_state) == 'spatial'
                 and layer.layer.data.label == ref_data_label)
 
     def _get_spatial_subset_layers(self):
         return [ls for ls in self.state.layers if self._is_spatial_subset(ls)]
 
     def _is_spectral_subset(self, layer):
         ref_data_label = self.layers[0].layer.data.label
-        return (isinstance(getattr(layer.layer, 'subset_state', None), RangeSubsetState)
+
+        subset_state = getattr(layer.layer, 'subset_state', None)
+        if subset_state is None:
+            return False
+        return (get_subset_type(subset_state) == 'spectral'
                 and layer.layer.data.label == ref_data_label)
 
     def _get_spectral_subset_layers(self):
@@ -152,14 +160,15 @@ def _expected_subset_layer_default(self, layer_state):
         """
         super()._expected_subset_layer_default(layer_state)
 
-        if not (self._is_spatial_subset(layer_state) or self._is_spectral_subset(layer_state)):
+        subset_type = get_subset_type(layer_state.layer)
+        if subset_type is None:
             return
 
         this_mark = self._get_marks_for_layers([layer_state])[0]
 
         new_marks = []
 
-        if isinstance(layer_state.layer.subset_state, RoiSubsetState):
+        if subset_type == 'spatial':
             layer_state.linewidth = 1
 
             # need to add marks for every intersection between THIS spatial subset and ALL spectral

jdaviz/configs/default/plugins/viewers.py

@@ -1,6 +1,5 @@
 import numpy as np
 
-from glue.core.subset import RoiSubsetState
 from glue_jupyter.bqplot.profile import BqplotProfileView
 from glue_jupyter.bqplot.image import BqplotImageView
 from glue_jupyter.bqplot.scatter.layer_artist import BqplotScatterLayerState
@@ -9,7 +8,7 @@
 from jdaviz.configs.imviz.helper import layer_is_image_data
 from jdaviz.components.toolbar_nested import NestedJupyterToolbar
 from jdaviz.core.registries import viewer_registry
-from jdaviz.utils import ColorCycler
+from jdaviz.utils import ColorCycler, get_subset_type
 
 __all__ = ['JdavizViewerMixin']
 
@@ -112,7 +111,8 @@ def _get_layer_info(layer):
             # want to include the collapse function *unless* the layer is a spectral subset
             for subset in self.jdaviz_app.data_collection.subset_groups:
                 if subset.label == layer.layer.label:
-                    if isinstance(subset.subset_state, RoiSubsetState):
+                    subset_type = get_subset_type(subset)
+                    if subset_type == 'spatial':
                         return "mdi-chart-scatter-plot", suffix
                     else:
                         return "mdi-chart-bell-curve", ""

jdaviz/configs/imviz/plugins/coords_info/coords_info.py

@@ -4,7 +4,6 @@
 
 from astropy import units as u
 from glue.core import BaseData
-from glue.core.subset import RoiSubsetState
 from glue.core.subset_group import GroupedSubset
 from glue_jupyter.bqplot.image.layer_artist import BqplotImageSubsetLayerArtist
 
@@ -17,6 +16,7 @@
 from jdaviz.core.marks import PluginScatter
 from jdaviz.core.registries import tool_registry
 from jdaviz.core.template_mixin import TemplateMixin, DatasetSelectMixin
+from jdaviz.utils import get_subset_type
 
 __all__ = ['CoordsInfo']
 
@@ -443,7 +443,11 @@ def _copy_axes_to_spectral():
                 continue
 
             if isinstance(lyr.layer, GroupedSubset):
-                if not isinstance(lyr.layer.subset_state, RoiSubsetState):
+                subset_state = getattr(lyr.layer, 'subset_state', None)
+                if subset_state is None:
+                    continue
+                subset_type = get_subset_type(subset_state)
+                if subset_type == 'spectral':
                     # then this is a SPECTRAL subset
                     continue
                 # For use later in data retrieval

jdaviz/configs/specviz/plugins/viewers.py

@@ -5,6 +5,7 @@
 from astropy import units as u
 from glue.core import BaseData
 from glue.core.subset import Subset
+from glue.core.subset_group import GroupedSubset
 from glue.config import data_translator
 from glue_jupyter.bqplot.profile import BqplotProfileView
 from glue.core.exceptions import IncompatibleAttribute
@@ -17,6 +18,7 @@
 from jdaviz.core.linelists import load_preset_linelist, get_available_linelists
 from jdaviz.core.freezable_state import FreezableProfileViewerState
 from jdaviz.configs.default.plugins.viewers import JdavizViewerMixin
+from jdaviz.utils import get_subset_type
 
 __all__ = ['SpecvizProfileView']
 
@@ -377,7 +379,9 @@ def add_data(self, data, color=None, alpha=None, **layer_state):
         # that new data entries (from model fitting or gaussian smooth, etc) will only be spectra
         # and all subsets affected will be spectral
         for layer in self.state.layers:
-            if "Subset" in layer.layer.label and layer.layer.data.label == data.label:
+            if (isinstance(layer.layer, GroupedSubset)
+                    and get_subset_type(layer.layer) == 'spectral'
+                    and layer.layer.data.label == data.label):
                 layer.linewidth = 3
 
         return result

jdaviz/core/template_mixin.py

@@ -28,6 +28,7 @@
 from jdaviz.core.events import (AddDataMessage, RemoveDataMessage,
                                 ViewerAddedMessage, ViewerRemovedMessage)
 from jdaviz.core.user_api import UserApiWrapper, PluginUserApi
+from jdaviz.utils import get_subset_type
 
 
 __all__ = ['show_widget', 'TemplateMixin', 'PluginTemplateMixin',
@@ -1014,8 +1015,8 @@ def _subset_to_dict(self, subset):
             for layer in viewer.layers:
                 if layer.layer.label == subset.label:
                     color = layer.state.color
-                    subset_type = _subset_type(subset)
-                    return {"label": subset.label, "color": color, "type": subset_type}
+                    type = get_subset_type(subset)
+                    return {"label": subset.label, "color": color, "type": type}
         return {"label": subset.label, "color": False, "type": False}
 
     def _delete_subset(self, subset):
@@ -1511,7 +1512,7 @@ def _cubeviz_include_spatial_subsets(self):
         # for changes in style, etc, we'll try to filter out extra messages in advance.
         def _subset_update(msg):
             if msg.attribute == 'subset_state':
-                if _subset_type(msg.subset) == 'spatial':
+                if get_subset_type(msg.subset) == 'spatial':
                     self._on_data_changed()
 
         self.hub.subscribe(self, SubsetUpdateMessage,
@@ -1624,7 +1625,7 @@ def _dc_to_dict(data):
         if getattr(self, '_include_spatial_subsets', False):
             # allow for spatial subsets to be listed
             self.items = self.items + [_dc_to_dict(subset) for subset in self.app.data_collection.subset_groups  # noqa
-                                       if _subset_type(subset) == 'spatial']
+                                       if get_subset_type(subset) == 'spatial']
         self._apply_default_selection()
         # future improvement: only clear cache if the selected data entry was changed?
         self._clear_cache(*self._cached_properties)

jdaviz/tests/test_subsets.py

@@ -5,12 +5,14 @@
 from glue.core import Data
 from glue.core.roi import CircularROI, CircularAnnulusROI, EllipticalROI, RectangularROI, XRangeROI
 from glue.core.edit_subset_mode import AndMode, AndNotMode, OrMode, XorMode
+from glue.core.subset_group import GroupedSubset
 from regions import (PixCoord, CirclePixelRegion, RectanglePixelRegion, EllipsePixelRegion,
                      CircleAnnulusPixelRegion)
 from numpy.testing import assert_allclose
 from specutils import SpectralRegion, Spectrum1D
 
 from jdaviz.core.marks import ShadowSpatialSpectral
+from jdaviz.utils import get_subset_type
 
 
 def test_region_from_subset_2d(cubeviz_helper):
@@ -199,7 +201,7 @@ def test_region_spectral_spatial(cubeviz_helper, spectral_cube_wcs):
     flux_viewer.toolbar.active_tool = flux_viewer.toolbar.tools['bqplot:rectangle']
     flux_viewer.apply_roi(RectangularROI(1, 3.5, -0.2, 3.3))
 
-    assert len([m for m in spectrum_viewer.figure.marks if isinstance(m, ShadowSpatialSpectral)]) == 1  # noqa
+    assert len([m for m in spectrum_viewer.figure.marks if isinstance(m, ShadowSpatialSpectral)]) == 2  # noqa
 
     subsets = cubeviz_helper.app.get_subsets(spectral_only=True)
     reg = subsets.get('Subset 1')
@@ -226,7 +228,7 @@ def test_region_spectral_spatial(cubeviz_helper, spectral_cube_wcs):
     spectrum_viewer.toolbar.active_tool = spectrum_viewer.toolbar.tools['jdaviz:panzoom']
     spectrum_viewer.toolbar.active_tool = spectrum_viewer.toolbar.tools['bqplot:xrange']
     spectrum_viewer.apply_roi(XRangeROI(3, 16))
-    assert len([m for m in spectrum_viewer.figure.marks if isinstance(m, ShadowSpatialSpectral)]) == 2  # noqa
+    assert len([m for m in spectrum_viewer.figure.marks if isinstance(m, ShadowSpatialSpectral)]) == 4  # noqa
 
 
 def test_disjoint_spatial_subset(cubeviz_helper, spectral_cube_wcs):
@@ -409,6 +411,10 @@ def test_composite_region_with_consecutive_and_not_states(cubeviz_helper):
     assert reg[1]['subset_state'].roi.xmin == 25
     assert reg[1]['subset_state'].roi.xmax == 30
 
+    for layer in viewer.state.layers:
+        if isinstance(layer.layer, GroupedSubset):
+            assert get_subset_type(layer.layer.subset_state) == 'spatial'
+
 
 def test_composite_region_with_imviz(imviz_helper, image_2d_wcs):
     arr = np.ones((10, 10))
@@ -503,6 +509,10 @@ def test_composite_region_from_subset_2d(specviz_helper, spectrum1d):
     subset_plugin.vue_simplify_subset()
     assert subset_plugin.glue_state_types == ["RangeSubsetState", "OrState"]
 
+    for layer in viewer.state.layers:
+        if isinstance(layer.layer, GroupedSubset):
+            assert get_subset_type(layer.layer.subset_state) == 'spectral'
+
 
 def test_edit_composite_spectral_subset(specviz_helper, spectrum1d):
     specviz_helper.load_data(spectrum1d)

jdaviz/utils.py

@@ -7,9 +7,11 @@
 from astropy.io import fits
 from ipyvue import watch
 from glue.config import settings
+from glue.core.subset import RangeSubsetState, RoiSubsetState
+
 
 __all__ = ['SnackbarQueue', 'enable_hot_reloading', 'bqplot_clear_figure',
-           'standardize_metadata', 'ColorCycler', 'alpha_index']
+           'standardize_metadata', 'ColorCycler', 'alpha_index', 'get_subset_type']
 
 # For Metadata Viewer plugin internal use only.
 PRIHDR_KEY = '_primary_header'
@@ -259,3 +261,33 @@ def __call__(self):
 
     def reset(self):
         self.counter = -1
+
+
+def get_subset_type(subset):
+    """
+    Determine the subset type of a subset or layer
+
+    Parameters
+    ----------
+    subset : glue.core.subset.Subset or glue.core.subset_group.GroupedSubset
+        should have ``subset_state`` as an attribute, otherwise will return ``None``.
+
+    Returns
+    -------
+    subset_type : str or None
+        'spatial', 'spectral', or None
+    """
+    if not hasattr(subset, 'subset_state'):
+        return None
+
+    while hasattr(subset.subset_state, 'state1'):
+        # this assumes no mixing between spatial and spectral subsets and just
+        # taking the first component (down the hierarchical tree) to determine the type
+        subset = subset.subset_state.state1
+
+    if isinstance(subset.subset_state, RoiSubsetState):
+        return 'spatial'
+    elif isinstance(subset.subset_state, RangeSubsetState):
+        return 'spectral'
+    else:
+        return None