Debug orientation operations with multiple viewers (#2759)

* Experimenting with fixes

* Fix wrong viewer's orientation being updated in plugin from viewer data menu change

* Fix orientation not being selected in viewer after deleting previous selection

* Add PR number to changelog

* Use base_wcs_layer_label

* Add test for orientation change

* Better fix for ref data deletion bug

* Only do this if wcs linked, just in case

* Add test for last bug

* Revert unneeded change

* Apply suggestions from code review

Co-authored-by: P. L. Lim <2090236+pllim@users.noreply.github.com>

---------

Co-authored-by: P. L. Lim <2090236+pllim@users.noreply.github.com>

CHANGES.rst

@@ -43,7 +43,7 @@ Imviz
 ^^^^^
 
 - There is now option for image rotation in Orientation (was Links Control) plugin.
-  This feature requires WCS linking. [#2179, #2673, #2699, #2734]
+  This feature requires WCS linking. [#2179, #2673, #2699, #2734, #2759]
 
 - Add "Random" colormap for visualizing image segmentation maps. [#2671]
 

jdaviz/app.py

@@ -2043,13 +2043,19 @@ def vue_data_item_remove(self, event):
         data = self.data_collection[data_label]
         orientation_plugin = self._jdaviz_helper.plugins.get("Orientation")
         if orientation_plugin is not None:
+            from jdaviz.configs.imviz.helper import base_wcs_layer_label
             orient = orientation_plugin.orientation.selected
+            if orient == data_label:
+                orient = base_wcs_layer_label
             self._reparent_subsets(data, new_parent=orient)
         else:
             self._reparent_subsets(data)
 
-        # Make sure the data isn't loaded in any viewers
-        for viewer_id in self._viewer_store:
+        # Make sure the data isn't loaded in any viewers and isn't the selected orientation
+        for viewer_id, viewer in self._viewer_store.items():
+            if orientation_plugin is not None and self._link_type == 'wcs':
+                if viewer.state.reference_data.label == data_label:
+                    self._change_reference_data(base_wcs_layer_label, viewer_id)
             self.remove_data_from_viewer(viewer_id, data_label)
 
         self.data_collection.remove(self.data_collection[data_label])

jdaviz/configs/imviz/plugins/orientation/orientation.py

@@ -359,6 +359,9 @@ def add_orientation(self, rotation_angle=None, east_left=None, label=None,
             self._add_data_to_viewer(label, viewer_ref)
 
         if set_on_create:
+            # Not sure why this is sometimes missing, but we can add it here
+            if label not in self.orientation.choices:
+                self.orientation.choices += [label]
             self.orientation.selected = label
 
     def _add_data_to_viewer(self, data_label, viewer_id):
@@ -414,7 +417,7 @@ def _change_reference_data(self, *args, **kwargs):
                 self.orientation.selected, viewer_id=self.viewer.selected
             )
             viewer_item = self.app._viewer_item_by_id(self.viewer.selected)
-            if viewer_item != self.orientation.selected:
+            if viewer_item['reference_data_label'] != self.orientation.selected:
                 viewer_item['reference_data_label'] = self.orientation.selected
 
     def _on_refdata_change(self, msg):
@@ -434,7 +437,8 @@ def _on_refdata_change(self, msg):
             if msg.data.label not in self.orientation.choices:
                 return
 
-            self.orientation.selected = msg.data.label
+            if msg.viewer_id == self.viewer.selected:
+                self.orientation.selected = msg.data.label
 
             # we never want to highlight subsets of pixels within WCS-only layers,
             # so if this layer is an ImageSubsetLayerState on a WCS-only layer,

jdaviz/configs/imviz/tests/test_orientation.py

@@ -131,11 +131,17 @@ def test_N_up_multi_viewer(self):
         viewer_2 = self.imviz.create_image_viewer()
         self.imviz.app.add_data_to_viewer("imviz-1", "has_wcs_1[SCI,1]")
         lc_plugin.viewer = "imviz-1"
+
         lc_plugin._obj.create_north_up_east_right(set_on_create=True)
 
         assert self.viewer.state.reference_data.label == "North-up, East-left"
         assert viewer_2.state.reference_data.label == "North-up, East-right"
 
+        # Change orientation in imviz-1 from UI and ensure plugin selection is the same
+        lc_plugin.viewer.selected = "imviz-0"
+        self.imviz.app._change_reference_data("Default orientation", "imviz-1")
+        assert lc_plugin.orientation.selected == "North-up, East-left"
+
         # Both viewers should revert back to same reference when pixel-linked.
         lc_plugin.link_type = 'Pixels'
         assert self.viewer.state.reference_data.label == "has_wcs_1[SCI,1]"
@@ -152,6 +158,7 @@ def test_custom_orientation(self):
         lc_plugin = self.imviz.plugins['Orientation']
         lc_plugin.link_type = 'WCS'
         lc_plugin.viewer = "imviz-0"
+
         lc_plugin.rotation_angle = 42  # Triggers auto-label
         lc_plugin._obj.add_orientation(rotation_angle=None, east_left=True, label=None,
                                        set_on_create=True, wrt_data=None)
@@ -160,6 +167,24 @@ def test_custom_orientation(self):
 
 class TestDeleteOrientation(BaseImviz_WCS_WCS):
 
+    def test_delete_orientation_multi_viewer(self):
+        lc_plugin = self.imviz.plugins['Orientation']
+        lc_plugin.link_type = 'WCS'
+
+        # Should automatically be applied as reference to first viewer.
+        lc_plugin._obj.create_north_up_east_left(set_on_create=True)
+
+        # This would set a different reference to second viewer.
+        viewer_2 = self.imviz.create_image_viewer()
+        self.imviz.app.add_data_to_viewer("imviz-1", "has_wcs_1[SCI,1]")
+        lc_plugin.viewer = "imviz-1"
+        lc_plugin.orientation = "North-up, East-left"
+
+        self.imviz.app.vue_data_item_remove({"item_name": "North-up, East-left"})
+
+        assert self.viewer.state.reference_data.label == "Default orientation"
+        assert viewer_2.state.reference_data.label == "Default orientation"
+
     @pytest.mark.parametrize("klass", [EllipseSkyRegion, RectangleSkyRegion])
     @pytest.mark.parametrize(
         ("angle", "sbst_theta"),