From 2b304f4641d136285fd4d53c30d29b7f8992a443 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Mon, 15 May 2023 16:08:31 -0400
Subject: [PATCH 1/6]  generalize handling of composite subsets when
 determining subset type

---
 jdaviz/configs/cubeviz/plugins/viewers.py | 19 +++++++++----
 jdaviz/configs/default/plugins/viewers.py |  6 ++---
 jdaviz/core/template_mixin.py             | 15 ++++++-----
 jdaviz/utils.py                           | 33 ++++++++++++++++++++++-
 4 files changed, 57 insertions(+), 16 deletions(-)

diff --git a/jdaviz/configs/cubeviz/plugins/viewers.py b/jdaviz/configs/cubeviz/plugins/viewers.py
index 7dfc57353b..b92ec10187 100644
--- a/jdaviz/configs/cubeviz/plugins/viewers.py
+++ b/jdaviz/configs/cubeviz/plugins/viewers.py
@@ -1,5 +1,4 @@
 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
@@ -7,6 +6,7 @@
 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,7 +108,11 @@ 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):
@@ -116,7 +120,11 @@ def _get_spatial_subset_layers(self):
 
     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
diff --git a/jdaviz/configs/default/plugins/viewers.py b/jdaviz/configs/default/plugins/viewers.py
index 04148dcacd..353e5a12fa 100644
--- a/jdaviz/configs/default/plugins/viewers.py
+++ b/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", ""
diff --git a/jdaviz/core/template_mixin.py b/jdaviz/core/template_mixin.py
index ca06dbadad..7b1fc5a1d1 100644
--- a/jdaviz/core/template_mixin.py
+++ b/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):
@@ -1026,9 +1027,6 @@ def _delete_subset(self, subset):
             self._apply_default_selection()
 
     def _is_valid_item(self, subset):
-        def is_spectral(subset):
-            return _subset_type(subset) == 'spectral'
-
         def is_spatial(subset):
             return _subset_type(subset) == 'spatial'
 
@@ -1044,6 +1042,9 @@ def is_not_annulus(subset):
         return super()._is_valid_item(subset, locals())
 
     def _update_subset(self, subset, attribute=None):
+        if self._allowed_type is not None and get_subset_type(subset) != self._allowed_type:
+            return
+
         if subset.label not in self.labels:
             # NOTE: this logic will need to be revisited if generic renaming of subsets is added
             # see https://github.com/spacetelescope/jdaviz/pull/1175#discussion_r829372470
@@ -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)
diff --git a/jdaviz/utils.py b/jdaviz/utils.py
index 59ad0cf814..d5f92174e6 100644
--- a/jdaviz/utils.py
+++ b/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,32 @@ 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 : 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

From ca00f2e40127c463fa458d562b4e8337dfdcad6d Mon Sep 17 00:00:00 2001
From: Jesse Averbukh <javerbukh@gmail.com>
Date: Mon, 10 Jul 2023 15:26:20 -0400
Subject: [PATCH 2/6] Fix composite subset not applying coords info or shadow
 mark

---
 jdaviz/configs/imviz/plugins/coords_info/coords_info.py | 8 ++++++--
 jdaviz/configs/specviz/plugins/viewers.py               | 3 ++-
 jdaviz/tests/test_subsets.py                            | 4 ++--
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/jdaviz/configs/imviz/plugins/coords_info/coords_info.py b/jdaviz/configs/imviz/plugins/coords_info/coords_info.py
index 65c51ab88c..aa5bd6f77b 100644
--- a/jdaviz/configs/imviz/plugins/coords_info/coords_info.py
+++ b/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:
+                    return False
+                subset_type = get_subset_type(subset_state)
+                if subset_type == 'spectral':
                     # then this is a SPECTRAL subset
                     continue
                 # For use later in data retrieval
diff --git a/jdaviz/configs/specviz/plugins/viewers.py b/jdaviz/configs/specviz/plugins/viewers.py
index 2d5dc13b23..a988198bf1 100644
--- a/jdaviz/configs/specviz/plugins/viewers.py
+++ b/jdaviz/configs/specviz/plugins/viewers.py
@@ -17,6 +17,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 +378,7 @@ 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 get_subset_type(layer.layer) == 'spectral' and layer.layer.data.label == data.label:
                 layer.linewidth = 3
 
         return result
diff --git a/jdaviz/tests/test_subsets.py b/jdaviz/tests/test_subsets.py
index 2c9acbcd80..766cc3102e 100644
--- a/jdaviz/tests/test_subsets.py
+++ b/jdaviz/tests/test_subsets.py
@@ -199,7 +199,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 +226,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):

From 6af11b2b6d7057eddba8d9eaf6fe6e0557c3de87 Mon Sep 17 00:00:00 2001
From: Jesse Averbukh <javerbukh@gmail.com>
Date: Mon, 10 Jul 2023 15:30:42 -0400
Subject: [PATCH 3/6] Add change

---
 CHANGES.rst | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/CHANGES.rst b/CHANGES.rst
index d7f6d67152..119a4d22b4 100644
--- a/CHANGES.rst
+++ b/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
 ^^^^^
 

From c4d17a815ead9fe80a3147b0bba506a07cf18c4c Mon Sep 17 00:00:00 2001
From: Jesse Averbukh <javerbukh@gmail.com>
Date: Tue, 11 Jul 2023 09:46:13 -0400
Subject: [PATCH 4/6] Add code from #2287 that was accidentally removed during
 rebase

---
 jdaviz/core/template_mixin.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/jdaviz/core/template_mixin.py b/jdaviz/core/template_mixin.py
index 7b1fc5a1d1..4b27787724 100644
--- a/jdaviz/core/template_mixin.py
+++ b/jdaviz/core/template_mixin.py
@@ -1027,6 +1027,9 @@ def _delete_subset(self, subset):
             self._apply_default_selection()
 
     def _is_valid_item(self, subset):
+        def is_spectral(subset):
+            return _subset_type(subset) == 'spectral'
+
         def is_spatial(subset):
             return _subset_type(subset) == 'spatial'
 
@@ -1042,9 +1045,6 @@ def is_not_annulus(subset):
         return super()._is_valid_item(subset, locals())
 
     def _update_subset(self, subset, attribute=None):
-        if self._allowed_type is not None and get_subset_type(subset) != self._allowed_type:
-            return
-
         if subset.label not in self.labels:
             # NOTE: this logic will need to be revisited if generic renaming of subsets is added
             # see https://github.com/spacetelescope/jdaviz/pull/1175#discussion_r829372470

From 7a88777ce7994110d1c65fc2a2975ed7e2b46378 Mon Sep 17 00:00:00 2001
From: Jesse Averbukh <javerbukh@gmail.com>
Date: Wed, 12 Jul 2023 08:48:10 -0400
Subject: [PATCH 5/6] Address review comments and add tests

---
 .../configs/imviz/plugins/coords_info/coords_info.py   |  2 +-
 jdaviz/configs/specviz/plugins/viewers.py              |  5 ++++-
 jdaviz/tests/test_subsets.py                           | 10 ++++++++++
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/jdaviz/configs/imviz/plugins/coords_info/coords_info.py b/jdaviz/configs/imviz/plugins/coords_info/coords_info.py
index aa5bd6f77b..126d1856be 100644
--- a/jdaviz/configs/imviz/plugins/coords_info/coords_info.py
+++ b/jdaviz/configs/imviz/plugins/coords_info/coords_info.py
@@ -445,7 +445,7 @@ def _copy_axes_to_spectral():
             if isinstance(lyr.layer, GroupedSubset):
                 subset_state = getattr(lyr.layer, 'subset_state', None)
                 if subset_state is None:
-                    return False
+                    continue
                 subset_type = get_subset_type(subset_state)
                 if subset_type == 'spectral':
                     # then this is a SPECTRAL subset
diff --git a/jdaviz/configs/specviz/plugins/viewers.py b/jdaviz/configs/specviz/plugins/viewers.py
index a988198bf1..89640f1ca1 100644
--- a/jdaviz/configs/specviz/plugins/viewers.py
+++ b/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
@@ -378,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 get_subset_type(layer.layer) == 'spectral' 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
diff --git a/jdaviz/tests/test_subsets.py b/jdaviz/tests/test_subsets.py
index 766cc3102e..21c6c93f68 100644
--- a/jdaviz/tests/test_subsets.py
+++ b/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):
@@ -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)

From 10ba0932416469cd03b9bfa177a8175bb1f61e0d Mon Sep 17 00:00:00 2001
From: Jesse Averbukh <javerbukh@gmail.com>
Date: Wed, 12 Jul 2023 13:38:13 -0400
Subject: [PATCH 6/6] Update jdaviz/utils.py

Co-authored-by: Brett M. Morris <morrisbrettm@gmail.com>
---
 jdaviz/utils.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/jdaviz/utils.py b/jdaviz/utils.py
index d5f92174e6..7552e99bca 100644
--- a/jdaviz/utils.py
+++ b/jdaviz/utils.py
@@ -269,7 +269,8 @@ def get_subset_type(subset):
 
     Parameters
     ----------
-    subset : should have ``subset_state`` as an attribute, otherwise will return ``None``.
+    subset : glue.core.subset.Subset or glue.core.subset_group.GroupedSubset
+        should have ``subset_state`` as an attribute, otherwise will return ``None``.
 
     Returns
     -------