Feature: Add AstroPy Support

README.rst

@@ -104,7 +104,20 @@ UltraPlot is published on `PyPi <https://pypi.org/project/ultraplot/>`__ and
    pip install ultraplot
    conda install -c conda-forge ultraplot
 
-The default install includes optional features (for example, pyCirclize-based plots).
+The PyPI install keeps optional integrations out of the default dependency set.
+Install extras only when you need them:
+
+.. code-block:: bash
+
+   pip install "ultraplot[astro]"   # Astropy/WCS integration
+   pip install "ultraplot[all]"     # Optional runtime integrations
+   pip install "ultraplot[astro,geo]"  # Combine extras
+   pip install "ultraplot[astropy,cartopy]"  # Package-name aliases also work
+
+The Conda package includes UltraPlot's supported integrations because Conda does
+not provide an equivalent to pip extras. No additional Astropy installation is
+needed for WCS axes when UltraPlot is installed from conda-forge.
+
 For a minimal install, use ``--no-deps`` and install the core requirements:
 
 .. code-block:: bash

docs/contributing.rst

@@ -88,6 +88,10 @@ To build the documentation locally, use the following commands:
    # Install dependencies to the base conda environment..
    conda env update -f environment.yml
    pip install -e ".[docs]"
+   # Add optional runtime integrations when needed
+   # pip install -e ".[astro]"
+   # pip install -e ".[all]"
+   # pip install -e ".[astro,geo]"
    # ...or create a new conda environment
    # conda env create -n ultraplot-dev --file docs/environment.yml
    # source activate ultraplot-dev

docs/usage.rst

@@ -135,10 +135,10 @@ functionality to existing figure and axes methods.
 Integration features
 ====================
 
-UltraPlot includes *optional* integration features with four external
+UltraPlot includes *optional* integration features with several external
 packages: the `pandas`_ and `xarray`_ packages, used for working with annotated
-tables and arrays, and the `cartopy`_ and `basemap`_ geographic
-plotting packages.
+tables and arrays, the `cartopy`_ and `basemap`_ geographic plotting packages,
+and Astropy WCS axes support for astronomical images.
 
 * The :class:`~ultraplot.axes.GeoAxes` class uses the `cartopy`_ or
   `basemap`_ packages to :ref:`plot geophysical data <ug_geoplot>`,
@@ -147,6 +147,10 @@ plotting packages.
   provides a simpler, cleaner interface than the original `cartopy`_ and `basemap`_
   interfaces. Figures can be filled with :class:`~ultraplot.axes.GeoAxes` by passing the
   `proj` keyword to :func:`~ultraplot.ui.subplots`.
+* UltraPlot can create native Astropy-backed WCS axes for astronomical plots.
+  When Astropy is installed, WCS objects and the ``"wcs"`` / ``"astropy"`` projection
+  aliases resolve to :class:`~ultraplot.axes.AstroAxes`, preserving Astropy
+  transforms like :meth:`~astropy.visualization.wcsaxes.core.WCSAxes.get_transform`.
 * If you pass a :class:`~pandas.Series`, :class:`~pandas.DataFrame`, or :class:`~xarray.DataArray`
   to any plotting command, the axis labels, tick labels, titles, colorbar
   labels, and legend labels are automatically applied from the metadata. If
@@ -158,6 +162,41 @@ plotting packages.
 Since these features are optional,
 UltraPlot can be used without installing any of these packages.
 
+.. _ug_astro_axes:
+
+Astropy WCS axes
+----------------
+
+UltraPlot can integrate with Astropy's WCSAxes for astronomical images and
+world-coordinate overlays. Install Astropy only when you need that support:
+
+.. code-block:: bash
+
+   pip install "ultraplot[astro]"
+
+The conda-forge package includes Astropy and the other supported integrations,
+since Conda has no equivalent to pip extras.
+
+You can then create WCS-aware axes either by passing a WCS object directly or
+by using the ``"wcs"`` projection alias with an explicit ``wcs=...`` argument:
+
+.. code-block:: python
+
+   from astropy.wcs import WCS
+   import ultraplot as uplt
+
+   wcs = WCS(naxis=2)
+   fig, axs = uplt.subplots(ncols=2, proj=[wcs, "wcs"], wcs=wcs)
+   axs[0].imshow(image)
+   axs[1].imshow(image)
+   axs[1].plot(ra, dec, transform=axs[1].get_transform("icrs"))
+
+These axes are represented by :class:`~ultraplot.axes.AstroAxes`, so they
+participate in UltraPlot sharing, panel layout, and figure-level formatting.
+UltraPlot's :meth:`~ultraplot.axes.Axes.format` supports a focused subset of WCS
+label, tick, and grid controls, while Astropy-specific coordinate customization
+is still available directly through ``ax.coords[...]``.
+
 .. _ug_external_axes:
 
 External axes containers (mpltern, others)

environment.yml

@@ -5,6 +5,7 @@ dependencies:
   - python>=3.10,<3.15
   - numpy
   - matplotlib>=3.9
+  - astropy
   - basemap >=1.4.1
   - cartopy
   - xarray

pyproject.toml

@@ -70,7 +70,30 @@ filterwarnings = [
 ]
 mpl-default-style = { axes.prop_cycle = "cycler('color', ['#4c72b0ff', '#55a868ff', '#c44e52ff', '#8172b2ff', '#ccb974ff', '#64b5cdff'])" }
 [project.optional-dependencies]
+astro = [
+    "astropy",
+]
+astropy = [
+    "astropy",
+]
+geo = [
+    "cartopy",
+    "basemap",
+]
+ternary = [
+    "mpltern",
+]
+mpltern = [
+    "mpltern",
+]
+all = [
+    "astropy",
+    "cartopy",
+    "basemap",
+    "mpltern",
+]
 docs = [
+    "astropy",
     "jupyter",
     "jupytext",
     "lxml-html-clean",

ultraplot/_subplots.py

@@ -7,13 +7,11 @@
 
 import matplotlib.axes as maxes
 import matplotlib.gridspec as mgridspec
-import matplotlib.projections as mproj
 import numpy as np
 
-from . import axes as paxes
-from . import constructor
 from . import gridspec as pgridspec
 from .internals import _not_none, _pop_params, warnings
+from .internals.projections import resolve_projection_kwargs
 
 if TYPE_CHECKING:
     from .figure import Figure
@@ -85,68 +83,16 @@ def parse_proj(
         """
         Translate user-input projection into a registered matplotlib axes class.
         """
-        # Parse arguments
         proj = _not_none(proj=proj, projection=projection, default="cartesian")
         proj_kw = _not_none(proj_kw=proj_kw, projection_kw=projection_kw, default={})
         backend = self.parse_backend(backend, basemap)
-        if isinstance(proj, str):
-            proj = proj.lower()
-
-        # Search axes projections
-        name = None
-
-        # Handle cartopy/basemap Projection objects directly
-        # These should be converted to Ultraplot GeoAxes
-        if not isinstance(proj, str):
-            if constructor.Projection is not object and isinstance(
-                proj, constructor.Projection
-            ):
-                name = "ultraplot_cartopy"
-                kwargs["map_projection"] = proj
-            elif constructor.Basemap is not object and isinstance(
-                proj, constructor.Basemap
-            ):
-                name = "ultraplot_basemap"
-                kwargs["map_projection"] = proj
-
-        if name is None and isinstance(proj, str):
-            try:
-                mproj.get_projection_class("ultraplot_" + proj)
-            except (KeyError, ValueError):
-                pass
-            else:
-                name = "ultraplot_" + proj
-        if name is None and isinstance(proj, str):
-            # Try geographic projections first if cartopy/basemap available
-            if (
-                constructor.Projection is not object
-                or constructor.Basemap is not object
-            ):
-                try:
-                    proj_obj = constructor.Proj(
-                        proj, backend=backend, include_axes=True, **proj_kw
-                    )
-                    name = "ultraplot_" + proj_obj._proj_backend
-                    kwargs["map_projection"] = proj_obj
-                except ValueError:
-                    pass  # not a geographic projection, try matplotlib registry below
-
-            # If not geographic, check if registered globally in matplotlib
-            # (e.g., 'ternary', 'polar', '3d')
-            if name is None and proj in mproj.get_projection_names():
-                name = proj
-
-        if name is None and isinstance(proj, str):
-            raise ValueError(
-                f"Invalid projection name {proj!r}. If you are trying to generate a "
-                "GeoAxes with a cartopy.crs.Projection or mpl_toolkits.basemap.Basemap "
-                "then cartopy or basemap must be installed. Otherwise the known axes "
-                f"subclasses are:\n{paxes._cls_table}"
-            )
-
-        if name is not None:
-            kwargs["projection"] = name
-        return kwargs
+        return resolve_projection_kwargs(
+            self.figure,
+            proj,
+            proj_kw=proj_kw,
+            backend=backend,
+            kwargs=kwargs,
+        )
 
     def add_subplot(self, *args, **kwargs):
         """
@@ -229,45 +175,11 @@ def add_subplot(self, *args, **kwargs):
         kwargs.setdefault("number", 1 + max(self.subplot_dict, default=0))
         kwargs.pop("refwidth", None)  # TODO: remove this
 
-        # Use container approach for external projections to make them
-        # ultraplot-compatible. Skip projections that start with "ultraplot_"
-        # as these are already Ultraplot axes classes.
-        projection_name = kwargs.get("projection")
-        external_axes_class = None
-        external_axes_kwargs = {}
-
-        if projection_name and isinstance(projection_name, str):
-            if not projection_name.startswith("ultraplot_"):
-                try:
-                    proj_class = mproj.get_projection_class(projection_name)
-                    if not issubclass(proj_class, paxes.Axes):
-                        external_axes_class = proj_class
-                        external_axes_kwargs["projection"] = projection_name
-
-                        from .axes.container import create_external_axes_container
-
-                        container_name = f"_ultraplot_container_{projection_name}"
-                        if container_name not in mproj.get_projection_names():
-                            container_class = create_external_axes_container(
-                                proj_class, projection_name=container_name
-                            )
-                            mproj.register_projection(container_class)
-
-                        kwargs["projection"] = container_name
-                        kwargs["external_axes_class"] = external_axes_class
-                        kwargs["external_axes_kwargs"] = external_axes_kwargs
-                except (KeyError, ValueError):
-                    pass
-
         kwargs.pop("_subplot_spec", None)
 
-        # NOTE: Skip past Figure.add_subplot (which routes back here) to the
-        # matplotlib implementation. Using super() rather than naming the
-        # matplotlib class keeps any mixin between Figure and matplotlib's
-        # Figure in a subclass MRO from being bypassed.
-        from .figure import Figure
-
-        ax = super(Figure, fig).add_subplot(ss, **kwargs)
+        # Figure.add_subplot routes back here, so use the dedicated hook that
+        # subclasses can override before dispatching to matplotlib.
+        ax = fig._add_subplot_mpl(ss, **kwargs)
         if ax.number:
             self.subplot_dict[ax.number] = ax
         return ax

ultraplot/axes/__init__.py

@@ -20,6 +20,10 @@
 from .taylor import TaylorAxes
 from .three import ThreeAxes  # noqa: F401
 
+_ASTRO_AXES_CLASS = None
+_ASTROPY_WCS_TYPES = ()
+_ASTRO_LOADED = False
+
 # Prevent importing module names and set order of appearance for objects
 __all__ = [
     "Axes",
@@ -36,6 +40,28 @@
 # NOTE: We integrate with cartopy and basemap rather than using matplotlib's
 # native projection system. Therefore axes names are not part of public API.
 _cls_dict = {}  # track valid names
+
+
+def _refresh_cls_table():
+    global _cls_table
+    _cls_table = "\n".join(
+        " "
+        + key
+        + " " * (max(map(len, _cls_dict)) - len(key) + 7)
+        + ("GeoAxes" if cls.__name__[:1] == "_" else cls.__name__)
+        for key, cls in _cls_dict.items()
+    )
+
+
+def _register_projection_class(_cls):
+    for _name in (_cls._name, *_cls._name_aliases):
+        with context._state_context(_cls, name="ultraplot_" + _name):
+            if "ultraplot_" + _name not in mproj.get_projection_names():
+                mproj.register_projection(_cls)
+            _cls_dict[_name] = _cls
+    _refresh_cls_table()
+
+
 for _cls in (
     CartesianAxes,
     PolarAxes,
@@ -44,14 +70,45 @@
     _BasemapAxes,
     ThreeAxes,
 ):
-    for _name in (_cls._name, *_cls._name_aliases):
-        with context._state_context(_cls, name="ultraplot_" + _name):
-            mproj.register_projection(_cls)
-            _cls_dict[_name] = _cls
-_cls_table = "\n".join(
-    " "
-    + key
-    + " " * (max(map(len, _cls_dict)) - len(key) + 7)
-    + ("GeoAxes" if cls.__name__[:1] == "_" else cls.__name__)
-    for key, cls in _cls_dict.items()
-)
+    _register_projection_class(_cls)
+
+
+def _load_astro_axes():
+    global _ASTROPY_WCS_TYPES, _ASTRO_AXES_CLASS, _ASTRO_LOADED
+    if _ASTRO_LOADED:
+        return _ASTRO_AXES_CLASS
+    try:
+        from .astro import ASTROPY_WCS_TYPES as _types, AstroAxes as _astro_axes
+    except ImportError as exc:
+        raise ImportError(
+            "AstroAxes support requires astropy. Install it with "
+            '`pip install "ultraplot[astro]"` or `pip install astropy`.'
+        ) from exc
+
+    _ASTRO_LOADED = True
+    _ASTROPY_WCS_TYPES = _types
+    _ASTRO_AXES_CLASS = _astro_axes
+    if "AstroAxes" not in __all__:
+        __all__.append("AstroAxes")
+    _register_projection_class(_ASTRO_AXES_CLASS)
+    return _ASTRO_AXES_CLASS
+
+
+def get_astro_axes_class(*, load=False):
+    if load:
+        _load_astro_axes()
+    return _ASTRO_AXES_CLASS
+
+
+def get_astropy_wcs_types(*, load=False):
+    if load:
+        _load_astro_axes()
+    return _ASTROPY_WCS_TYPES
+
+
+def __getattr__(name):
+    if name == "AstroAxes":
+        return get_astro_axes_class(load=True)
+    if name == "ASTROPY_WCS_TYPES":
+        return get_astropy_wcs_types(load=True)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")