Repository: astropy/astropy

Issue:
Support custom mapping from UCD to WSCAxes properties for custom coordinate frames
### What is the problem this feature will solve?

The overall aim is to have nice plot with proper `ax.coords` setting e.g. `coord_wrap`, `format_unit`, `coord_type` for custom Coordinate frames.

Currently half of the is supported, `fitswcs.py` supports custom ctype to `UCD` mappings via `CTYPE_TO_UCD1_CUSTOM` but then on the other side `wscapi.py` has a hard coded list in `transform_coord_meta_from_wcs` therefore `WCSAxes` doesn't recognise the custom UCD.

So currently this can be done
```
MYFRAME_CTYPE_TO_UCD1 = {
    "SXLT": "custom:pos.myframe.lat",
    "SXLN": "custom:pos.myframe.lon",
    "SXRZ": "custom:pos.myframe.z"
}
astropy.wcs.wcsapi.fitswcs.CTYPE_TO_UCD1_CUSTOM.update(MYFRAME_CTYPE_TO_UCD1)
```
but there is no way to include settings for `custom:pos.myframe.X` in this section of code https://github.com/astropy/astropy/blob/cf99b9e7761c45c15ddc060a76b563574f525df9/astropy/visualization/wcsaxes/wcsapi.py#L54-L71

where Id like to add something like
```
            if "custom:pos.MYFRAME.lon" in axis_type:
                coord_wrap = 180.0 * u.deg
                format_unit = u.arcsec
                coord_type = "longitude"
            elif "custom:pos.MYFARME.lat" in axis_type:
                format_unit = u.arcsec
                coord_type = "latitude"
```

### Describe the desired outcome

Be able to register or other wise include custom plot setting for custom UCDs in a similar manner to how custom CTYPE to UCD mapping can be included.

### Additional context

_No response_

Discussion:
github-actions[bot]: Thank you for your contribution to Astropy! 🌌 This checklist is meant to remind the package maintainers who will review this pull request of some common things to look for.

  - [ ] Do the proposed changes actually accomplish desired goals?
  - [ ] Do the proposed changes follow the [Astropy coding guidelines](https://docs.astropy.org/en/latest/development/codeguide.html)?
  - [ ] Are tests added/updated as required? If so, do they follow the [Astropy testing guidelines](https://docs.astropy.org/en/latest/development/testguide.html)?
  - [ ] Are docs added/updated as required? If so, do they follow the [Astropy documentation guidelines](https://docs.astropy.org/en/latest/development/docguide.html#astropy-documentation-rules-and-guidelines)?
  - [ ] Is rebase and/or squash necessary? If so, please provide the author with appropriate instructions. Also see instructions for [rebase](https://docs.astropy.org/en/latest/development/workflow/development_workflow.html#rebase-if-necessary) and [squash](https://docs.astropy.org/en/latest/development/workflow/development_workflow.html#squash-if-necessary).
  - [ ] Did the CI pass? If no, are the failures related? If you need to run daily and weekly cron jobs as part of the PR, please apply the "Extra CI" label. Codestyle issues can be fixed by the [bot](https://docs.astropy.org/en/latest/development/workflow/development_workflow.html#pre-commit).
  - [ ] Is a change log needed? If yes, did the change log check pass? If no, add the "no-changelog-entry-needed" label. If this is a manual backport, use the "skip-changelog-checks" label unless special changelog handling is necessary.
  - [ ] Is this a big PR that makes a "What's new?" entry worthwhile and if so, is (1) a "what's new" entry included in this PR and (2) the "whatsnew-needed" label applied?
  - [ ] At the time of adding the milestone, if the milestone set requires a backport to release branch(es), apply the appropriate "backport-X.Y.x" label(s) *before* merge.
github-actions[bot]: 👋 Thank you for your draft pull request! Do you know that you can use `[ci skip]` or `[skip ci]` in your commit messages to skip running continuous integration tests until you are ready?
github-actions[bot]:  Hi humans :wave: - this pull request hasn't had any new commits for approximately 4 months. **I plan to close this in 30 days if the pull request doesn't have any new commits by then.**

In lieu of a stalled pull request, please consider closing this and open an issue instead if a reminder is needed to revisit in the future. Maintainers may also choose to add **keep-open** label to keep this PR open but it is discouraged unless absolutely necessary.

If this PR still needs to be reviewed, as an author, you can rebase it to reset the clock.

*If you believe I commented on this pull request incorrectly, please report this [here](https://github.com/pllim/action-astropy-stalebot/issues).*

Patch:
diff --git a/astropy/visualization/wcsaxes/__init__.py b/astropy/visualization/wcsaxes/__init__.py
index a47400df96..18bf323780 100644
--- a/astropy/visualization/wcsaxes/__init__.py
+++ b/astropy/visualization/wcsaxes/__init__.py
@@ -18,6 +18,7 @@ from .coordinates_map import CoordinatesMap
 from .core import *
 from .helpers import *
 from .patches import *
+from .wcsapi import custom_ucd_coord_meta_mapping
 
 
 class Conf(_config.ConfigNamespace):
diff --git a/astropy/visualization/wcsaxes/wcsapi.py b/astropy/visualization/wcsaxes/wcsapi.py
index 1772b77806..c3eb27b4ae 100644
--- a/astropy/visualization/wcsaxes/wcsapi.py
+++ b/astropy/visualization/wcsaxes/wcsapi.py
@@ -1,5 +1,7 @@
 # Functions/classes for WCSAxes related to APE14 WCSes
 
+from contextlib import contextmanager
+
 import numpy as np
 
 from astropy import units as u
@@ -15,6 +17,7 @@ __all__ = [
     "transform_coord_meta_from_wcs",
     "WCSWorld2PixelTransform",
     "WCSPixel2WorldTransform",
+    "custom_ucd_coord_meta_mapping",
 ]
 
 IDENTITY = WCS(naxis=2)
@@ -23,6 +26,94 @@ IDENTITY.wcs.crval = [0.0, 0.0]
 IDENTITY.wcs.crpix = [1.0, 1.0]
 IDENTITY.wcs.cdelt = [1.0, 1.0]
 
+UCD_COORD_META_MAPPING = {
+    "lon": {"coord_type": "longitude"},
+    "lat": {"coord_type": "latitude"},
+    "ra": {"coord_type": "longitude", "format_unit": u.hourangle},
+    "dec": {"coord_type": "latitude"},
+    "alt": {"coord_type": "longitude"},
+    "az": {"coord_type": "latitude"},
+    "long": {"coord_type": "longitude"},
+}
+
+CUSTOM_UCD_COORD_META_MAPPING = {
+    "pos.helioprojective.lon": {
+        "coord_wrap": 180.0 * u.deg,
+        "format_unit": u.arcsec,
+        "coord_type": "longitude",
+    },
+    "pos.helioprojective.lat": {"format_unit": u.arcsec, "coord_type": "latitude"},
+    "pos.heliographic.stonyhurst.lon": {
+        "coord_wrap": 180.0 * u.deg,
+        "format_unit": u.deg,
+        "coord_type": "longitude",
+    },
+    "pos.heliographic.stonyhurst.lat": {"format_unit": u.deg, "coord_type": "latitude"},
+    "pos.heliographic.carrington.lon": {
+        "coord_wrap": 360.0 * u.deg,
+        "format_unit": u.deg,
+        "coord_type": "longitude",
+    },
+    "pos.heliographic.carrington.lat": {"format_unit": u.deg, "coord_type": "latitude"},
+}
+
+
+@contextmanager
+def custom_ucd_coord_meta_mapping(mapping, *, overwrite=False):
+    """
+    A context manager that makes it possible to temporarily add new UCD+ to WCS coordinate
+    plot metadata mappings.
+
+    Parameters
+    ----------
+    mapping : dict
+        A dictionary mapping a UCD to coordinate plot metadata.
+        Note that custom UCD names have their "custom:" prefix stripped.
+    overwrite : bool
+        If `True` overwrite existing entries with ``mapping``.
+
+    Examples
+    --------
+    >>> from matplotlib import pyplot as plt
+    >>> from astropy.visualization.wcsaxes.wcsapi import custom_ucd_coord_meta_mapping
+    >>> from astropy.wcs.wcsapi.fitswcs import custom_ctype_to_ucd_mapping
+    >>> wcs = WCS(naxis=1)
+    >>> wcs.wcs.ctype = ["eggs"]
+    >>> wcs.wcs.cunit = ["deg"]
+    >>> custom_mapping = {"eggs": "custom:pos.eggs"}
+    >>> with custom_ctype_to_ucd_mapping(custom_mapping):
+    ...     custom_meta = {
+    ...         "pos.eggs": {
+    ...             "coord_wrap": 360.0 * u.deg,
+    ...             "format_unit": u.arcsec,
+    ...             "coord_type": "longitude",
+    ...         }
+    ...     }
+    ...     with custom_ucd_coord_meta_mapping(custom_meta):
+    ...        fig = plt.figure()
+    ...        ax = fig.add_subplot(111, projection=wcs)
+    ...        ax.coords
+    <CoordinatesMap with 1 world coordinates:
+    <BLANKLINE>
+      index       aliases           type   unit  wrap format_unit visible
+                                                 deg
+      ----- -------------------- --------- ---- ----- ----------- -------
+          0 custom:pos.eggs eggs longitude  deg 360.0      arcsec     yes
+    <BLANKLINE>
+    >
+    """
+    for k, v in mapping.items():
+        k = k.removeprefix("custom:")
+        if not overwrite and k in CUSTOM_UCD_COORD_META_MAPPING:
+            raise ValueError(f"UCD metadata mapping {k} already exists.")
+        CUSTOM_UCD_COORD_META_MAPPING[k] = v
+
+    yield
+
+    for k in mapping.keys():
+        k = k.removeprefix("custom:")
+        del CUSTOM_UCD_COORD_META_MAPPING[k]
+
 
 def transform_coord_meta_from_wcs(wcs, frame_class, slices=None):
     if slices is not None:
@@ -56,56 +147,35 @@ def transform_coord_meta_from_wcs(wcs, frame_class, slices=None):
         axis_unit = u.Unit(wcs.world_axis_units[idx])
         coord_wrap = None
         format_unit = axis_unit
-
         coord_type = "scalar"
 
+        dim_meta = {
+            "coord_type": coord_type,
+            "coord_wrap": coord_wrap,
+            "format_unit": format_unit,
+            "axis_unit": axis_unit,
+        }
+
         if axis_type is not None:
             axis_type_split = axis_type.split(".")
+
             if len(axis_type_split):
                 axis_type_split[0] = axis_type_split[0].replace("custom:", "")
 
-            if "pos.helioprojective.lon" in axis_type:
-                coord_wrap = 180.0 * u.deg
-                format_unit = u.arcsec
-                coord_type = "longitude"
-            elif "pos.helioprojective.lat" in axis_type:
-                format_unit = u.arcsec
-                coord_type = "latitude"
-            elif "pos.heliographic.stonyhurst.lon" in axis_type:
-                coord_wrap = 180.0 * u.deg
-                format_unit = u.deg
-                coord_type = "longitude"
-            elif "pos.heliographic.stonyhurst.lat" in axis_type:
-                format_unit = u.deg
-                coord_type = "latitude"
-            elif "pos.heliographic.carrington.lon" in axis_type:
-                coord_wrap = 360.0 * u.deg
-                format_unit = u.deg
-                coord_type = "longitude"
-            elif "pos.heliographic.carrington.lat" in axis_type:
-                format_unit = u.deg
-                coord_type = "latitude"
-            elif "pos" in axis_type_split:
-                if "lon" in axis_type_split:
-                    coord_type = "longitude"
-                elif "lat" in axis_type_split:
-                    coord_type = "latitude"
-                elif "ra" in axis_type_split:
-                    coord_type = "longitude"
-                    format_unit = u.hourangle
-                elif "dec" in axis_type_split:
-                    coord_type = "latitude"
-                elif "alt" in axis_type_split:
-                    coord_type = "longitude"
-                elif "az" in axis_type_split:
-                    coord_type = "latitude"
-                elif "long" in axis_type_split:
-                    coord_type = "longitude"
-
-        coord_meta["type"].append(coord_type)
-        coord_meta["wrap"].append(coord_wrap)
-        coord_meta["format_unit"].append(format_unit)
-        coord_meta["unit"].append(axis_unit)
+            for ucd, meta in CUSTOM_UCD_COORD_META_MAPPING.items():
+                if ucd in axis_type:
+                    dim_meta.update(meta)
+                    break
+            else:
+                for ucd, meta in UCD_COORD_META_MAPPING.items():
+                    if ucd == axis_type_split[-1]:
+                        dim_meta.update(meta)
+                        break
+
+        coord_meta["type"].append(dim_meta["coord_type"])
+        coord_meta["wrap"].append(dim_meta["coord_wrap"])
+        coord_meta["format_unit"].append(dim_meta["format_unit"])
+        coord_meta["unit"].append(dim_meta["axis_unit"])
 
         # For FITS-WCS, for backward-compatibility, we need to make sure that we
         # provide aliases based on CTYPE for the name.
diff --git a/docs/changes/visualization/16347.feature.rst b/docs/changes/visualization/16347.feature.rst
new file mode 100644
index 0000000000..a65c158e34
--- /dev/null
+++ b/docs/changes/visualization/16347.feature.rst
@@ -0,0 +1,2 @@
+Add support for custom coordinate frames for ``WCSAxes`` through a context
+manager ``astropy.visualization.wcsaxes.custom_ucd_coord_meta_mapping``.


Test patch:
diff --git a/astropy/visualization/wcsaxes/tests/test_wcsapi.py b/astropy/visualization/wcsaxes/tests/test_wcsapi.py
index 6d6d926157..e7c310dbde 100644
--- a/astropy/visualization/wcsaxes/tests/test_wcsapi.py
+++ b/astropy/visualization/wcsaxes/tests/test_wcsapi.py
@@ -19,6 +19,7 @@ from astropy.visualization.wcsaxes.frame import RectangularFrame, RectangularFra
 from astropy.visualization.wcsaxes.wcsapi import (
     WCSWorld2PixelTransform,
     apply_slices,
+    custom_ucd_coord_meta_mapping,
     transform_coord_meta_from_wcs,
 )
 from astropy.wcs import WCS
@@ -266,6 +267,127 @@ def test_coord_type_from_ctype(cube_wcs):
     assert coord_meta["format_unit"] == [u.one, u.one]
     assert coord_meta["wrap"] == [None, None]
 
+    myframe_mapping = {
+        "custom:pos.myframe.lon": {
+            "coord_wrap": 180.0 * u.deg,
+            "format_unit": u.arcsec,
+            "coord_type": "longitude",
+        },
+        "custom:pos.myframe.lat": {"format_unit": u.arcsec, "coord_type": "latitude"},
+    }
+
+
+def test_custom_coord_type_from_ctype():
+    wcs = WCS(naxis=1)
+    wcs.wcs.ctype = ["eggs"]
+    wcs.wcs.cunit = ["deg"]
+
+    custom_mapping = {
+        "eggs": "custom:pos.eggs",
+    }
+    with custom_ctype_to_ucd_mapping(custom_mapping):
+        fig = plt.figure()
+        ax = fig.add_subplot(111, projection=wcs)
+        assert ax.coords["eggs"].coord_type == "scalar"
+        assert ax.coords["eggs"].coord_wrap == None
+        assert ax.coords["eggs"].get_format_unit() == u.deg
+
+        custom_meta = {
+            "pos.eggs": {
+                "coord_wrap": 360.0 * u.deg,
+                "format_unit": u.arcsec,
+                "coord_type": "longitude",
+            }
+        }
+        with custom_ucd_coord_meta_mapping(custom_meta):
+            ax = fig.add_subplot(111, projection=wcs)
+            ax.coords
+            assert ax.coords["eggs"].coord_type == "longitude"
+            assert ax.coords["eggs"].coord_wrap == 360 * u.deg
+            assert ax.coords["eggs"].get_format_unit() == u.arcsec
+
+        fig = plt.figure()
+        ax = fig.add_subplot(111, projection=wcs)
+        assert ax.coords["eggs"].coord_type == "scalar"
+        assert ax.coords["eggs"].coord_wrap == None
+        assert ax.coords["eggs"].get_format_unit() == u.deg
+
+
+def test_custom_coord_type_from_ctype_nested():
+    wcs = WCS(naxis=2)
+    wcs.wcs.ctype = ["eggs", "spam"]
+    wcs.wcs.cunit = ["deg", "deg"]
+
+    custom_mapping = {
+        "eggs": "custom:pos.eggs",
+        "spam": "custom:pos.spam",
+    }
+
+    with custom_ctype_to_ucd_mapping(custom_mapping):
+        fig = plt.figure()
+        custom_meta_1 = {
+            "pos.eggs": {
+                "coord_wrap": 360.0 * u.deg,
+                "format_unit": u.arcsec,
+                "coord_type": "longitude",
+            }
+        }
+        with custom_ucd_coord_meta_mapping(custom_meta_1):
+            custom_meta_2 = {
+                "pos.spam": {
+                    "format_unit": u.deg,
+                    "coord_type": "latitude",
+                }
+            }
+            with custom_ucd_coord_meta_mapping(custom_meta_2):
+                ax = fig.add_subplot(111, projection=wcs)
+                ax.coords
+                assert ax.coords["eggs"].coord_type == "longitude"
+                assert ax.coords["eggs"].coord_wrap == 360 * u.deg
+                assert ax.coords["eggs"].get_format_unit() == u.arcsec
+                assert ax.coords["spam"].coord_type == "latitude"
+                assert ax.coords["spam"].get_format_unit() == u.deg
+
+        # Now test the mappings have been removed
+        fig2 = plt.figure()
+        ax = fig.add_subplot(111, projection=wcs)
+        assert ax.coords["eggs"].coord_type == "scalar"
+        assert ax.coords["eggs"].coord_wrap == None
+        assert ax.coords["eggs"].get_format_unit() == u.deg
+        assert ax.coords["spam"].coord_type == "scalar"
+        assert ax.coords["spam"].coord_wrap == None
+
+
+def test_custom_coord_type_1d_2d_wcs_overwrite():
+    wcs = WCS(naxis=2)
+    wcs.wcs.ctype = ["HGLN-TAN", "HGLT-TAN"]
+    wcs.wcs.crpix = [256.0] * 2
+    wcs.wcs.cdelt = [-0.05] * 2
+    wcs.wcs.crval = [50.0] * 2
+    wcs.wcs.set()
+
+    custom_meta = {
+        "custom:pos.heliographic.stonyhurst.lon": {
+            "format_unit": u.arcsec,
+            # This also tests that we don't overwrite the custom meta with the
+            # stock meta that will set to longitude when the UCD ends in lon
+            "coord_type": "latitude",
+        }
+    }
+
+    with pytest.raises(
+        ValueError, match="pos.heliographic.stonyhurst.lon already exists"
+    ):
+        with custom_ucd_coord_meta_mapping(custom_meta):
+            _, coord_meta = transform_coord_meta_from_wcs(wcs, RectangularFrame)
+
+    with custom_ucd_coord_meta_mapping(custom_meta, overwrite=True):
+        _, coord_meta = transform_coord_meta_from_wcs(wcs, RectangularFrame)
+
+    assert coord_meta["type"] == ["latitude", "latitude"]
+    assert coord_meta["format_unit"] == [u.arcsec, u.deg]
+    assert coord_meta["wrap"] == [None, None]
+
 
 def test_coord_type_1d_1d_wcs():
     wcs = WCS(naxis=1)