Repository: spacetelescope/jdaviz

Issue:
BUG: WCS to Pixels to WCS gives wrong reference data
This is a follow-up of:

* #2179 

And discovered in:

* #2712

Workflow to reproduce:

1. Use the `ImvizDitheredExample.ipynb` example notebook. Load both data.
2. Link by WCS on the Orientation plugin.
3. Click the "N-up E-left" icon so now the images in first viewer is aligned by N-up E-left.
4. Create a second viewer and load the first image only.
5. With second viewer (`imviz-1`) selected in Orientation, click on the "N-up E-right" icon so now the second viewer image is aligned by N-up E-right.
6. Now select "Pixels" in Orientation. Both viewers will go back to pixel linking and no longer appear rotated.
7. Now click "WCS" again in Orientation. See that both images do not go back to being aligned by one of the 3 possible WCS layers (see screenshot below). Despite that, when you blink, the stars still aligned (despite the exposures being dithered) as if #2179 never happened.

![Screenshot 2024-02-22 162201](https://github.com/spacetelescope/jdaviz/assets/2090236/8e4590d0-14c5-46e4-8da9-bf18e54fec9e)

[🐱](https://jira.stsci.edu/browse/JDAT-4224)

Discussion:
pllim: Is this going in circles?

https://github.com/spacetelescope/jdaviz/blob/deda388472169d11604db304a369869e2e13d00e/jdaviz/configs/imviz/plugins/orientation/orientation.py#L410-L413

triggers

https://github.com/spacetelescope/jdaviz/blob/deda388472169d11604db304a369869e2e13d00e/jdaviz/app.py#L546

that is being listened back in

https://github.com/spacetelescope/jdaviz/blob/deda388472169d11604db304a369869e2e13d00e/jdaviz/configs/imviz/plugins/orientation/orientation.py#L122-L123

that calls this

https://github.com/spacetelescope/jdaviz/blob/deda388472169d11604db304a369869e2e13d00e/jdaviz/configs/imviz/plugins/orientation/orientation.py#L437

that would trigger this again?

https://github.com/spacetelescope/jdaviz/blob/deda388472169d11604db304a369869e2e13d00e/jdaviz/configs/imviz/plugins/orientation/orientation.py#L410-L413
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/spacetelescope/jdaviz/pull/2751?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope) Report
Attention: Patch coverage is `81.08108%` with `21 lines` in your changes are missing coverage. Please review.
> Project coverage is 88.65%. Comparing base [(`9515b7a`)](https://app.codecov.io/gh/spacetelescope/jdaviz/commit/9515b7ac4c6cd7c42939ae35a52aa05f4118a6c5?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope) to head [(`22d6eb4`)](https://app.codecov.io/gh/spacetelescope/jdaviz/pull/2751?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope).
> Report is 3 commits behind head on main.

| [Files](https://app.codecov.io/gh/spacetelescope/jdaviz/pull/2751?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope) | Patch % | Lines |
|---|---|---|
| [...z/configs/imviz/plugins/orientation/orientation.py](https://app.codecov.io/gh/spacetelescope/jdaviz/pull/2751?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope#diff-amRhdml6L2NvbmZpZ3MvaW12aXovcGx1Z2lucy9vcmllbnRhdGlvbi9vcmllbnRhdGlvbi5weQ==) | 80.73% | [21 Missing :warning: ](https://app.codecov.io/gh/spacetelescope/jdaviz/pull/2751?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #2751      +/-   ##
==========================================
- Coverage   88.68%   88.65%   -0.04%     
==========================================
  Files         108      108              
  Lines       16214    16203      -11     
==========================================
- Hits        14379    14364      -15     
- Misses       1835     1839       +4     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/spacetelescope/jdaviz/pull/2751?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=spacetelescope).

rosteen: > Is this going in circles?

I think it could potentially get called one extra time, but the second time `self.orientation.selected` wouldn't actually change so it wouldn't trigger another loop.
pllim: Theoretically Orientation is the only plugin that is allowed to change the ref data, right? There is no reason for it to listen to its own message?
rosteen: You can also change the ref data through the viewer data menu, so the Orientation plugin needs to see those events to stay in synch.
pllim: Right... 😒  Thanks for the reminder!
rosteen: Drats, I thought there was a good chance that this also fixed the bug I was trying to fix on Friday, but it didn't. I did confirm that it fixes the intended use case though. I'll take a look at the actual code before approving.
pllim: > the actual code

Sorry for the messy diff.
pllim: Code suggestion accepted. Thanks!
pllim: OK looks like the lines are some safeguards that didn't get triggered by test suite. I can look into that as part of this PR but I should probably rebase first. 🤞 
pllim: Change log failure is unrelated but something I need to fix anyway: https://github.com/scientific-python/action-check-changelogfile/issues/13

Patch:
diff --git a/jdaviz/app.py b/jdaviz/app.py
index ffae9f42..e9276b9e 100644
--- a/jdaviz/app.py
+++ b/jdaviz/app.py
@@ -2096,7 +2096,7 @@ class Application(VuetifyTemplate, HubListener):
         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
+            from jdaviz.configs.imviz.plugins.orientation.orientation import base_wcs_layer_label
             orient = orientation_plugin.orientation.selected
             if orient == data_label:
                 orient = base_wcs_layer_label
diff --git a/jdaviz/configs/imviz/helper.py b/jdaviz/configs/imviz/helper.py
index 0d83f754..63758185 100644
--- a/jdaviz/configs/imviz/helper.py
+++ b/jdaviz/configs/imviz/helper.py
@@ -4,23 +4,15 @@ import warnings
 from copy import deepcopy
 
 import numpy as np
-import astropy.units as u
-from astropy.wcs.wcsapi import BaseHighLevelWCS
 from glue.core import BaseData
 from glue.core.link_helpers import LinkSame
-from glue.plugins.wcs_autolinking.wcs_autolinking import WCSLink, NoAffineApproximation
 
-from jdaviz.core.events import SnackbarMessage, NewViewerMessage, LinkUpdatedMessage
+from jdaviz.core.events import SnackbarMessage, NewViewerMessage
 from jdaviz.core.helpers import ImageConfigHelper
-from jdaviz.configs.imviz.wcs_utils import (
-    _get_rotated_nddata_from_label, get_compass_info
-)
 from jdaviz.utils import data_has_valid_wcs, _wcs_only_label
 
 __all__ = ['Imviz']
 
-base_wcs_layer_label = 'Default orientation'
-
 
 class Imviz(ImageConfigHelper):
     """Imviz Helper class."""
@@ -436,7 +428,8 @@ def get_top_layer_index(viewer):
 
 def get_reference_image_data(app, viewer_id=None):
     """
-    Return the reference data in the first image viewer and its index
+    Return the current reference data in the given image viewer and its index.
+    By default, the first viewer is used.
     """
     if viewer_id is None:
         refdata = app._jdaviz_helper.default_viewer._obj.state.reference_data
@@ -448,240 +441,4 @@ def get_reference_image_data(app, viewer_id=None):
         iref = app.data_collection.index(refdata)
         return refdata, iref
 
-    # if reference data not found above, fall back on old method:
-    for i, data in enumerate(app.data_collection):
-        if layer_is_image_data(data):
-            iref = i
-            refdata = data
-            break
-    if refdata is None:
-        raise ValueError(f'No valid reference data found in collection: {app.data_collection}')
-    return refdata, iref
-
-
-# TODO: This is not really public API, so we can move what Orientation uses here into the plugin
-#       and remove this function from helper.py module in the future. Also move base_wcs_layer_label
-#       and remove update_plugin keyword when that happens.
-def link_image_data(app, link_type='pixels', wcs_fallback_scheme=None, wcs_use_affine=True,
-                    error_on_fail=False, update_plugin=True):
-    """(Re)link loaded data in Imviz with the desired link type.
-
-    .. note::
-
-        Any markers added in Imviz will need to be removed manually before changing linking type.
-        You can add back the markers using
-        :meth:`~jdaviz.core.astrowidgets_api.AstrowidgetsImageViewerMixin.add_markers`
-        for the relevant viewer(s).
-
-    Parameters
-    ----------
-    app : `~jdaviz.app.Application`
-        Application associated with Imviz, e.g., ``imviz.app``.
-
-    link_type : {'pixels', 'wcs'}
-        Choose to link by pixels or WCS.
-
-    wcs_fallback_scheme : {None, 'pixels'}
-        If WCS linking failed, choose to fall back to linking by pixels or not at all.
-        This is only used when ``link_type='wcs'``.
-        Choosing `None` may result in some Imviz functionality not working properly.
-
-    wcs_use_affine : bool
-        Use an affine transform to represent the offset between images if possible
-        (requires that the approximation is accurate to within 1 pixel with the
-        full WCS transformations). If approximation fails, it will automatically
-        fall back to full WCS transformation. This is only used when ``link_type='wcs'``.
-        Affine approximation is much more performant at the cost of accuracy.
-
-    error_on_fail : bool
-        If `True`, any failure in linking will raise an exception.
-        If `False`, warnings will be emitted as snackbar messages.
-        When only warnings are emitted and no links are assigned,
-        some Imviz functionality may not work properly.
-
-    update_plugin : bool
-        Whether to update the state of the "Orientation" plugin, if available.
-
-    Raises
-    ------
-    ValueError
-        Invalid inputs or reference data.
-
-    """
-    if len(app.data_collection) <= 1 and link_type != 'wcs':  # No need to link, we are done.
-        return
-
-    if link_type not in ('pixels', 'wcs'):
-        raise ValueError(f"link_type must be 'pixels' or 'wcs', got {link_type}")
-    if link_type == 'wcs' and wcs_fallback_scheme not in (None, 'pixels'):
-        raise ValueError("wcs_fallback_scheme must be None or 'pixels', "
-                         f"got {wcs_fallback_scheme}")
-    if link_type == 'wcs':
-        at_least_one_data_have_wcs = len([
-            hasattr(d, 'coords') and isinstance(d.coords, BaseHighLevelWCS)
-            for d in app.data_collection
-        ]) > 0
-        if not at_least_one_data_have_wcs:
-            if wcs_fallback_scheme is None:
-                if error_on_fail:
-                    raise ValueError("link_type can only be 'wcs' when wcs_fallback_scheme "
-                                     "is 'None' if all data have valid WCS.")
-                else:
-                    return
-            else:
-                # fall back on pixel linking
-                link_type = 'pixels'
-
-    # default reference layer is the first-loaded image in default viewer:
-    default_reference_layer = app._jdaviz_helper.default_viewer._obj.first_loaded_data
-    if default_reference_layer is None:  # No data in viewer, just use first in collection
-        default_reference_layer = app.data_collection[0]
-
-    # if the plugin exists, send a message so that the plugin's state is updated and spinner
-    # is shown (the plugin will make a call back here)
-    if 'imviz-orientation' in [item['name'] for item in app.state.tray_items]:
-        link_plugin = app.get_tray_item_from_name('imviz-orientation')
-        if update_plugin:
-            link_plugin.linking_in_progress = True
-    else:
-        link_plugin = None
-
-    data_already_linked = []
-    if link_type == app._link_type and wcs_use_affine == app._wcs_use_affine:
-        for link in app.data_collection.external_links:
-            if link.data1.label != _wcs_only_label:
-                data_already_linked.append(link.data2)
-    else:
-        for viewer in app._viewer_store.values():
-            if len(viewer._marktags):
-                raise ValueError(f"cannot change link_type (from '{app._link_type}' to "
-                                 f"'{link_type}') when markers are present. "
-                                 f" Clear markers with viewer.reset_markers() first")
-
-    old_link_type = getattr(app, '_link_type', None)
-
-    # if linking via WCS, add WCS-only reference data layer:
-    insert_base_wcs_layer = (
-        link_type == 'wcs' and
-        base_wcs_layer_label not in [d.label for d in app.data_collection]
-    )
-
-    if insert_base_wcs_layer:
-        degn = get_compass_info(default_reference_layer.coords, default_reference_layer.shape)[-3]
-        # Default rotation is the same orientation as the original reference data:
-        rotation_angle = -degn * u.deg
-        ndd = _get_rotated_nddata_from_label(
-            app, default_reference_layer.label, rotation_angle
-        )
-        app._jdaviz_helper.load_data(ndd, base_wcs_layer_label)
-
-        # set base layer to reference data in all viewers:
-        for viewer_id in app.get_viewer_ids():
-            app._change_reference_data(
-                base_wcs_layer_label, viewer_id=viewer_id
-            )
-
-    refdata, iref = get_reference_image_data(app)
-
-    # set internal tracking of link_type before changing reference data for anything that is
-    # subscribed to a change in reference data
-    app._link_type = link_type
-    app._wcs_use_affine = wcs_use_affine
-
-    if link_type == 'pixels' and old_link_type == 'wcs':
-        for viewer_id in app.get_viewer_ids():
-            app._change_reference_data(
-                default_reference_layer.label, viewer_id=viewer_id
-            )
-
-    links_list = []
-    ids0 = default_reference_layer.pixel_component_ids
-    ndim_range = range(default_reference_layer.ndim)
-
-    for i, data in enumerate(app.data_collection):
-        # Do not link with self
-        if i == iref:
-            continue
-
-        # We are not touching any existing Subsets. They keep their own links.
-        if not layer_is_2d(data):
-            continue
-
-        if data in data_already_linked:
-            # links already exist for this entry and we're not changing the type
-            continue
-
-        # We are not touching fake WCS layers in pixel linking.
-        if link_type == 'pixels' and data.meta.get('_WCS_ONLY'):
-            continue
-
-        ids1 = data.pixel_component_ids
-        new_links = []
-        try:
-            if link_type == 'pixels':
-                new_links = [LinkSame(ids0[i], ids1[i]) for i in ndim_range]
-            # otherwise if linking by WCS *and* this data entry has WCS:
-            elif hasattr(data.coords, 'pixel_to_world'):
-                wcslink = WCSLink(data1=refdata, data2=data, cids1=ids0, cids2=ids1)
-                if wcs_use_affine:
-                    try:
-                        new_links = [wcslink.as_affine_link()]
-                    except NoAffineApproximation:  # pragma: no cover
-                        new_links = [wcslink]
-                else:
-                    new_links = [wcslink]
-        except Exception as e:
-            if link_type == 'wcs' and wcs_fallback_scheme == 'pixels':
-                try:
-                    new_links = [LinkSame(ids0[i], ids1[i]) for i in ndim_range]
-                except Exception as e:  # pragma: no cover
-                    if error_on_fail:
-                        raise
-                    else:
-                        app.hub.broadcast(SnackbarMessage(
-                            f"Error linking '{data.label}' to '{refdata.label}': "
-                            f"{repr(e)}", color="warning", timeout=8000, sender=app))
-                        continue
-            else:
-                if error_on_fail:
-                    raise
-                else:
-                    app.hub.broadcast(SnackbarMessage(
-                        f"Error linking '{data.label}' to '{refdata.label}': "
-                        f"{repr(e)}", color="warning", timeout=8000, sender=app))
-                    continue
-        links_list += new_links
-
-    if len(links_list) > 0:
-        with app.data_collection.delay_link_manager_update():
-            if len(data_already_linked):
-                app.data_collection.add_link(links_list)
-            else:
-                app.data_collection.set_links(links_list)
-
-        app.hub.broadcast(SnackbarMessage(
-            'Images successfully relinked', color='success', timeout=8000, sender=app))
-
-    for viewer in app._viewer_store.values():
-        wcs_linked = link_type == 'wcs'
-        # viewer-state needs to know link type for reset_limits behavior
-        viewer.state.linked_by_wcs = wcs_linked
-        # also need to store a copy in the viewer item for the data dropdown to access
-        viewer_item = app._get_viewer_item(viewer.reference)
-
-        viewer_item['reference_data_label'] = refdata.label
-        viewer_item['linked_by_wcs'] = wcs_linked
-
-        # if changing from one link type to another, reset the limits:
-        if link_type != old_link_type:
-            viewer.state.reset_limits()
-
-    if link_plugin is not None:
-        # Only broadcast after success.
-        app.hub.broadcast(LinkUpdatedMessage(link_type,
-                                             wcs_fallback_scheme == 'pixels',
-                                             wcs_use_affine,
-                                             sender=app))
-
-        # reset the progress spinner
-        link_plugin.linking_in_progress = False
+    return None, -1
diff --git a/jdaviz/configs/imviz/plugins/orientation/orientation.py b/jdaviz/configs/imviz/plugins/orientation/orientation.py
index a3e7b016..e15fa701 100644
--- a/jdaviz/configs/imviz/plugins/orientation/orientation.py
+++ b/jdaviz/configs/imviz/plugins/orientation/orientation.py
@@ -1,22 +1,25 @@
-from traitlets import List, Unicode, Bool, Dict, observe
-
+from astropy import units as u
+from astropy.wcs.wcsapi import BaseHighLevelWCS
+from glue.core.link_helpers import LinkSame
 from glue.core.message import (
     DataCollectionAddMessage, SubsetCreateMessage, SubsetDeleteMessage
 )
 from glue.core.subset import Subset
 from glue.core.subset_group import GroupedSubset
+from glue.plugins.wcs_autolinking.wcs_autolinking import WCSLink, NoAffineApproximation
+from traitlets import List, Unicode, Bool, Dict, observe
 
-import astropy.units as u
-from jdaviz.configs.imviz.helper import link_image_data, base_wcs_layer_label
+from jdaviz.configs.imviz.helper import get_reference_image_data, layer_is_2d
 from jdaviz.configs.imviz.wcs_utils import (
     get_compass_info, _get_rotated_nddata_from_label
 )
+from jdaviz.core.custom_traitlets import FloatHandleEmpty
 from jdaviz.core.events import (
-    LinkUpdatedMessage, ExitBatchLoadMessage, ChangeRefDataMessage,
+    ExitBatchLoadMessage, ChangeRefDataMessage,
     AstrowidgetMarkersChangedMessage, MarkersPluginUpdate,
-    SnackbarMessage, ViewerAddedMessage, AddDataMessage
+    SnackbarMessage, ViewerAddedMessage, AddDataMessage, LinkUpdatedMessage
 )
-from jdaviz.core.custom_traitlets import FloatHandleEmpty
+
 from jdaviz.core.registries import tray_registry
 from jdaviz.core.template_mixin import (
     PluginTemplateMixin, SelectPluginComponent, LayerSelect, ViewerSelectMixin, AutoTextField
@@ -26,6 +29,7 @@ from jdaviz.utils import _wcs_only_label
 
 __all__ = ['Orientation']
 
+base_wcs_layer_label = 'Default orientation'
 link_type_msg_to_trait = {'pixels': 'Pixels', 'wcs': 'WCS'}
 
 
@@ -104,9 +108,6 @@ class Orientation(PluginTemplateMixin, ViewerSelectMixin):
             self, 'new_layer_label', 'new_layer_label_default', 'new_layer_label_auto', None
         )
 
-        self.hub.subscribe(self, LinkUpdatedMessage,
-                           handler=self._on_link_updated)
-
         self.hub.subscribe(self, DataCollectionAddMessage,
                            handler=self._on_new_app_data)
 
@@ -147,20 +148,24 @@ class Orientation(PluginTemplateMixin, ViewerSelectMixin):
             )
         )
 
-    def _on_link_updated(self, msg):
-        self.link_type.selected = link_type_msg_to_trait[msg.link_type]
-        self.linking_in_progress = True
-        self.wcs_use_fallback = msg.wcs_use_fallback
-        self.wcs_use_affine = msg.wcs_use_affine
-
     def _link_image_data(self):
-        link_image_data(
-            self.app,
-            link_type=self.link_type.selected.lower(),
-            wcs_fallback_scheme='pixels' if self.wcs_use_fallback else None,
-            wcs_use_affine=self.wcs_use_affine,
-            error_on_fail=False,
-            update_plugin=False)
+        self.linking_in_progress = True
+        try:
+            link_type = self.link_type.selected.lower()
+            link_image_data(
+                self.app,
+                link_type=link_type,
+                wcs_fallback_scheme='pixels' if self.wcs_use_fallback else None,
+                wcs_use_affine=self.wcs_use_affine,
+                error_on_fail=False)
+        except Exception:  # pragma: no cover
+            raise
+        else:
+            # Only broadcast after success.
+            self.app.hub.broadcast(LinkUpdatedMessage(
+                link_type, self.wcs_use_fallback, self.wcs_use_affine, sender=self.app))
+        finally:
+            self.linking_in_progress = False
 
     def _check_if_data_with_wcs_exists(self):
         for data in self.app.data_collection:
@@ -203,14 +208,15 @@ class Orientation(PluginTemplateMixin, ViewerSelectMixin):
         if self.linking_in_progress:
             return
 
+        self.linking_in_progress = True  # Prevent recursion
+
         if self.need_clear_subsets:
+            self.linking_in_progress = False
             raise ValueError("Link type can only be changed after existing subsets "
                              f"are deleted, but {len(self.app.data_collection.subset_groups)} "
                              f"subset(s) still exist. To delete them, you can use "
                              f"`delete_subsets()` from the plugin API.")
 
-        self.linking_in_progress = True
-
         if self.need_clear_astrowidget_markers:
             setattr(self, msg.get('name'), msg.get('old'))
             self.linking_in_progress = False
@@ -220,6 +226,7 @@ class Orientation(PluginTemplateMixin, ViewerSelectMixin):
         if self.link_type.selected == 'Pixels':
             self.wcs_use_affine = True
 
+        self.linking_in_progress = False
         self._link_image_data()
 
         # load data into the viewer that are now compatible with the
@@ -244,7 +251,6 @@ class Orientation(PluginTemplateMixin, ViewerSelectMixin):
         if wcs_linked:
             self._send_wcs_layers_to_all_viewers()
 
-        self.linking_in_progress = False
         self._update_layer_label_default()
 
         # Clear previous zoom limits because they no longer mean anything.
@@ -519,3 +525,222 @@ class Orientation(PluginTemplateMixin, ViewerSelectMixin):
             f'CCW {self.rotation_angle_deg():.2f} ' +
             ('(E-left)' if self.east_left else '(E-right)')
         )
+
+
+def link_image_data(app, link_type='pixels', wcs_fallback_scheme=None, wcs_use_affine=True,
+                    error_on_fail=False):
+    """(Re)link loaded data in Imviz with the desired link type.
+
+    .. note::
+
+        Any markers added in Imviz will need to be removed manually before changing linking type.
+        You can add back the markers using
+        :meth:`~jdaviz.core.astrowidgets_api.AstrowidgetsImageViewerMixin.add_markers`
+        for the relevant viewer(s).
+
+    Parameters
+    ----------
+    app : `~jdaviz.app.Application`
+        Application associated with Imviz, e.g., ``imviz.app``.
+
+    link_type : {'pixels', 'wcs'}
+        Choose to link by pixels or WCS.
+
+    wcs_fallback_scheme : {None, 'pixels'}
+        If WCS linking failed, choose to fall back to linking by pixels or not at all.
+        This is only used when ``link_type='wcs'``.
+        Choosing `None` may result in some Imviz functionality not working properly.
+
+    wcs_use_affine : bool
+        Use an affine transform to represent the offset between images if possible
+        (requires that the approximation is accurate to within 1 pixel with the
+        full WCS transformations). If approximation fails, it will automatically
+        fall back to full WCS transformation. This is only used when ``link_type='wcs'``.
+        Affine approximation is much more performant at the cost of accuracy.
+
+    error_on_fail : bool
+        If `True`, any failure in linking will raise an exception.
+        If `False`, warnings will be emitted as snackbar messages.
+        When only warnings are emitted and no links are assigned,
+        some Imviz functionality may not work properly.
+
+    Raises
+    ------
+    ValueError
+        Invalid inputs or reference data.
+
+    """
+    if len(app.data_collection) <= 1 and link_type != 'wcs':  # No need to link, we are done.
+        return
+
+    if link_type not in ('pixels', 'wcs'):  # pragma: no cover
+        raise ValueError(f"link_type must be 'pixels' or 'wcs', got {link_type}")
+    if link_type == 'wcs' and wcs_fallback_scheme not in (None, 'pixels'):  # pragma: no cover
+        raise ValueError("wcs_fallback_scheme must be None or 'pixels', "
+                         f"got {wcs_fallback_scheme}")
+    if link_type == 'wcs':
+        at_least_one_data_have_wcs = len([
+            hasattr(d, 'coords') and isinstance(d.coords, BaseHighLevelWCS)
+            for d in app.data_collection
+        ]) > 0
+        if not at_least_one_data_have_wcs:  # pragma: no cover
+            if wcs_fallback_scheme is None:
+                if error_on_fail:
+                    raise ValueError("link_type can only be 'wcs' when wcs_fallback_scheme "
+                                     "is 'None' if at least one image has a valid WCS.")
+                else:
+                    return
+            else:
+                # fall back on pixel linking
+                link_type = 'pixels'
+
+    old_link_type = getattr(app, '_link_type', None)
+
+    # In WCS linking, changing orientation layer is done within Orientation plugin,
+    # so here we assume viewer.state.reference_data is already the desired
+    # orientation by the time this function is called.
+    #
+    # In pixels linking, Affine approximation does not matter.
+    #
+    # data1 = reference, data2 = actual data
+    data_already_linked = []
+    if (link_type == old_link_type and
+            (link_type == "pixels" or wcs_use_affine == app._wcs_use_affine)):
+        # We are only here to link new data with existing configuration,
+        # so no need to relink existing data.
+        for link in app.data_collection.external_links:
+            data_already_linked.append(link.data2)
+    else:  # pragma: no cover
+        # Everything has to be relinked.
+        for viewer in app._viewer_store.values():
+            if len(viewer._marktags):
+                raise ValueError(f"cannot change link_type (from '{app._link_type}' to "
+                                 f"'{link_type}') when markers are present. "
+                                 f" Clear markers with viewer.reset_markers() first")
+
+    # set internal tracking of link_type before changing reference data for anything that is
+    # subscribed to a change in reference data
+    app._link_type = link_type
+    app._wcs_use_affine = wcs_use_affine
+
+    # wcs -> pixels: First loaded real data will be reference.
+    if link_type == 'pixels' and old_link_type == 'wcs':
+        # default reference layer is the first-loaded image in default viewer:
+        refdata = app._jdaviz_helper.default_viewer._obj.first_loaded_data
+        if refdata is None:  # No data in viewer, just use first in collection  # pragma: no cover
+            iref = 0
+            refdata = app.data_collection[iref]
+        else:
+            iref = app.data_collection.index(refdata)
+
+        # set default layer to reference data in all viewers:
+        for viewer_id in app.get_viewer_ids():
+            app._change_reference_data(refdata.label, viewer_id=viewer_id)
+
+    # pixels -> wcs: Always the default orientation
+    elif link_type == 'wcs' and old_link_type == 'pixels':
+        # Have to create the default orientation first.
+        if base_wcs_layer_label not in app.data_collection.labels:
+            default_reference_layer = (app._jdaviz_helper.default_viewer._obj.first_loaded_data
+                                       or app.data_collection[0])
+            degn = get_compass_info(default_reference_layer.coords, default_reference_layer.shape)[-3]  # noqa: E501
+            # Default rotation is the same orientation as the original reference data:
+            rotation_angle = -degn * u.deg
+            ndd = _get_rotated_nddata_from_label(app, default_reference_layer.label, rotation_angle)
+            app._jdaviz_helper.load_data(ndd, base_wcs_layer_label)
+
+        # set default layer to reference data in all viewers:
+        for viewer_id in app.get_viewer_ids():
+            app._change_reference_data(base_wcs_layer_label, viewer_id=viewer_id)
+
+        refdata = app.data_collection[base_wcs_layer_label]
+        iref = app.data_collection.index(refdata)
+
+    # Re-use current reference data.
+    else:
+        refdata, iref = get_reference_image_data(app)
+        # App just loaded, nothing yet, so take first image.
+        if refdata is None:
+            refdata = app._jdaviz_helper.default_viewer._obj.first_loaded_data
+            if refdata is None:  # No data in viewer, just use first in collection
+                iref = 0
+                refdata = app.data_collection[iref]
+            else:   # pragma: no cover
+                iref = app.data_collection.index(refdata)
+
+    # With reference data changed, if needed, now we relink as needed.
+
+    links_list = []
+    ids0 = refdata.pixel_component_ids
+    ndim_range = range(2)  # We only support 2D
+
+    for i, data in enumerate(app.data_collection):
+        # 1. Do not link with self.
+        # 2. We are not touching any existing Subsets or Table. They keep their own links.
+        # 3. Links already exist for this entry and we're not changing the type.
+        # 4. We are not touching fake WCS layers in pixel linking.
+        # 5. We are not touching data without WCS in WCS linking.
+        if ((i == iref) or (not layer_is_2d(data)) or (data in data_already_linked) or
+                (link_type == "pixels" and data.meta.get(_wcs_only_label)) or
+                (link_type == "wcs" and not hasattr(data.coords, 'pixel_to_world'))):
+            continue
+
+        ids1 = data.pixel_component_ids
+        new_links = []
+        try:
+            if link_type == 'pixels':
+                new_links = [LinkSame(ids0[i], ids1[i]) for i in ndim_range]
+            else:  # wcs
+                wcslink = WCSLink(data1=refdata, data2=data, cids1=ids0, cids2=ids1)
+                if wcs_use_affine:
+                    try:
+                        new_links = [wcslink.as_affine_link()]
+                    except NoAffineApproximation:  # pragma: no cover
+                        new_links = [wcslink]
+                else:
+                    new_links = [wcslink]
+        except Exception as e:  # pragma: no cover
+            if link_type == 'wcs' and wcs_fallback_scheme == 'pixels':
+                try:
+                    new_links = [LinkSame(ids0[i], ids1[i]) for i in ndim_range]
+                except Exception as e:  # pragma: no cover
+                    if error_on_fail:
+                        raise
+                    else:
+                        app.hub.broadcast(SnackbarMessage(
+                            f"Error linking '{data.label}' to '{refdata.label}': "
+                            f"{repr(e)}", color="warning", timeout=8000, sender=app))
+                        continue
+            else:  # pragma: no cover
+                if error_on_fail:
+                    raise
+                else:
+                    app.hub.broadcast(SnackbarMessage(
+                        f"Error linking '{data.label}' to '{refdata.label}': "
+                        f"{repr(e)}", color="warning", timeout=8000, sender=app))
+                    continue
+        links_list += new_links
+
+    if len(links_list) > 0:
+        with app.data_collection.delay_link_manager_update():
+            if len(data_already_linked):
+                app.data_collection.add_link(links_list)  # Add to existing
+            else:
+                app.data_collection.set_links(links_list)  # Redo all links
+
+        app.hub.broadcast(SnackbarMessage(
+            'Images successfully relinked', color='success', timeout=8000, sender=app))
+
+    for viewer in app._viewer_store.values():
+        wcs_linked = link_type == 'wcs'
+        # viewer-state needs to know link type for reset_limits behavior
+        viewer.state.linked_by_wcs = wcs_linked
+        # also need to store a copy in the viewer item for the data dropdown to access
+        viewer_item = app._get_viewer_item(viewer.reference)
+
+        viewer_item['reference_data_label'] = refdata.label
+        viewer_item['linked_by_wcs'] = wcs_linked
+
+        # if changing from one link type to another, reset the limits:
+        if link_type != old_link_type:
+            viewer.state.reset_limits()


Test patch:
diff --git a/jdaviz/configs/default/plugins/markers/tests/test_markers_plugin.py b/jdaviz/configs/default/plugins/markers/tests/test_markers_plugin.py
index 19b51623..f388f7f9 100644
--- a/jdaviz/configs/default/plugins/markers/tests/test_markers_plugin.py
+++ b/jdaviz/configs/default/plugins/markers/tests/test_markers_plugin.py
@@ -20,7 +20,7 @@ def _assert_dict_allclose(dict1, dict2):
             assert v == dict2.get(k)
 
 
-def test_markers_cubeviz(cubeviz_helper, spectrum1d_cube):
+def test_markers_cubeviz(tmp_path, cubeviz_helper, spectrum1d_cube):
     cubeviz_helper.load_data(spectrum1d_cube, "test")
     fv = cubeviz_helper.app.get_viewer('flux-viewer')
     sv = cubeviz_helper.app.get_viewer('spectrum-viewer')
@@ -130,6 +130,7 @@ def test_markers_cubeviz(cubeviz_helper, spectrum1d_cube):
 
     # appears as option in export plugin and exports successfully
     assert "Markers:table" in exp.table.choices
+    exp.filename = str(tmp_path / "cubeviz_export.ecsv")
     exp.table = "Markers:table"
     exp.export()
 
diff --git a/jdaviz/configs/imviz/tests/test_linking.py b/jdaviz/configs/imviz/tests/test_linking.py
index c3bdf8f2..21e69bf2 100644
--- a/jdaviz/configs/imviz/tests/test_linking.py
+++ b/jdaviz/configs/imviz/tests/test_linking.py
@@ -286,8 +286,9 @@ class TestLink_GWCS_GWCS(BaseImviz_GWCS_GWCS):
 
 
 def test_imviz_no_data(imviz_helper):
-    with pytest.raises(ValueError, match='No valid reference data'):
-        get_reference_image_data(imviz_helper.app)
+    refdata, iref = get_reference_image_data(imviz_helper.app)
+    assert refdata is None
+    assert iref == -1
 
     imviz_helper.link_data()  # Just no-op, do not crash
     links = imviz_helper.app.data_collection.external_links
diff --git a/jdaviz/configs/imviz/tests/test_orientation.py b/jdaviz/configs/imviz/tests/test_orientation.py
index c580ff5b..bbeceeaa 100644
--- a/jdaviz/configs/imviz/tests/test_orientation.py
+++ b/jdaviz/configs/imviz/tests/test_orientation.py
@@ -147,12 +147,9 @@ class TestNonDefaultOrientation(BaseImviz_WCS_WCS):
         assert self.viewer.state.reference_data.label == "has_wcs_1[SCI,1]"
         assert viewer_2.state.reference_data.label == "has_wcs_1[SCI,1]"
 
-        # FIXME: spacetelescope/jdaviz#2724 (remove AssertionError)
         lc_plugin.link_type = 'WCS'
-        with pytest.raises(AssertionError):
-            assert self.viewer.state.reference_data.label == "Default orientation"
-        with pytest.raises(AssertionError):
-            assert viewer_2.state.reference_data.label == "Default orientation"
+        assert self.viewer.state.reference_data.label == "Default orientation"
+        assert viewer_2.state.reference_data.label == "Default orientation"
 
     def test_custom_orientation(self):
         lc_plugin = self.imviz.plugins['Orientation']
diff --git a/jdaviz/configs/imviz/tests/test_wcs_utils.py b/jdaviz/configs/imviz/tests/test_wcs_utils.py
index 0db7de54..5e0b05d3 100644
--- a/jdaviz/configs/imviz/tests/test_wcs_utils.py
+++ b/jdaviz/configs/imviz/tests/test_wcs_utils.py
@@ -9,7 +9,7 @@ from gwcs import coordinate_frames as cf
 from numpy.testing import assert_allclose
 
 from jdaviz.configs.imviz import wcs_utils
-from jdaviz.configs.imviz.helper import base_wcs_layer_label
+from jdaviz.configs.imviz.plugins.orientation.orientation import base_wcs_layer_label
 from jdaviz.configs.imviz.tests.utils import BaseImviz_WCS_GWCS, create_example_gwcs