From 488c09b853a4a03041b7f203188793911b672d06 Mon Sep 17 00:00:00 2001
From: "Brett M. Morris" <bmmorris@stsci.edu>
Date: Tue, 19 Mar 2024 13:01:49 -0400
Subject: [PATCH 1/5] implement associations between Data layers

---
 jdaviz/app.py                           | 63 +++++++++++++++++++++++--
 jdaviz/configs/imviz/plugins/parsers.py | 16 +++----
 jdaviz/core/template_mixin.py           | 12 +++++
 jdaviz/tests/test_app.py                | 15 ++++++
 4 files changed, 95 insertions(+), 11 deletions(-)

diff --git a/jdaviz/app.py b/jdaviz/app.py
index 2a462fc546..295c329c79 100644
--- a/jdaviz/app.py
+++ b/jdaviz/app.py
@@ -356,6 +356,9 @@ def __init__(self, configuration=None, *args, **kwargs):
         self.hub.subscribe(self, SubsetUpdateMessage,
                            handler=lambda msg: self._clear_object_cache(msg.subset.label))
 
+        # Store for associations between Data entries:
+        self._data_associations = self._init_data_associations()
+
         # Subscribe to messages that result in changes to the layers
         self.hub.subscribe(self, AddDataMessage,
                            handler=self._on_layers_changed)
@@ -473,9 +476,11 @@ def _on_layers_changed(self, msg):
         if hasattr(msg, 'data'):
             layer_name = msg.data.label
             is_wcs_only = msg.data.meta.get(_wcs_only_label, False)
+            is_not_child = self._get_assoc_data_parent(layer_name) is None
         elif hasattr(msg, 'subset'):
             layer_name = msg.subset.label
             is_wcs_only = False
+            is_not_child = True
         else:
             raise NotImplementedError(f"cannot recognize new layer from {msg}")
 
@@ -490,13 +495,25 @@ def _on_layers_changed(self, msg):
                 self.state.layer_icons = {**self.state.layer_icons,
                                           layer_name: orientation_icons.get(layer_name,
                                                                             wcs_only_refdata_icon)}
-            else:
+            elif is_not_child:
                 self.state.layer_icons = {
                     **self.state.layer_icons,
                     layer_name: alpha_index(len([ln for ln, ic in self.state.layer_icons.items()
-                                                 if not ic.startswith('mdi-')]))
+                                                 if not ic.startswith('mdi-') and
+                                                 self._get_assoc_data_parent(ln) is None]))
                 }
 
+        # all remaining layers at this point have a parent:
+        for layer_name in self.state.layer_icons:
+            children_layers = self._get_assoc_data_children(layer_name)
+            if children_layers is not None:
+                parent_icon = self.state.layer_icons[layer_name]
+                for i, child_layer in enumerate(children_layers, start=1):
+                    self.state.layer_icons = {
+                        **self.state.layer_icons,
+                        child_layer: f'{parent_icon}{i}'
+                    }
+
     def _change_reference_data(self, new_refdata_label, viewer_id=None):
         """
         Change reference data to Data with ``data_label``.
@@ -1251,7 +1268,7 @@ def merge_overlapping_spectral_regions(self, subset_name, att):
 
         return new_state
 
-    def add_data(self, data, data_label=None, notify_done=True):
+    def add_data(self, data, data_label=None, notify_done=True, parent=None):
         """
         Add data to the Glue ``DataCollection``.
 
@@ -1280,6 +1297,11 @@ def add_data(self, data, data_label=None, notify_done=True):
 
         self.data_collection[data_label] = data
 
+        # manage associated Data entries:
+        self._add_assoc_data_as_parent(data_label)
+        if parent is not None:
+            self._set_assoc_data_as_child(data_label, new_parent_label=parent)
+
         # Send out a toast message
         if notify_done:
             snackbar_message = SnackbarMessage(
@@ -1998,6 +2020,17 @@ def set_data_visibility(self, viewer_reference, data_label, visible=True, replac
                 if layer.layer.data.label != data_label:
                     layer.visible = False
 
+        # if Data has children, update their visibilities to match Data:
+        assoc_children = self._get_assoc_data_children(data_label)
+        for layer in viewer.layers:
+            for data_label in assoc_children:
+                if layer.layer.data.label == data_label:
+                    if visible and not layer.visible:
+                        layer.visible = True
+                        layer.update()
+                    else:
+                        layer.visible = visible
+
         # update data menu - selected_data_items should be READ ONLY, not modified by the user/UI
         selected_items = viewer_item['selected_data_items']
         data_id = self._data_id_from_label(data_label)
@@ -2577,3 +2610,27 @@ def get_tray_item_from_name(self, name):
             raise KeyError(f'{name} not found in app.state.tray_items')
 
         return tray_item
+
+    def _init_data_associations(self):
+        # assume all Data are parents:
+        data_associations = {
+            data.label: {'parent': None, 'children': []}
+            for data in self.data_collection
+        }
+        return data_associations
+
+    def _add_assoc_data_as_parent(self, data_label):
+        self._data_associations[data_label] = {'parent': None, 'children': []}
+
+    def _set_assoc_data_as_child(self, data_label, new_parent_label):
+        # Data has a new parent:
+        self._data_associations[data_label]['parent'] = new_parent_label
+        # parent has a new child:
+        self._data_associations[new_parent_label]['children'].append(data_label)
+
+    def _get_assoc_data_children(self, data_label):
+        # intentionally not recursive for now, just one generation:
+        return self._data_associations.get(data_label, {}).get('children', [])
+
+    def _get_assoc_data_parent(self, data_label):
+        return self._data_associations.get(data_label, {}).get('parent')
diff --git a/jdaviz/configs/imviz/plugins/parsers.py b/jdaviz/configs/imviz/plugins/parsers.py
index 7228ae5079..97558ed568 100644
--- a/jdaviz/configs/imviz/plugins/parsers.py
+++ b/jdaviz/configs/imviz/plugins/parsers.py
@@ -30,7 +30,7 @@
 
 
 @data_parser_registry("imviz-data-parser")
-def parse_data(app, file_obj, ext=None, data_label=None):
+def parse_data(app, file_obj, ext=None, data_label=None, parent=None):
     """Parse a data file into Imviz.
 
     Parameters
@@ -74,17 +74,17 @@ def parse_data(app, file_obj, ext=None, data_label=None):
             else:  # Assume RGB
                 pf = rgb2gray(im)
             pf = pf[::-1, :]  # Flip it
-            _parse_image(app, pf, data_label, ext=ext)
+            _parse_image(app, pf, data_label, ext=ext, parent=parent)
 
         elif file_obj_lower.endswith('.asdf'):
             try:
                 if HAS_ROMAN_DATAMODELS:
                     with rdd.open(file_obj) as pf:
-                        _parse_image(app, pf, data_label, ext=ext)
+                        _parse_image(app, pf, data_label, ext=ext, parent=parent)
             except TypeError:
                 # if roman_datamodels cannot parse the file, load it with asdf:
                 with asdf.open(file_obj) as af:
-                    _parse_image(app, af, data_label, ext=ext)
+                    _parse_image(app, af, data_label, ext=ext, parent=parent)
 
         elif file_obj_lower.endswith('.reg'):
             # This will load DS9 regions as Subset but only if there is already data.
@@ -92,9 +92,9 @@ def parse_data(app, file_obj, ext=None, data_label=None):
 
         else:  # Assume FITS
             with fits.open(file_obj) as pf:
-                _parse_image(app, pf, data_label, ext=ext)
+                _parse_image(app, pf, data_label, ext=ext, parent=parent)
     else:
-        _parse_image(app, file_obj, data_label, ext=ext)
+        _parse_image(app, file_obj, data_label, ext=ext, parent=parent)
 
 
 def get_image_data_iterator(app, file_obj, data_label, ext=None):
@@ -168,7 +168,7 @@ def get_image_data_iterator(app, file_obj, data_label, ext=None):
     return data_iter
 
 
-def _parse_image(app, file_obj, data_label, ext=None):
+def _parse_image(app, file_obj, data_label, ext=None, parent=None):
     if app is None:
         raise ValueError("app is None, cannot proceed")
     if data_label is None:
@@ -186,7 +186,7 @@ def _parse_image(app, file_obj, data_label, ext=None):
             data.coords.bounding_box = None
         if not data.meta.get(_wcs_only_label, False):
             data_label = app.return_data_label(data_label, alt_name="image_data")
-        app.add_data(data, data_label)
+        app.add_data(data, data_label, parent=parent)
 
     # Do not link image data here. We do it at the end in Imviz.load_data()
 
diff --git a/jdaviz/core/template_mixin.py b/jdaviz/core/template_mixin.py
index 1a4626b42c..f5de99412a 100644
--- a/jdaviz/core/template_mixin.py
+++ b/jdaviz/core/template_mixin.py
@@ -1365,6 +1365,12 @@ def __init__(self, plugin, items, selected, viewer,
         self._update_layer_items()
         self.update_wcs_only_filter(only_wcs_layers)
 
+        # ignore layers that are children in associations:
+        def is_parent(data):
+            return self.app._get_assoc_data_parent(data.label) is None
+
+        self.add_filter(is_parent)
+
     def _get_viewer(self, viewer):
         # newer will likely be the viewer name in most cases, but viewer id in the case
         # of additional viewers in imviz.
@@ -2923,6 +2929,12 @@ def __init__(self, plugin, items, selected,
         # initialize items from original viewers
         self._on_data_changed()
 
+        # ignore layers that are children in associations:
+        def is_parent(data):
+            return self.app._get_assoc_data_parent(data.label) is None
+
+        self.add_filter(is_parent)
+
     def _cubeviz_include_spatial_subsets(self):
         """
         Call this method to prepend spatial subsets to the list of datasets (and listen for newly
diff --git a/jdaviz/tests/test_app.py b/jdaviz/tests/test_app.py
index adbcd29da3..2658568d35 100644
--- a/jdaviz/tests/test_app.py
+++ b/jdaviz/tests/test_app.py
@@ -1,5 +1,6 @@
 import pytest
 
+import numpy as np
 from jdaviz import Application, Specviz
 from jdaviz.configs.default.plugins.gaussian_smooth.gaussian_smooth import GaussianSmooth
 
@@ -170,3 +171,17 @@ def test_viewer_renaming_imviz(imviz_helper):
             old_reference='non-existent',
             new_reference='this-is-forbidden'
         )
+
+
+def test_data_associations(imviz_helper):
+    shape = (10, 10)
+
+    data_parent = np.ones(shape, dtype=float)
+    data_child = np.zeros(shape, dtype=int)
+
+    imviz_helper.load_data(data_parent, data_label='parent_data')
+    imviz_helper.load_data(data_child, data_label='child_data', parent='parent_data')
+
+    assert imviz_helper.app._get_assoc_data_children('parent_data') == ['child_data']
+    assert imviz_helper.app._get_assoc_data_parent('child_data') == 'parent_data'
+

From fe40115396c23658a8a931cf3d484eeed1c60674 Mon Sep 17 00:00:00 2001
From: "Brett M. Morris" <bmmorris@stsci.edu>
Date: Tue, 19 Mar 2024 15:10:15 -0400
Subject: [PATCH 2/5] support for Imviz blinking, glue 'bad' masks

---
 jdaviz/configs/imviz/plugins/parsers.py | 10 ++++++++++
 jdaviz/configs/imviz/plugins/viewers.py | 18 +++++++++++++++---
 jdaviz/tests/test_app.py                |  1 -
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/jdaviz/configs/imviz/plugins/parsers.py b/jdaviz/configs/imviz/plugins/parsers.py
index 97558ed568..06cbb44263 100644
--- a/jdaviz/configs/imviz/plugins/parsers.py
+++ b/jdaviz/configs/imviz/plugins/parsers.py
@@ -186,6 +186,16 @@ def _parse_image(app, file_obj, data_label, ext=None, parent=None):
             data.coords.bounding_box = None
         if not data.meta.get(_wcs_only_label, False):
             data_label = app.return_data_label(data_label, alt_name="image_data")
+
+        # TODO: generalize/centralize this for use in other configs too
+        if parent is not None and ext == 'DQ':
+            print('data', data.main_components)
+            # nans are used to mark "good" flags in the DQ colormap, so
+            # convert DQ array to float to support nans:
+            cid = data.get_component("DQ")
+            data_arr = np.float32(cid.data)
+            data_arr[data_arr == 0] = np.nan
+            data.update_components({cid: data_arr})
         app.add_data(data, data_label, parent=parent)
 
     # Do not link image data here. We do it at the end in Imviz.load_data()
diff --git a/jdaviz/configs/imviz/plugins/viewers.py b/jdaviz/configs/imviz/plugins/viewers.py
index 285d801b77..788ae57fef 100644
--- a/jdaviz/configs/imviz/plugins/viewers.py
+++ b/jdaviz/configs/imviz/plugins/viewers.py
@@ -86,9 +86,16 @@ def blink_once(self, reversed=False):
         # Simple blinking of images - this will make it so that only one
         # layer is visible at a time and cycles through the layers.
 
-        # Exclude Subsets: They are global
+        # Exclude Subsets (they are global) and children via associated data
+
+        def is_parent(data):
+            return self.session.jdaviz_app._get_assoc_data_parent(data.label) is None
+
         valid = [ilayer for ilayer, layer in enumerate(self.state.layers)
-                 if layer_is_image_data(layer.layer)]
+                 if layer_is_image_data(layer.layer) and is_parent(layer.layer)]
+        children = [ilayer for ilayer, layer in enumerate(self.state.layers)
+                    if layer_is_image_data(layer.layer) and not is_parent(layer.layer)]
+
         n_layers = len(valid)
 
         if n_layers == 1:
@@ -116,7 +123,12 @@ def blink_once(self, reversed=False):
                 next_layer = valid[(valid.index(visible[-1]) + delta) % n_layers]
                 self.state.layers[next_layer].visible = True
 
-                for ilayer in (set(valid) - set([next_layer])):
+                # make invisible all parent layers other than the next layer:
+                layers_to_set_not_visible = set(valid) - set([next_layer])
+                # no child layers are visible by default:
+                layers_to_set_not_visible.update(set(children))
+
+                for ilayer in layers_to_set_not_visible:
                     self.state.layers[ilayer].visible = False
 
                 # We can display the active data label in Compass plugin.
diff --git a/jdaviz/tests/test_app.py b/jdaviz/tests/test_app.py
index 2658568d35..9179af4350 100644
--- a/jdaviz/tests/test_app.py
+++ b/jdaviz/tests/test_app.py
@@ -184,4 +184,3 @@ def test_data_associations(imviz_helper):
 
     assert imviz_helper.app._get_assoc_data_children('parent_data') == ['child_data']
     assert imviz_helper.app._get_assoc_data_parent('child_data') == 'parent_data'
-

From 5afa49fd552c5e5d763b78fd5b7431dc68abd1b7 Mon Sep 17 00:00:00 2001
From: "Brett M. Morris" <bmmorris@stsci.edu>
Date: Wed, 20 Mar 2024 12:13:36 -0400
Subject: [PATCH 3/5] prevent children of children in data associations

---
 jdaviz/app.py                           | 11 ++++++++++-
 jdaviz/configs/imviz/plugins/parsers.py |  1 -
 jdaviz/tests/test_app.py                |  4 ++++
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/jdaviz/app.py b/jdaviz/app.py
index 295c329c79..68a308c949 100644
--- a/jdaviz/app.py
+++ b/jdaviz/app.py
@@ -1283,10 +1283,12 @@ def add_data(self, data, data_label=None, notify_done=True, parent=None):
             The name associated with this data. If none is given, label is pulled
             from the input data (if `~glue.core.data.Data`) or a generic name is
             generated.
-        notify_done: bool
+        notify_done : bool
             Flag controlling whether a snackbar message is set when the data is
             added to the app. Set to False to avoid overwhelming the user if
             lots of data is getting loaded at once.
+        parent : str, optional
+            Associate the added Data entry as the child of layer ``parent``.
         """
 
         if not data_label and hasattr(data, "label"):
@@ -1300,6 +1302,13 @@ def add_data(self, data, data_label=None, notify_done=True, parent=None):
         # manage associated Data entries:
         self._add_assoc_data_as_parent(data_label)
         if parent is not None:
+            # Does the parent Data have a parent? If so, raise error:
+            parent_of_parent = self._get_assoc_data_parent(parent)
+            if parent_of_parent is not None:
+                raise NotImplementedError('Data associations are currently supported '
+                                          'between root layers (without parents) and their '
+                                          f'children, but the proposed parent "{parent}" has '
+                                          f'parent "{parent_of_parent}".')
             self._set_assoc_data_as_child(data_label, new_parent_label=parent)
 
         # Send out a toast message
diff --git a/jdaviz/configs/imviz/plugins/parsers.py b/jdaviz/configs/imviz/plugins/parsers.py
index 06cbb44263..b6ea47938d 100644
--- a/jdaviz/configs/imviz/plugins/parsers.py
+++ b/jdaviz/configs/imviz/plugins/parsers.py
@@ -189,7 +189,6 @@ def _parse_image(app, file_obj, data_label, ext=None, parent=None):
 
         # TODO: generalize/centralize this for use in other configs too
         if parent is not None and ext == 'DQ':
-            print('data', data.main_components)
             # nans are used to mark "good" flags in the DQ colormap, so
             # convert DQ array to float to support nans:
             cid = data.get_component("DQ")
diff --git a/jdaviz/tests/test_app.py b/jdaviz/tests/test_app.py
index 9179af4350..9311bd41d3 100644
--- a/jdaviz/tests/test_app.py
+++ b/jdaviz/tests/test_app.py
@@ -184,3 +184,7 @@ def test_data_associations(imviz_helper):
 
     assert imviz_helper.app._get_assoc_data_children('parent_data') == ['child_data']
     assert imviz_helper.app._get_assoc_data_parent('child_data') == 'parent_data'
+
+    with pytest.raises(NotImplementedError):
+        # we don't (yet) allow children of children:
+        imviz_helper.load_data(data_child, data_label='grandchild_data', parent='child_data')

From 55d2e688c44d5ef4a86d624a978619c300f9ff87 Mon Sep 17 00:00:00 2001
From: "Brett M. Morris" <bmmorris@stsci.edu>
Date: Wed, 20 Mar 2024 16:07:14 -0400
Subject: [PATCH 4/5] ensure requested parent in data associations exists

---
 jdaviz/app.py            | 5 +++++
 jdaviz/tests/test_app.py | 4 ++++
 2 files changed, 9 insertions(+)

diff --git a/jdaviz/app.py b/jdaviz/app.py
index 68a308c949..e0dfeef67a 100644
--- a/jdaviz/app.py
+++ b/jdaviz/app.py
@@ -1302,6 +1302,11 @@ def add_data(self, data, data_label=None, notify_done=True, parent=None):
         # manage associated Data entries:
         self._add_assoc_data_as_parent(data_label)
         if parent is not None:
+            data_collection_labels = [data.label for data in self.data_collection]
+            if parent not in data_collection_labels:
+                raise ValueError(f'parent "{parent}" is not a valid data label in '
+                                 f'the data collection: {data_collection_labels}.')
+
             # Does the parent Data have a parent? If so, raise error:
             parent_of_parent = self._get_assoc_data_parent(parent)
             if parent_of_parent is not None:
diff --git a/jdaviz/tests/test_app.py b/jdaviz/tests/test_app.py
index 9311bd41d3..a81ad5ac42 100644
--- a/jdaviz/tests/test_app.py
+++ b/jdaviz/tests/test_app.py
@@ -188,3 +188,7 @@ def test_data_associations(imviz_helper):
     with pytest.raises(NotImplementedError):
         # we don't (yet) allow children of children:
         imviz_helper.load_data(data_child, data_label='grandchild_data', parent='child_data')
+
+    with pytest.raises(ValueError):
+        # ensure the parent actually exists:
+        imviz_helper.load_data(data_child, data_label='child_data', parent='absent parent')

From 00785d7f7386f68777e359a0c22748b4ce7fe4a3 Mon Sep 17 00:00:00 2001
From: "Brett M. Morris" <bmmorris@stsci.edu>
Date: Wed, 20 Mar 2024 16:21:28 -0400
Subject: [PATCH 5/5] more efficient layer icon updates

---
 jdaviz/app.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/jdaviz/app.py b/jdaviz/app.py
index e0dfeef67a..1da83663be 100644
--- a/jdaviz/app.py
+++ b/jdaviz/app.py
@@ -477,10 +477,13 @@ def _on_layers_changed(self, msg):
             layer_name = msg.data.label
             is_wcs_only = msg.data.meta.get(_wcs_only_label, False)
             is_not_child = self._get_assoc_data_parent(layer_name) is None
+            children_layers = self._get_assoc_data_children(layer_name)
+
         elif hasattr(msg, 'subset'):
             layer_name = msg.subset.label
             is_wcs_only = False
             is_not_child = True
+            children_layers = []
         else:
             raise NotImplementedError(f"cannot recognize new layer from {msg}")
 
@@ -504,15 +507,18 @@ def _on_layers_changed(self, msg):
                 }
 
         # all remaining layers at this point have a parent:
+        child_layer_icons = {}
         for layer_name in self.state.layer_icons:
             children_layers = self._get_assoc_data_children(layer_name)
             if children_layers is not None:
                 parent_icon = self.state.layer_icons[layer_name]
                 for i, child_layer in enumerate(children_layers, start=1):
-                    self.state.layer_icons = {
-                        **self.state.layer_icons,
-                        child_layer: f'{parent_icon}{i}'
-                    }
+                    child_layer_icons[child_layer] = f'{parent_icon}{i}'
+
+        self.state.layer_icons = {
+            **self.state.layer_icons,
+            **child_layer_icons
+        }
 
     def _change_reference_data(self, new_refdata_label, viewer_id=None):
         """