From e8f2452a47b5f25d2819d022ab36f228f4ff62e6 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Sun, 8 Mar 2026 08:31:18 +1000 Subject: [PATCH 01/12] Handle WCS projection objects in external containers Closes #606 --- ultraplot/_subplots.py | 34 ++---------- ultraplot/axes/container.py | 15 ++++++ ultraplot/figure.py | 49 +++++++++++++++++ .../tests/test_external_container_mocked.py | 53 +++++++++++++++++++ 4 files changed, 122 insertions(+), 29 deletions(-) diff --git a/ultraplot/_subplots.py b/ultraplot/_subplots.py index 4cd34ef03..22ceb7856 100644 --- a/ultraplot/_subplots.py +++ b/ultraplot/_subplots.py @@ -144,6 +144,8 @@ def parse_proj( if name is not None: kwargs["projection"] = name + elif not isinstance(proj, str): + kwargs["projection"] = proj return kwargs def add_subplot(self, *args, **kwargs): @@ -227,35 +229,9 @@ 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 + # Wrap Matplotlib-native or third-party projection classes so UltraPlot + # can preserve its own axes bookkeeping while delegating rendering. + kwargs = fig._wrap_external_projection(**kwargs) kwargs.pop("_subplot_spec", None) diff --git a/ultraplot/axes/container.py b/ultraplot/axes/container.py index 98bbcba88..e91cde762 100644 --- a/ultraplot/axes/container.py +++ b/ultraplot/axes/container.py @@ -700,6 +700,21 @@ def get_external_child(self): """ return self.get_external_axes() + def get_transform(self, *args, **kwargs): + """ + Delegate projection-specific transform lookups to the external axes. + + Some external axes classes (for example WCSAxes) accept extra arguments + like ``frame`` on ``get_transform()``. Without an explicit override here, + the container inherits ``Artist.get_transform()`` and masks that API. + """ + if self._external_axes is not None: + ext_get_transform = getattr(type(self._external_axes), "get_transform", None) + base_get_transform = getattr(maxes.Axes, "get_transform", None) + if args or kwargs or ext_get_transform is not base_get_transform: + return self._external_axes.get_transform(*args, **kwargs) + return super().get_transform() + def clear(self): """Clear the container and mark external axes as stale.""" # Mark external axes as stale before clearing diff --git a/ultraplot/figure.py b/ultraplot/figure.py index 60cf99014..52d54c1db 100644 --- a/ultraplot/figure.py +++ b/ultraplot/figure.py @@ -1733,6 +1733,54 @@ def _parse_proj(self, *args, **kwargs): """Delegate to SubplotManager.""" return self._subplots.parse_proj(*args, **kwargs) + def _wrap_external_projection(self, **kwargs): + """Wrap non-ultraplot projection classes in an external container.""" + projection = kwargs.get("projection") + if projection is None: + return kwargs + + external_axes_class = None + external_axes_kwargs = {} + if isinstance(projection, str): + if projection.startswith("ultraplot_"): + return kwargs + try: + external_axes_class = mproj.get_projection_class(projection) + except (KeyError, ValueError): + return kwargs + elif hasattr(projection, "_as_mpl_axes"): + try: + external_axes_class, external_axes_kwargs = ( + self._process_projection_requirements(projection=projection) + ) + except Exception: + return kwargs + else: + return kwargs + + if issubclass(external_axes_class, paxes.Axes): + return kwargs + + from .axes.container import create_external_axes_container + + container_token = ( + f"{external_axes_class.__module__}_{external_axes_class.__name__}" + ) + container_name = ( + "_ultraplot_container_" + + container_token.replace(".", "_").replace("-", "_").lower() + ) + if container_name not in mproj.get_projection_names(): + container_class = create_external_axes_container( + external_axes_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"] = dict(external_axes_kwargs) + return kwargs + def _get_align_axes(self, side): """ Return the main axes along the edge of the figure. @@ -3251,6 +3299,7 @@ def add_axes(self, rect, **kwargs): %(figure.axes)s """ kwargs = self._parse_proj(**kwargs) + kwargs = self._wrap_external_projection(**kwargs) return super().add_axes(rect, **kwargs) @docstring._concatenate_inherited diff --git a/ultraplot/tests/test_external_container_mocked.py b/ultraplot/tests/test_external_container_mocked.py index bb2c30305..4454c2b01 100644 --- a/ultraplot/tests/test_external_container_mocked.py +++ b/ultraplot/tests/test_external_container_mocked.py @@ -233,6 +233,29 @@ def get_tightbbox(self, renderer): return super().get_tightbbox(renderer) +class MockProjectionTransformAxes(MockExternalAxes): + """Mock external axes with a projection-aware get_transform API.""" + + def __init__(self, fig, *args, transform_id=None, **kwargs): + self.transform_id = transform_id + self.transform_calls = [] + super().__init__(fig, *args, **kwargs) + + def get_transform(self, frame=None): + self.transform_calls.append(frame) + return (self.transform_id, frame) + + +class MockProjectionObject: + """Projection-like object resolved by Matplotlib via _as_mpl_axes.""" + + def __init__(self, transform_id="mock"): + self.transform_id = transform_id + + def _as_mpl_axes(self): + return MockProjectionTransformAxes, {"transform_id": self.transform_id} + + # Tests @@ -261,6 +284,36 @@ def test_container_creation_with_external_axes(): assert isinstance(ax.get_external_child(), MockExternalAxes) +def test_add_axes_wraps_projection_object_and_delegates_get_transform(): + """Projection objects should be wrapped and keep custom transform APIs.""" + fig = uplt.figure() + ax = fig.add_axes( + [0.1, 0.1, 0.8, 0.8], projection=MockProjectionObject("mock-wcs") + ) + + assert ax.has_external_child() + child = ax.get_external_child() + assert isinstance(child, MockProjectionTransformAxes) + + transform = ax.get_transform("icrs") + assert transform == ("mock-wcs", "icrs") + assert child.transform_calls == ["icrs"] + + +def test_add_subplot_wraps_projection_object_and_delegates_get_transform(): + """Subplots should also wrap projection objects via the external container.""" + fig = uplt.figure() + ax = fig.add_subplot(111, projection=MockProjectionObject("subplot-wcs")) + + assert ax.has_external_child() + child = ax.get_external_child() + assert isinstance(child, MockProjectionTransformAxes) + + transform = ax.get_transform("fk5") + assert transform == ("subplot-wcs", "fk5") + assert child.transform_calls == ["fk5"] + + def test_external_axes_removed_from_figure_axes(): """Test that external axes is removed from figure axes list.""" fig = uplt.figure() From 2af15e608fe76cf57be169bc9706783c85216c8a Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Sun, 8 Mar 2026 08:39:12 +1000 Subject: [PATCH 02/12] Black --- ultraplot/axes/container.py | 4 +++- ultraplot/tests/test_external_container_mocked.py | 4 +--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/ultraplot/axes/container.py b/ultraplot/axes/container.py index e91cde762..5a922c1ed 100644 --- a/ultraplot/axes/container.py +++ b/ultraplot/axes/container.py @@ -709,7 +709,9 @@ def get_transform(self, *args, **kwargs): the container inherits ``Artist.get_transform()`` and masks that API. """ if self._external_axes is not None: - ext_get_transform = getattr(type(self._external_axes), "get_transform", None) + ext_get_transform = getattr( + type(self._external_axes), "get_transform", None + ) base_get_transform = getattr(maxes.Axes, "get_transform", None) if args or kwargs or ext_get_transform is not base_get_transform: return self._external_axes.get_transform(*args, **kwargs) diff --git a/ultraplot/tests/test_external_container_mocked.py b/ultraplot/tests/test_external_container_mocked.py index 4454c2b01..5e3bef78c 100644 --- a/ultraplot/tests/test_external_container_mocked.py +++ b/ultraplot/tests/test_external_container_mocked.py @@ -287,9 +287,7 @@ def test_container_creation_with_external_axes(): def test_add_axes_wraps_projection_object_and_delegates_get_transform(): """Projection objects should be wrapped and keep custom transform APIs.""" fig = uplt.figure() - ax = fig.add_axes( - [0.1, 0.1, 0.8, 0.8], projection=MockProjectionObject("mock-wcs") - ) + ax = fig.add_axes([0.1, 0.1, 0.8, 0.8], projection=MockProjectionObject("mock-wcs")) assert ax.has_external_child() child = ax.get_external_child() From 3aeb6a8e16dd96a15f168896d601910805a822da Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Tue, 10 Mar 2026 18:51:06 +1000 Subject: [PATCH 03/12] Add native AstroAxes WCS integration Route Astropy WCS projections through a native AstroAxes bridge instead of the external container, preserve per-axes transform dispatch for SubplotGrid, and keep shared panel labels stable for both cartesian and astro axes. --- ultraplot/_subplots.py | 69 +----- ultraplot/axes/__init__.py | 10 +- ultraplot/axes/astro.py | 372 +++++++++++++++++++++++++++++ ultraplot/axes/base.py | 2 + ultraplot/figure.py | 185 ++++++++------ ultraplot/gridspec.py | 24 +- ultraplot/internals/projections.py | 285 ++++++++++++++++++++++ ultraplot/tests/test_astro_axes.py | 192 +++++++++++++++ ultraplot/tests/test_subplots.py | 77 ++++++ 9 files changed, 1071 insertions(+), 145 deletions(-) create mode 100644 ultraplot/axes/astro.py create mode 100644 ultraplot/internals/projections.py create mode 100644 ultraplot/tests/test_astro_axes.py diff --git a/ultraplot/_subplots.py b/ultraplot/_subplots.py index 22ceb7856..5911b70a1 100644 --- a/ultraplot/_subplots.py +++ b/ultraplot/_subplots.py @@ -14,6 +14,7 @@ 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 @@ -86,67 +87,13 @@ def parse_proj( 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 - constructor._warn_basemap_deprecated() - - 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 - elif not isinstance(proj, str): - kwargs["projection"] = proj - return kwargs + return resolve_projection_kwargs( + self.figure, + proj, + proj_kw=proj_kw, + backend=backend, + kwargs=kwargs, + ) def add_subplot(self, *args, **kwargs): """ diff --git a/ultraplot/axes/__init__.py b/ultraplot/axes/__init__.py index 37effe7d8..321998a77 100644 --- a/ultraplot/axes/__init__.py +++ b/ultraplot/axes/__init__.py @@ -14,6 +14,7 @@ _BasemapAxes, _CartopyAxes, ) +from .astro import ASTROPY_WCS_TYPES, AstroAxes from .plot import PlotAxes # noqa: F401 from .polar import PolarAxes from .shared import _SharedAxes # noqa: F401 @@ -31,19 +32,24 @@ "ThreeAxes", "ExternalAxesContainer", ] +if AstroAxes is not None: + __all__.append("AstroAxes") # Register projections with package prefix to avoid conflicts # 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 -for _cls in ( +_projection_classes = [ CartesianAxes, PolarAxes, TaylorAxes, _CartopyAxes, _BasemapAxes, ThreeAxes, -): +] +if AstroAxes is not None: + _projection_classes.append(AstroAxes) +for _cls in _projection_classes: for _name in (_cls._name, *_cls._name_aliases): with context._state_context(_cls, name="ultraplot_" + _name): mproj.register_projection(_cls) diff --git a/ultraplot/axes/astro.py b/ultraplot/axes/astro.py new file mode 100644 index 000000000..8169b9bdd --- /dev/null +++ b/ultraplot/axes/astro.py @@ -0,0 +1,372 @@ +#!/usr/bin/env python3 +""" +Astropy WCS axes integration. +""" + +import inspect +import numbers +from collections.abc import Iterable + +from ..config import rc +from ..utils import _not_none +from . import base +from .cartesian import CartesianAxes + +try: + from astropy.visualization.wcsaxes.core import WCSAxes + from astropy.wcs.wcsapi import BaseHighLevelWCS, BaseLowLevelWCS +except ImportError: # pragma: no cover + WCSAxes = None + ASTROPY_WCS_TYPES = () +else: + ASTROPY_WCS_TYPES = (BaseLowLevelWCS, BaseHighLevelWCS) + + +if WCSAxes is not None: + + class AstroAxes(base.Axes, WCSAxes): + """ + Native UltraPlot wrapper for Astropy WCS axes. + """ + + _name = "astro" + _name_aliases = ("astropy", "wcs") + + def _update_background(self, **kwargs): + kw_face, kw_edge = rc._get_background_props(**kwargs) + self.patch.update(kw_face) + self.patch.update(kw_edge) + + def _get_coord_helper(self, axis): + index = {"x": 0, "y": 1}[axis] + try: + return self.coords[index] + except IndexError: + return None + + def _share_coord_signature(self, axis): + coord = self._get_coord_helper(axis) + if coord is None: + return None + unit = getattr(coord, "coord_unit", None) + if unit is not None and hasattr(unit, "to_string"): + unit = unit.to_string() + return ( + getattr(coord, "coord_type", None), + unit, + getattr(coord, "default_label", None), + ) + + def _update_coord_locator(self, axis, locator): + coord = self._get_coord_helper(axis) + if coord is None or locator is None: + return + if isinstance(locator, numbers.Real) and not isinstance(locator, bool): + coord.set_ticks(number=locator) + return + if isinstance(locator, Iterable) and not isinstance(locator, (str, bytes)): + coord.set_ticks(values=locator) + return + raise TypeError( + "AstroAxes.format only supports numeric or iterable tick locators. " + f"Received {locator!r}. Use ax.coords[...] for advanced locator setup." + ) + + def _update_coord_formatter(self, axis, formatter): + coord = self._get_coord_helper(axis) + if coord is None or formatter is None: + return + if isinstance(formatter, str) or callable(formatter): + coord.set_major_formatter(formatter) + return + raise TypeError( + "AstroAxes.format only supports string or callable tick formatters. " + f"Received {formatter!r}. Use ax.coords[...] for advanced formatter setup." + ) + + def _update_coord_ticks( + self, + axis, + *, + grid=None, + gridcolor=None, + tickcolor=None, + ticklen=None, + tickwidth=None, + tickdir=None, + ticklabelpad=None, + ticklabelcolor=None, + ticklabelsize=None, + ticklabelweight=None, + tickminor=None, + ): + coord = self._get_coord_helper(axis) + if coord is None: + return + if tickminor is not None: + coord.display_minor_ticks(bool(tickminor)) + major = {} + if ticklen is not None: + major["length"] = ticklen + if tickwidth is not None: + major["width"] = tickwidth + if tickcolor is not None: + major["color"] = tickcolor + if tickdir is not None: + major["direction"] = tickdir + if ticklabelpad is not None: + major["pad"] = ticklabelpad + if ticklabelcolor is not None: + major["labelcolor"] = ticklabelcolor + if ticklabelsize is not None: + major["labelsize"] = ticklabelsize + if major: + coord.tick_params(**major) + if ticklabelweight is not None: + coord.set_ticklabel(weight=ticklabelweight) + if grid is not None or gridcolor is not None: + kw = {} + if gridcolor is not None: + kw["color"] = gridcolor + coord.grid(draw_grid=grid, **kw) + + def _update_axis_label( + self, + axis, + *, + label=None, + labelpad=None, + labelcolor=None, + labelsize=None, + labelweight=None, + label_kw=None, + ): + coord = self._get_coord_helper(axis) + if coord is None: + return + if label is None and not any( + value is not None for value in (labelpad, labelcolor, labelsize, labelweight) + ): + return + setter = getattr(self, f"set_{axis}label") + getter = getattr(self, f"get_{axis}label") + kw = dict(label_kw or {}) + if labelcolor is not None: + kw["color"] = labelcolor + if labelsize is not None: + kw["size"] = labelsize + if labelweight is not None: + kw["weight"] = labelweight + if labelpad is not None: + kw["labelpad"] = labelpad + setter(getter() if label is None else label, **kw) + + def _update_limits(self, axis, *, lim=None, min_=None, max_=None, reverse=None): + lo = hi = None + if lim is not None: + lo, hi = lim + lo = _not_none(min_=min_, lim_0=lo) + hi = _not_none(max_=max_, lim_1=hi) + if lo is not None or hi is not None: + get_lim = getattr(self, f"get_{axis}lim") + set_lim = getattr(self, f"set_{axis}lim") + cur_lo, cur_hi = get_lim() + set_lim((_not_none(lo, cur_lo), _not_none(hi, cur_hi))) + if reverse is not None: + inverted = getattr(self, f"{axis}axis_inverted")() + if bool(reverse) != bool(inverted): + getattr(self, f"invert_{axis}axis")() + + def _share_axis_limits(self, other, which): + self._shared_axes[which].join(self, other) + axis = getattr(self, f"{which}axis") + other_axis = getattr(other, f"{which}axis") + setattr(self, f"_share{which}", other) + axis.major = other_axis.major + axis.minor = other_axis.minor + get_lim = getattr(other, f"get_{which}lim") + set_lim = getattr(self, f"set_{which}lim") + get_auto = getattr(other, f"get_autoscale{which}_on") + set_lim(*get_lim(), emit=False, auto=get_auto()) + axis._scale = other_axis._scale + + def _sharex_setup(self, sharex, *, labels=True, limits=True): + super()._sharex_setup(sharex) + level = ( + 3 + if self._panel_sharex_group and self._is_panel_group_member(sharex) + else self.figure._sharex + ) + if level not in range(5): + raise ValueError(f"Invalid sharing level sharex={level!r}.") + if sharex in (None, self) or not isinstance(sharex, AstroAxes): + return + if level > 0 and labels: + self._sharex = sharex + if level > 1 and limits: + self._share_axis_limits(sharex, "x") + + def _sharey_setup(self, sharey, *, labels=True, limits=True): + super()._sharey_setup(sharey) + level = ( + 3 + if self._panel_sharey_group and self._is_panel_group_member(sharey) + else self.figure._sharey + ) + if level not in range(5): + raise ValueError(f"Invalid sharing level sharey={level!r}.") + if sharey in (None, self) or not isinstance(sharey, AstroAxes): + return + if level > 0 and labels: + self._sharey = sharey + if level > 1 and limits: + self._share_axis_limits(sharey, "y") + + def _is_ticklabel_on(self, side: str) -> bool: + axis = "x" if side in ("labelbottom", "labeltop") else "y" + coord = self._get_coord_helper(axis) + if coord is None or not coord.get_ticklabel_visible(): + return False + positions = coord.get_ticklabel_position() + tokens = { + "labelbottom": "b", + "labeltop": "t", + "labelleft": "l", + "labelright": "r", + "bottom": "b", + "top": "t", + "left": "l", + "right": "r", + } + token = tokens.get(side, side) + if token in positions: + return True + if "#" in positions: + return token == ("b" if axis == "x" else "l") + return False + + def _apply_ticklabel_state(self, axis: str, state: dict): + coord = self._get_coord_helper(axis) + if coord is None: + return + positions = [] + for side in ("bottom", "top") if axis == "x" else ("left", "right"): + if state.get(f"label{side}", False): + positions.append(side[0]) + position = "".join(positions) + coord.set_ticklabel_position(position) + coord.set_axislabel_position(position) + coord.set_ticklabel_visible(bool(positions)) + + def format( + self, + *, + aspect=None, + xreverse=None, + yreverse=None, + xlim=None, + ylim=None, + xmin=None, + ymin=None, + xmax=None, + ymax=None, + xformatter=None, + yformatter=None, + xlocator=None, + ylocator=None, + xtickminor=None, + ytickminor=None, + xtickcolor=None, + ytickcolor=None, + xticklen=None, + yticklen=None, + xtickwidth=None, + ytickwidth=None, + xtickdir=None, + ytickdir=None, + xticklabelpad=None, + yticklabelpad=None, + xticklabelcolor=None, + yticklabelcolor=None, + xticklabelsize=None, + yticklabelsize=None, + xticklabelweight=None, + yticklabelweight=None, + xlabel=None, + ylabel=None, + xlabelpad=None, + ylabelpad=None, + xlabelcolor=None, + ylabelcolor=None, + xlabelsize=None, + ylabelsize=None, + xlabelweight=None, + ylabelweight=None, + xgrid=None, + ygrid=None, + xgridcolor=None, + ygridcolor=None, + xlabel_kw=None, + ylabel_kw=None, + **kwargs, + ): + if aspect is not None: + self.set_aspect(aspect) + self._update_limits("x", lim=xlim, min_=xmin, max_=xmax, reverse=xreverse) + self._update_limits("y", lim=ylim, min_=ymin, max_=ymax, reverse=yreverse) + self._update_coord_locator("x", xlocator) + self._update_coord_locator("y", ylocator) + self._update_coord_formatter("x", xformatter) + self._update_coord_formatter("y", yformatter) + self._update_coord_ticks( + "x", + grid=xgrid, + gridcolor=xgridcolor, + tickcolor=xtickcolor, + ticklen=xticklen, + tickwidth=xtickwidth, + tickdir=xtickdir, + ticklabelpad=xticklabelpad, + ticklabelcolor=xticklabelcolor, + ticklabelsize=xticklabelsize, + ticklabelweight=xticklabelweight, + tickminor=xtickminor, + ) + self._update_coord_ticks( + "y", + grid=ygrid, + gridcolor=ygridcolor, + tickcolor=ytickcolor, + ticklen=yticklen, + tickwidth=ytickwidth, + tickdir=ytickdir, + ticklabelpad=yticklabelpad, + ticklabelcolor=yticklabelcolor, + ticklabelsize=yticklabelsize, + ticklabelweight=yticklabelweight, + tickminor=ytickminor, + ) + self._update_axis_label( + "x", + label=xlabel, + labelpad=xlabelpad, + labelcolor=xlabelcolor, + labelsize=xlabelsize, + labelweight=xlabelweight, + label_kw=xlabel_kw, + ) + self._update_axis_label( + "y", + label=ylabel, + labelpad=ylabelpad, + labelcolor=ylabelcolor, + labelsize=ylabelsize, + labelweight=ylabelweight, + label_kw=ylabel_kw, + ) + return base.Axes.format(self, **kwargs) + + + AstroAxes._format_signatures[AstroAxes] = inspect.signature(CartesianAxes.format) +else: # pragma: no cover + AstroAxes = None diff --git a/ultraplot/axes/base.py b/ultraplot/axes/base.py index 19a429d15..d82b9e1cf 100644 --- a/ultraplot/axes/base.py +++ b/ultraplot/axes/base.py @@ -1089,6 +1089,8 @@ def _add_inset_axes( if proj is None: if self._name in ("cartopy", "basemap"): proj = copy.copy(self.projection) + elif self._name == "astro" and getattr(self, "wcs", None) is not None: + proj = self.wcs else: proj = self._name kwargs = self.figure._parse_proj(proj, **kwargs) diff --git a/ultraplot/figure.py b/ultraplot/figure.py index 52d54c1db..95b4e22d0 100644 --- a/ultraplot/figure.py +++ b/ultraplot/figure.py @@ -29,6 +29,7 @@ from . import gridspec as pgridspec from . import legend as plegend from .config import rc, rc_matplotlib +from .internals.projections import finalize_projection_kwargs from .internals import ( _alias_kwargs, _not_none, @@ -1175,6 +1176,14 @@ def _share_axes_compatible(self, ref, other, which: str): ): return False, "different Geo projection classes" + ref_astro = paxes.AstroAxes is not None and isinstance(ref, paxes.AstroAxes) + other_astro = paxes.AstroAxes is not None and isinstance(other, paxes.AstroAxes) + if ref_astro or other_astro: + if not (ref_astro and other_astro): + return False, "astro and non-astro axes cannot be shared" + if ref._share_coord_signature(which) != other._share_coord_signature(which): + return False, "different Astro coordinate families" + # Polar and non-polar should not share. ref_polar = isinstance(ref, paxes.PolarAxes) other_polar = isinstance(other, paxes.PolarAxes) @@ -1182,7 +1191,12 @@ def _share_axes_compatible(self, ref, other, which: str): return False, "polar and non-polar axes cannot be shared" # Non-geo external axes are generally Cartesian-like in UltraPlot. - if not ref_geo and not other_geo and not (ref_external or other_external): + if ( + not ref_geo + and not other_geo + and not (ref_external or other_external) + and not (ref_astro or other_astro) + ): if not ( isinstance(ref, paxes.CartesianAxes) and isinstance(other, paxes.CartesianAxes) @@ -1597,9 +1611,10 @@ def _compute_baseline_tick_state(self, group_axes, axis: str, label_keys): return {}, True # Supported axes types - if not isinstance( - axi, (paxes.CartesianAxes, paxes._CartopyAxes, paxes._BasemapAxes) - ): + supported = (paxes.CartesianAxes, paxes._CartopyAxes, paxes._BasemapAxes) + if paxes.AstroAxes is not None: + supported = (*supported, paxes.AstroAxes) + if not isinstance(axi, supported): warnings._warn_ultraplot( f"Tick label sharing not implemented for {type(axi)} subplots." ) @@ -1615,6 +1630,11 @@ def _compute_baseline_tick_state(self, group_axes, axis: str, label_keys): key = label_keys[f"label{side}"] if params.get(key): baseline[key] = params[key] + elif paxes.AstroAxes is not None and isinstance(axi, paxes.AstroAxes): + for side in sides: + key = f"label{side}" + if axi._is_ticklabel_on(key): + baseline[key] = True elif isinstance(axi, paxes.GeoAxes): for side in sides: key = f"label{side}" @@ -1697,6 +1717,19 @@ def _effective_share_level(self, axi, axis: str, sides: tuple[str, str]) -> int: return level + def _get_ticklabel_state(self, axi, axis: str): + """Read the visible ticklabel sides for cartesian, geo, and astro axes.""" + sides = ("top", "bottom") if axis == "x" else ("left", "right") + if isinstance(axi, paxes.GeoAxes): + return {f"label{side}": axi._is_ticklabel_on(f"label{side}") for side in sides} + if paxes.AstroAxes is not None and isinstance(axi, paxes.AstroAxes): + return {f"label{side}": axi._is_ticklabel_on(f"label{side}") for side in sides} + params = getattr(axi, f"{axis}axis").get_tick_params() + return { + f"label{side}": params.get(axi._label_key(f"label{side}"), False) + for side in sides + } + def _set_ticklabel_state(self, axi, axis: str, state: dict): """Apply the computed ticklabel state to cartesian or geo axes.""" if state: @@ -1704,6 +1737,8 @@ def _set_ticklabel_state(self, axi, axis: str, state: dict): cleaned = {k: (True if v in ("x", "y") else v) for k, v in state.items()} if isinstance(axi, paxes.GeoAxes): axi._toggle_gridliner_labels(**cleaned) + elif paxes.AstroAxes is not None and isinstance(axi, paxes.AstroAxes): + axi._apply_ticklabel_state(axis, cleaned) else: getattr(axi, f"{axis}axis").set_tick_params(**cleaned) @@ -1735,51 +1770,7 @@ def _parse_proj(self, *args, **kwargs): def _wrap_external_projection(self, **kwargs): """Wrap non-ultraplot projection classes in an external container.""" - projection = kwargs.get("projection") - if projection is None: - return kwargs - - external_axes_class = None - external_axes_kwargs = {} - if isinstance(projection, str): - if projection.startswith("ultraplot_"): - return kwargs - try: - external_axes_class = mproj.get_projection_class(projection) - except (KeyError, ValueError): - return kwargs - elif hasattr(projection, "_as_mpl_axes"): - try: - external_axes_class, external_axes_kwargs = ( - self._process_projection_requirements(projection=projection) - ) - except Exception: - return kwargs - else: - return kwargs - - if issubclass(external_axes_class, paxes.Axes): - return kwargs - - from .axes.container import create_external_axes_container - - container_token = ( - f"{external_axes_class.__module__}_{external_axes_class.__name__}" - ) - container_name = ( - "_ultraplot_container_" - + container_token.replace(".", "_").replace("-", "_").lower() - ) - if container_name not in mproj.get_projection_names(): - container_class = create_external_axes_container( - external_axes_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"] = dict(external_axes_kwargs) - return kwargs + return finalize_projection_kwargs(self, kwargs) def _get_align_axes(self, side): """ @@ -2129,31 +2120,24 @@ def _add_axes_panel( *getattr(ax, f"get_{'y' if side in ('left','right') else 'x'}lim")(), auto=True, ) + filled = kw.get("filled", False) + shared_state = None # Push main axes tick labels to the outside relative to the added panel # Skip this for filled panels (colorbars/legends) - if not kw.get("filled", False) and share: - if isinstance(ax, paxes.GeoAxes): - if side == "top": - ax._toggle_gridliner_labels(labeltop=False) - elif side == "bottom": - ax._toggle_gridliner_labels(labelbottom=False) - elif side == "left": - ax._toggle_gridliner_labels(labelleft=False) - elif side == "right": - ax._toggle_gridliner_labels(labelright=False) - else: - if side == "top": - ax.xaxis.set_tick_params(**{ax._label_key("labeltop"): False}) - elif side == "bottom": - ax.xaxis.set_tick_params(**{ax._label_key("labelbottom"): False}) - elif side == "left": - ax.yaxis.set_tick_params(**{ax._label_key("labelleft"): False}) - elif side == "right": - ax.yaxis.set_tick_params(**{ax._label_key("labelright"): False}) - - # Panel labels: prefer outside only for non-sharing top/right; otherwise keep off + if not filled and share: + shared_axis = "y" if side in ("left", "right") else "x" + shared_state = self._get_ticklabel_state(ax, shared_axis) + main_state = shared_state.copy() + main_state[f"label{side}"] = False + self._set_ticklabel_state(ax, shared_axis, main_state) + + # Panel labels: for non-sharing panels, keep labels on the outer edges of the + # full stack. For shared panels, only propagate the panel-side labels where + # the existing sharing logic expects them (top/right). if side == "top": - if not share: + if not share and not filled: + ax.xaxis.tick_bottom() + ax.xaxis.set_label_position("bottom") pax.xaxis.set_tick_params( **{ pax._label_key("labeltop"): True, @@ -2161,11 +2145,17 @@ def _add_axes_panel( } ) else: - on = ax.xaxis.get_tick_params()[ax._label_key("labeltop")] - pax.xaxis.set_tick_params(**{pax._label_key("labeltop"): on}) - ax.yaxis.set_tick_params(labeltop=False) + on = shared_state is not None and shared_state.get("labeltop", False) + pax.xaxis.set_tick_params( + **{ + pax._label_key("labeltop"): on, + pax._label_key("labelbottom"): False, + } + ) elif side == "right": - if not share: + if not share and not filled: + ax.yaxis.tick_left() + ax.yaxis.set_label_position("left") pax.yaxis.set_tick_params( **{ pax._label_key("labelright"): True, @@ -2173,9 +2163,47 @@ def _add_axes_panel( } ) else: - on = ax.yaxis.get_tick_params()[ax._label_key("labelright")] - pax.yaxis.set_tick_params(**{pax._label_key("labelright"): on}) - ax.yaxis.set_tick_params(**{ax._label_key("labelright"): False}) + on = shared_state is not None and shared_state.get("labelright", False) + pax.yaxis.set_tick_params( + **{ + pax._label_key("labelright"): on, + pax._label_key("labelleft"): False, + } + ) + elif side == "left" and not share and not filled: + ax.yaxis.tick_right() + ax.yaxis.set_label_position("right") + ax.yaxis.set_tick_params( + **{ + ax._label_key("labelleft"): False, + ax._label_key("labelright"): True, + } + ) + pax.yaxis.tick_left() + pax.yaxis.set_label_position("left") + pax.yaxis.set_tick_params( + **{ + pax._label_key("labelleft"): True, + pax._label_key("labelright"): False, + } + ) + elif side == "bottom" and not share and not filled: + ax.xaxis.tick_top() + ax.xaxis.set_label_position("top") + ax.xaxis.set_tick_params( + **{ + ax._label_key("labelbottom"): False, + ax._label_key("labeltop"): True, + } + ) + pax.xaxis.tick_bottom() + pax.xaxis.set_label_position("bottom") + pax.xaxis.set_tick_params( + **{ + pax._label_key("labelbottom"): True, + pax._label_key("labeltop"): False, + } + ) return pax @@ -3299,7 +3327,6 @@ def add_axes(self, rect, **kwargs): %(figure.axes)s """ kwargs = self._parse_proj(**kwargs) - kwargs = self._wrap_external_projection(**kwargs) return super().add_axes(rect, **kwargs) @docstring._concatenate_inherited diff --git a/ultraplot/gridspec.py b/ultraplot/gridspec.py index eb7fdd59f..f1c42b82a 100644 --- a/ultraplot/gridspec.py +++ b/ultraplot/gridspec.py @@ -39,6 +39,15 @@ __all__ = ["GridSpec", "SubplotGrid"] +class _GridCommandResult(tuple): + """ + Tuple subclass marking one result per axes from `SubplotGrid` dispatch. + """ + + def __new__(cls, values): + return super().__new__(cls, values) + + # Gridspec vector arguments # Valid for figure() and GridSpec() _shared_docstring = """ @@ -1892,11 +1901,20 @@ def __getattr__(self, attr): return objs[0] if len(self) == 1 else objs elif all(map(callable, objs)): + def _dispatch_value(obj, idx): + if isinstance(obj, _GridCommandResult): + return obj[idx] + return obj + @functools.wraps(objs[0]) def _iterate_subplots(*args, **kwargs): result = [] - for func in objs: - result.append(func(*args, **kwargs)) + for idx, func in enumerate(objs): + iargs = tuple(_dispatch_value(arg, idx) for arg in args) + ikwargs = { + key: _dispatch_value(val, idx) for key, val in kwargs.items() + } + result.append(func(*iargs, **ikwargs)) if len(self) == 1: return result[0] elif all(res is None for res in result): @@ -1904,7 +1922,7 @@ def _iterate_subplots(*args, **kwargs): elif all(isinstance(res, paxes.Axes) for res in result): return SubplotGrid(result, n=self._n, order=self._order) else: - return tuple(result) + return _GridCommandResult(result) _iterate_subplots.__doc__ = inspect.getdoc(objs[0]) return _iterate_subplots diff --git a/ultraplot/internals/projections.py b/ultraplot/internals/projections.py new file mode 100644 index 000000000..8b3116f95 --- /dev/null +++ b/ultraplot/internals/projections.py @@ -0,0 +1,285 @@ +#!/usr/bin/env python3 +""" +Projection binding registry used by figure axes creation. +""" + +from dataclasses import dataclass, field + +import matplotlib.projections as mproj + +from .. import constructor + + +@dataclass(frozen=True) +class ProjectionContext: + """ + Context passed to projection bindings. + """ + + figure: object + proj_kw: dict + backend: str | None + + +@dataclass(frozen=True) +class ProjectionResolution: + """ + Resolved projection plus any injected keyword arguments. + """ + + projection: object | str | None = None + kwargs: dict = field(default_factory=dict) + + def as_kwargs(self, kwargs=None): + merged = dict(kwargs or {}) + if self.projection is not None: + merged["projection"] = self.projection + merged.update(self.kwargs) + return merged + + +@dataclass(frozen=True) +class ProjectionBinding: + """ + Projection matcher and resolver pair. + """ + + name: str + matcher: object + resolver: object + + +_PROJECTION_BINDINGS = [] + + +def register_projection_binding(name, matcher, resolver=None): + """ + Register a projection binding. Can be used as a decorator. + """ + if resolver is None: + + def decorator(func): + _PROJECTION_BINDINGS.append(ProjectionBinding(name, matcher, func)) + return func + + return decorator + + _PROJECTION_BINDINGS.append(ProjectionBinding(name, matcher, resolver)) + return resolver + + +def iter_projection_bindings(): + """ + Return the registered projection bindings. + """ + return tuple(_PROJECTION_BINDINGS) + + +def _get_axes_module(): + from .. import axes as paxes + + return paxes + + +def _prefixed_projection_name(name): + if name.startswith("ultraplot_"): + return name if name in mproj.get_projection_names() else None + prefixed = "ultraplot_" + name + try: + mproj.get_projection_class(prefixed) + except (KeyError, ValueError): + return None + return prefixed + + +def _container_projection_name(external_axes_class): + token = f"{external_axes_class.__module__}_{external_axes_class.__name__}" + return "_ultraplot_container_" + token.replace(".", "_").replace("-", "_").lower() + + +def _wrap_external_projection(figure, projection): + if projection is None: + return ProjectionResolution() + + external_axes_class = None + external_axes_kwargs = {} + if isinstance(projection, str): + if projection.startswith("ultraplot_") or projection.startswith( + "_ultraplot_container_" + ): + return ProjectionResolution(projection=projection) + try: + external_axes_class = mproj.get_projection_class(projection) + except (KeyError, ValueError): + return ProjectionResolution(projection=projection) + elif hasattr(projection, "_as_mpl_axes"): + try: + external_axes_class, external_axes_kwargs = ( + figure._process_projection_requirements(projection=projection) + ) + except Exception: + return ProjectionResolution(projection=projection) + else: + return ProjectionResolution(projection=projection) + + paxes = _get_axes_module() + if issubclass(external_axes_class, paxes.Axes): + return ProjectionResolution( + projection=projection, + kwargs=dict(external_axes_kwargs), + ) + + from ..axes.container import create_external_axes_container + + container_name = _container_projection_name(external_axes_class) + if container_name not in mproj.get_projection_names(): + container_class = create_external_axes_container( + external_axes_class, projection_name=container_name + ) + mproj.register_projection(container_class) + + return ProjectionResolution( + projection=container_name, + kwargs={ + "external_axes_class": external_axes_class, + "external_axes_kwargs": dict(external_axes_kwargs), + }, + ) + + +@register_projection_binding( + "native_ultraplot_string", + lambda proj, context: isinstance(proj, str) + and _prefixed_projection_name(proj) is not None, +) +def _resolve_native_ultraplot_string(proj, context): + return ProjectionResolution(projection=_prefixed_projection_name(proj)) + + +@register_projection_binding( + "astropy_wcs_object", + lambda proj, context: ( + not isinstance(proj, str) + and bool(_get_axes_module().ASTROPY_WCS_TYPES) + and isinstance(proj, _get_axes_module().ASTROPY_WCS_TYPES) + ), +) +def _resolve_astropy_wcs_object(proj, context): + return ProjectionResolution(projection="ultraplot_astro", kwargs={"wcs": proj}) + + +@register_projection_binding( + "cartopy_projection_object", + lambda proj, context: ( + not isinstance(proj, str) + and constructor.Projection is not object + and isinstance(proj, constructor.Projection) + ), +) +def _resolve_cartopy_projection_object(proj, context): + return ProjectionResolution( + projection="ultraplot_cartopy", + kwargs={"map_projection": proj}, + ) + + +@register_projection_binding( + "basemap_projection_object", + lambda proj, context: ( + not isinstance(proj, str) + and constructor.Basemap is not object + and isinstance(proj, constructor.Basemap) + ), +) +def _resolve_basemap_projection_object(proj, context): + return ProjectionResolution( + projection="ultraplot_basemap", + kwargs={"map_projection": proj}, + ) + + +@register_projection_binding( + "geographic_projection_name", + lambda proj, context: isinstance(proj, str) + and (constructor.Projection is not object or constructor.Basemap is not object), +) +def _resolve_geographic_projection_name(proj, context): + try: + proj_obj = constructor.Proj( + proj, + backend=context.backend, + include_axes=True, + **context.proj_kw, + ) + except ValueError: + return ProjectionResolution() + return ProjectionResolution( + projection="ultraplot_" + proj_obj._proj_backend, + kwargs={"map_projection": proj_obj}, + ) + + +@register_projection_binding( + "registered_matplotlib_string", + lambda proj, context: isinstance(proj, str) and proj in mproj.get_projection_names(), +) +def _resolve_registered_matplotlib_string(proj, context): + return ProjectionResolution(projection=proj) + + +def resolve_projection(proj, *, figure, proj_kw=None, backend=None): + """ + Resolve a user projection spec to a final projection and kwargs. + """ + proj_kw = proj_kw or {} + if isinstance(proj, str): + proj = proj.lower() + context = ProjectionContext(figure=figure, proj_kw=proj_kw, backend=backend) + + resolution = None + for binding in _PROJECTION_BINDINGS: + if binding.matcher(proj, context): + resolution = binding.resolver(proj, context) + if resolution.projection is not None or resolution.kwargs: + break + + if resolution is None or (resolution.projection is None and not resolution.kwargs): + if isinstance(proj, str): + paxes = _get_axes_module() + 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}" + ) + resolution = ProjectionResolution(projection=proj) + + final = _wrap_external_projection(figure, resolution.projection) + merged_kwargs = dict(resolution.kwargs) + merged_kwargs.update(final.kwargs) + projection = final.projection if final.projection is not None else resolution.projection + return ProjectionResolution(projection=projection, kwargs=merged_kwargs) + + +def resolve_projection_kwargs(figure, proj, *, proj_kw=None, backend=None, kwargs=None): + """ + Resolve a projection and merge the result into an existing keyword dictionary. + """ + resolution = resolve_projection( + proj, + figure=figure, + proj_kw=proj_kw, + backend=backend, + ) + return resolution.as_kwargs(kwargs) + + +def finalize_projection_kwargs(figure, kwargs): + """ + Finalize an already-parsed projection dictionary. + """ + projection = kwargs.get("projection") + if projection is None: + return kwargs + final = _wrap_external_projection(figure, projection) + return final.as_kwargs(kwargs) diff --git a/ultraplot/tests/test_astro_axes.py b/ultraplot/tests/test_astro_axes.py new file mode 100644 index 000000000..8403263ff --- /dev/null +++ b/ultraplot/tests/test_astro_axes.py @@ -0,0 +1,192 @@ +import warnings + +import numpy as np +import pytest + +import ultraplot as uplt +from ultraplot import axes as paxes + +pytest.importorskip("astropy.visualization.wcsaxes") +from astropy.wcs import WCS + + +def _make_test_wcs(): + wcs = WCS(naxis=2) + wcs.wcs.crpix = [50.0, 50.0] + wcs.wcs.cdelt = [-0.066667, 0.066667] + wcs.wcs.crval = [0.0, -90.0] + wcs.wcs.ctype = ["RA---TAN", "DEC--TAN"] + wcs.wcs.cunit = ["deg", "deg"] + return wcs + + +def test_add_subplot_with_wcs_projection_returns_native_astro_axes(): + fig = uplt.figure() + ax = fig.add_subplot(111, projection=_make_test_wcs()) + + assert paxes.AstroAxes is not None + assert isinstance(ax, paxes.AstroAxes) + assert not (hasattr(ax, "has_external_axes") and ax.has_external_axes()) + assert ax.get_transform("icrs") is not None + + fig.canvas.draw() + bbox = ax.get_tightbbox(fig.canvas.get_renderer()) + assert bbox.width > 0 + assert bbox.height > 0 + + +def test_add_axes_with_wcs_projection_supports_basic_formatting(): + fig = uplt.figure() + ax = fig.add_axes([0.1, 0.1, 0.8, 0.8], projection=_make_test_wcs()) + + ax.format(xlabel="RA", ylabel="Dec", title="Sky", xgrid=True, ygrid=True) + + assert isinstance(ax, paxes.AstroAxes) + assert ax.get_xlabel() == "RA" + assert ax.get_ylabel() == "Dec" + assert ax.get_title() == "Sky" + + fig.canvas.draw() + bbox = ax.get_tightbbox(fig.canvas.get_renderer()) + assert bbox.width > 0 + assert bbox.height > 0 + + +def test_string_wcs_projection_uses_native_astro_axes(): + fig = uplt.figure() + ax = fig.add_subplot(111, projection="wcs", wcs=_make_test_wcs()) + + assert isinstance(ax, paxes.AstroAxes) + + +def test_same_family_astro_axes_can_share_without_warning(): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + fig, (ax1, ax2) = uplt.subplots( + nrows=2, + proj=[_make_test_wcs(), _make_test_wcs()], + sharex=2, + ) + + messages = [str(item.message) for item in caught] + assert not any("Skipping incompatible x-axis sharing" in msg for msg in messages) + assert ax1.get_shared_x_axes().joined(ax1, ax2) + + +def test_different_astro_coordinate_families_do_not_share(): + galactic = _make_test_wcs() + galactic.wcs.ctype = ["GLON-TAN", "GLAT-TAN"] + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + fig, (ax1, ax2) = uplt.subplots( + nrows=2, + proj=[_make_test_wcs(), galactic], + sharex=2, + ) + + messages = [str(item.message) for item in caught] + assert any("different Astro coordinate families" in msg for msg in messages) + assert not ax1.get_shared_x_axes().joined(ax1, ax2) + + +def test_subplot_grid_arrow_dispatches_per_axes_transforms(): + fig, axs = uplt.subplots( + ncols=2, + proj=[_make_test_wcs(), _make_test_wcs()], + share=0, + ) + axs.imshow(np.zeros((16, 16)), origin="lower") + + arrows = axs.arrow( + 0.0, + -89.95, + 0.0, + 0.02, + head_width=0, + head_length=0, + width=0.01, + transform=axs.get_transform("icrs"), + ) + + assert len(arrows) == 2 + fig.canvas.draw() + + +def test_astro_axes_share_ticklabels_without_hiding_outer_wcs_labels(): + fig, axs = uplt.subplots( + ncols=2, + proj=[_make_test_wcs(), _make_test_wcs()], + ) + axs.imshow(np.zeros((16, 16)), origin="lower") + + fig.canvas.draw() + + assert axs[0].coords[1].get_ticklabel_visible() + assert axs[0].coords[1].get_axislabel_position() + assert not axs[1].coords[1].get_ticklabel_visible() + assert not axs[1].coords[1].get_axislabel_position() + + +def test_astro_axes_preserve_shared_top_labels(): + fig, axs = uplt.subplots( + nrows=2, + proj=[_make_test_wcs(), _make_test_wcs()], + ) + axs.imshow(np.zeros((16, 16)), origin="lower") + for ax in axs: + ax.coords[0].set_ticklabel_position("t") + ax.coords[0].set_axislabel_position("t") + + fig.canvas.draw() + + assert axs[0].coords[0].get_ticklabel_position() == ["t"] + assert axs[0].coords[0].get_axislabel_position() == ["t"] + assert not axs[1].coords[0].get_ticklabel_position() + assert not axs[1].coords[0].get_axislabel_position() + + +def test_astro_axes_preserve_shared_right_labels(): + fig, axs = uplt.subplots( + ncols=2, + proj=[_make_test_wcs(), _make_test_wcs()], + ) + axs.imshow(np.zeros((16, 16)), origin="lower") + for ax in axs: + ax.coords[1].set_ticklabel_position("r") + ax.coords[1].set_axislabel_position("r") + + fig.canvas.draw() + + assert not axs[0].coords[1].get_ticklabel_position() + assert not axs[0].coords[1].get_axislabel_position() + assert axs[1].coords[1].get_ticklabel_position() == ["r"] + assert axs[1].coords[1].get_axislabel_position() == ["r"] + + +def test_astro_axes_panels_preserve_explicit_top_right_labels(): + fig, axs = uplt.subplots( + nrows=2, + ncols=2, + proj=[_make_test_wcs() for _ in range(4)], + ) + axs.imshow(np.zeros((16, 16)), origin="lower") + for ax in axs: + ax.coords[0].set_ticklabel_position("t") + ax.coords[0].set_axislabel_position("t") + ax.coords[1].set_ticklabel_position("r") + ax.coords[1].set_axislabel_position("r") + + pax_top = axs[0].panel("top") + pax_right = axs[1].panel("right") + fig.canvas.draw() + + assert not axs[0].coords[0].get_ticklabel_position() + assert not axs[0].coords[0].get_axislabel_position() + assert pax_top._is_ticklabel_on("labeltop") + assert not pax_top._is_ticklabel_on("labelbottom") + + assert not axs[1].coords[1].get_ticklabel_position() + assert not axs[1].coords[1].get_axislabel_position() + assert pax_right._is_ticklabel_on("labelright") + assert not pax_right._is_ticklabel_on("labelleft") diff --git a/ultraplot/tests/test_subplots.py b/ultraplot/tests/test_subplots.py index 458e9b902..08acfecf4 100644 --- a/ultraplot/tests/test_subplots.py +++ b/ultraplot/tests/test_subplots.py @@ -722,6 +722,30 @@ def assert_panel(axi_panel, side, share_flag): assert_panel(pax_bottom, "bottom", share_panels) +def test_shared_panels_preserve_explicit_top_right_labels(): + fig, axs = uplt.subplots(nrows=2, ncols=2) + for ax in axs: + ax.imshow(np.zeros((10, 10))) + ax.xaxis.tick_top() + ax.xaxis.set_label_position("top") + ax.xaxis.set_tick_params(labeltop=True, labelbottom=False) + ax.yaxis.tick_right() + ax.yaxis.set_label_position("right") + ax.yaxis.set_tick_params(labelright=True, labelleft=False) + + pax_top = axs[0].panel("top") + pax_right = axs[1].panel("right") + fig.canvas.draw() + + assert not axs[0]._is_ticklabel_on("labeltop") + assert pax_top._is_ticklabel_on("labeltop") + assert not pax_top._is_ticklabel_on("labelbottom") + + assert not axs[1]._is_ticklabel_on("labelright") + assert pax_right._is_ticklabel_on("labelright") + assert not pax_right._is_ticklabel_on("labelleft") + + def test_non_rectangular_outside_labels_top(): """ Check that non-rectangular layouts work with outside labels. @@ -806,6 +830,59 @@ def test_panel_share_flag_controls_group_membership(): assert ax2[0]._panel_sharex_group is False +def test_nonsharing_left_panel_moves_main_labels_outside(): + fig, axs = uplt.subplots() + ax = axs[0] + ax.format(ylabel="main ylabel") + pax = ax.panel("left", share=False) + pax.format(ylabel="panel ylabel") + + fig.canvas.draw() + + assert not ax._is_ticklabel_on("labelleft") + assert ax._is_ticklabel_on("labelright") + assert pax._is_ticklabel_on("labelleft") + assert not pax._is_ticklabel_on("labelright") + assert ax.yaxis.get_label_position() == "right" + assert pax.yaxis.get_label_position() == "left" + + +def test_nonsharing_bottom_panel_moves_main_labels_outside(): + fig, axs = uplt.subplots() + ax = axs[0] + ax.format(xlabel="main xlabel") + pax = ax.panel("bottom", share=False) + pax.format(xlabel="panel xlabel") + + fig.canvas.draw() + + assert not ax._is_ticklabel_on("labelbottom") + assert ax._is_ticklabel_on("labeltop") + assert pax._is_ticklabel_on("labelbottom") + assert not pax._is_ticklabel_on("labeltop") + assert ax.xaxis.get_label_position() == "top" + assert pax.xaxis.get_label_position() == "bottom" + + +def test_nonsharing_left_panel_gap_matches_right_panel(): + def _panel_gap(side): + fig, axs = uplt.subplots() + ax = axs[0] + ax.format(ylabel="main ylabel") + pax = ax.panel(side, share=False) + pax.format(xlabel="panel xlabel", ylabel="panel ylabel") + fig.canvas.draw() + main = ax.get_position().bounds + panel = pax.get_position().bounds + if side == "left": + return main[0] - (panel[0] + panel[2]) + return panel[0] - (main[0] + main[2]) + + gap_left = _panel_gap("left") + gap_right = _panel_gap("right") + assert abs(gap_left - gap_right) < 1e-3 + + def test_ticklabels_with_guides_share_true_cartesian(): """ With share=True, tick labels should only appear on bottom row and left column From 22fad05c9cefad8b2295c38699d52ef87f5e26d7 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Tue, 10 Mar 2026 18:58:58 +1000 Subject: [PATCH 04/12] Refactor ticklabel sharing adapters Move logical ticklabel state access into the axes classes so Figure can share and hand off panel labels without branching on cartesian, geo, and astro subplot types. --- ultraplot/axes/astro.py | 12 ++++++- ultraplot/axes/base.py | 23 ++++++++++++ ultraplot/axes/geo.py | 13 +++++++ ultraplot/figure.py | 80 +++++------------------------------------ 4 files changed, 55 insertions(+), 73 deletions(-) diff --git a/ultraplot/axes/astro.py b/ultraplot/axes/astro.py index 8169b9bdd..45edf69eb 100644 --- a/ultraplot/axes/astro.py +++ b/ultraplot/axes/astro.py @@ -245,7 +245,14 @@ def _is_ticklabel_on(self, side: str) -> bool: return token == ("b" if axis == "x" else "l") return False - def _apply_ticklabel_state(self, axis: str, state: dict): + def _get_ticklabel_state(self, axis: str) -> dict[str, bool]: + sides = ("top", "bottom") if axis == "x" else ("left", "right") + return { + f"label{side}": self._is_ticklabel_on(f"label{side}") + for side in sides + } + + def _set_ticklabel_state(self, axis: str, state: dict): coord = self._get_coord_helper(axis) if coord is None: return @@ -258,6 +265,9 @@ def _apply_ticklabel_state(self, axis: str, state: dict): coord.set_axislabel_position(position) coord.set_ticklabel_visible(bool(positions)) + def _apply_ticklabel_state(self, axis: str, state: dict): + self._set_ticklabel_state(axis, state) + def format( self, *, diff --git a/ultraplot/axes/base.py b/ultraplot/axes/base.py index d82b9e1cf..64a4b9750 100644 --- a/ultraplot/axes/base.py +++ b/ultraplot/axes/base.py @@ -3595,6 +3595,29 @@ def _is_ticklabel_on(self, side: str) -> bool: return axis.get_tick_params().get(self._label_key(side), False) + def _get_ticklabel_state(self, axis: str) -> dict[str, bool]: + """ + Return visible ticklabel sides for one logical axis. + """ + sides = ("top", "bottom") if axis == "x" else ("left", "right") + return { + f"label{side}": self._is_ticklabel_on(f"label{side}") + for side in sides + } + + def _set_ticklabel_state(self, axis: str, state: dict) -> None: + """ + Apply logical ticklabel visibility to one logical axis. + """ + cleaned = {k: (True if v in ("x", "y") else v) for k, v in state.items()} + mapped = { + self._label_key(key): value + for key, value in cleaned.items() + if key.startswith("label") + } + if mapped: + getattr(self, f"{axis}axis").set_tick_params(**mapped) + @docstring._snippet_manager def inset(self, *args, **kwargs): """ diff --git a/ultraplot/axes/geo.py b/ultraplot/axes/geo.py index 786fbff54..07238833f 100644 --- a/ultraplot/axes/geo.py +++ b/ultraplot/axes/geo.py @@ -2367,6 +2367,19 @@ def _is_ticklabel_on(self, side: str) -> bool: return False return adapter.is_label_on(side) + def _get_ticklabel_state(self, axis: str) -> dict[str, bool]: + sides = ("top", "bottom") if axis == "x" else ("left", "right") + return { + f"label{side}": self._is_ticklabel_on(f"label{side}") + for side in sides + } + + def _set_ticklabel_state(self, axis: str, state: dict) -> None: + sides = ("top", "bottom") if axis == "x" else ("left", "right") + self._toggle_gridliner_labels( + **{f"label{side}": state.get(f"label{side}", False) for side in sides} + ) + def _clear_edge_lon_labels(self) -> None: for label in self._edge_lon_labels: try: diff --git a/ultraplot/figure.py b/ultraplot/figure.py index 95b4e22d0..7a70821e7 100644 --- a/ultraplot/figure.py +++ b/ultraplot/figure.py @@ -1498,9 +1498,6 @@ def _share_ticklabels(self, *, axis: str) -> None: axes = list(self._iter_axes(panels=True, hidden=False)) groups = self._group_axes_by_axis(axes, axis) - # Version-dependent label name mapping for reading back params - label_keys = self._label_key_map() - # Process each group independently for _, group_axes in groups.items(): # Singleton groups can still need border masking reapplied for @@ -1523,9 +1520,7 @@ def _share_ticklabels(self, *, axis: str) -> None: continue # Build baseline from MAIN axes only (exclude panels) - baseline, skip_group = self._compute_baseline_tick_state( - group_axes, axis, label_keys - ) + baseline, skip_group = self._compute_baseline_tick_state(group_axes, axis) if skip_group: continue @@ -1539,27 +1534,10 @@ def _share_ticklabels(self, *, axis: str) -> None: continue # Apply to geo/cartesian appropriately - self._set_ticklabel_state(axi, axis, masked) + axi._set_ticklabel_state(axis, masked) self.stale = True - def _label_key_map(self): - """ - Return a mapping for version-dependent label keys for Matplotlib tick params. - """ - first_axi = next(self._iter_axes(panels=True), None) - if first_axi is None: - return { - "labelleft": "labelleft", - "labelright": "labelright", - "labeltop": "labeltop", - "labelbottom": "labelbottom", - } - return { - name: first_axi._label_key(name) - for name in ("labelleft", "labelright", "labeltop", "labelbottom") - } - def _group_axes_by_axis(self, axes, axis: str): """ Group axes by row (x) or column (y). Panels included; invalid subplotspec skipped. @@ -1580,7 +1558,7 @@ def _group_key(ax): groups[key].append(axi) return groups - def _compute_baseline_tick_state(self, group_axes, axis: str, label_keys): + def _compute_baseline_tick_state(self, group_axes, axis: str): """ Build a baseline ticklabel visibility dict from MAIN axes (panels excluded). Returns (baseline_dict, skip_group: bool). Emits warnings when encountering @@ -1624,22 +1602,9 @@ def _compute_baseline_tick_state(self, group_axes, axis: str, label_keys): subplot_types.add(type(axi)) # Collect label visibility state - if isinstance(axi, paxes.CartesianAxes): - params = getattr(axi, f"{axis}axis").get_tick_params() - for side in sides: - key = label_keys[f"label{side}"] - if params.get(key): - baseline[key] = params[key] - elif paxes.AstroAxes is not None and isinstance(axi, paxes.AstroAxes): - for side in sides: - key = f"label{side}" - if axi._is_ticklabel_on(key): - baseline[key] = True - elif isinstance(axi, paxes.GeoAxes): - for side in sides: - key = f"label{side}" - if axi._is_ticklabel_on(key): - baseline[key] = axi._is_ticklabel_on(key) + for key, value in axi._get_ticklabel_state(axis).items(): + if value: + baseline[key] = value if unsupported_found: return {}, True @@ -1658,16 +1623,12 @@ def _apply_border_mask( ): """ Apply figure-border constraints and panel opposite-side suppression. - Keeps label key mapping per-axis for cartesian. """ from .axes.cartesian import OPPOSITE_SIDE masked = baseline.copy() for side in sides: label = f"label{side}" - if isinstance(axi, paxes.CartesianAxes): - # Use per-axis version-mapped key when writing - label = axi._label_key(label) # Only keep labels on true figure borders if axi not in outer_axes[side]: @@ -1717,31 +1678,6 @@ def _effective_share_level(self, axi, axis: str, sides: tuple[str, str]) -> int: return level - def _get_ticklabel_state(self, axi, axis: str): - """Read the visible ticklabel sides for cartesian, geo, and astro axes.""" - sides = ("top", "bottom") if axis == "x" else ("left", "right") - if isinstance(axi, paxes.GeoAxes): - return {f"label{side}": axi._is_ticklabel_on(f"label{side}") for side in sides} - if paxes.AstroAxes is not None and isinstance(axi, paxes.AstroAxes): - return {f"label{side}": axi._is_ticklabel_on(f"label{side}") for side in sides} - params = getattr(axi, f"{axis}axis").get_tick_params() - return { - f"label{side}": params.get(axi._label_key(f"label{side}"), False) - for side in sides - } - - def _set_ticklabel_state(self, axi, axis: str, state: dict): - """Apply the computed ticklabel state to cartesian or geo axes.""" - if state: - # Normalize "x"/"y" values to booleans for both Geo and Cartesian axes - cleaned = {k: (True if v in ("x", "y") else v) for k, v in state.items()} - if isinstance(axi, paxes.GeoAxes): - axi._toggle_gridliner_labels(**cleaned) - elif paxes.AstroAxes is not None and isinstance(axi, paxes.AstroAxes): - axi._apply_ticklabel_state(axis, cleaned) - else: - getattr(axi, f"{axis}axis").set_tick_params(**cleaned) - def _context_adjusting(self, cache=True): """ Prevent re-running auto layout steps due to draws triggered by figure @@ -2126,10 +2062,10 @@ def _add_axes_panel( # Skip this for filled panels (colorbars/legends) if not filled and share: shared_axis = "y" if side in ("left", "right") else "x" - shared_state = self._get_ticklabel_state(ax, shared_axis) + shared_state = ax._get_ticklabel_state(shared_axis) main_state = shared_state.copy() main_state[f"label{side}"] = False - self._set_ticklabel_state(ax, shared_axis, main_state) + ax._set_ticklabel_state(shared_axis, main_state) # Panel labels: for non-sharing panels, keep labels on the outer edges of the # full stack. For shared panels, only propagate the panel-side labels where From dc4eb2229f6ff7dec84677aecb302febd3b0a379 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Tue, 10 Mar 2026 19:08:50 +1000 Subject: [PATCH 05/12] Lazy-load AstroAxes integration Make the Astropy WCS bridge a true optional import by loading and registering AstroAxes only on WCS-specific access, while keeping the panel/share refactors and import behavior covered by tests. --- ultraplot/axes/__init__.py | 84 +++++++++++++++++++++++------- ultraplot/axes/astro.py | 7 ++- ultraplot/axes/base.py | 5 +- ultraplot/axes/geo.py | 5 +- ultraplot/figure.py | 10 ++-- ultraplot/internals/projections.py | 28 ++++++++-- ultraplot/tests/test_imports.py | 36 +++++++++++++ 7 files changed, 136 insertions(+), 39 deletions(-) diff --git a/ultraplot/axes/__init__.py b/ultraplot/axes/__init__.py index 321998a77..dfa1e252f 100644 --- a/ultraplot/axes/__init__.py +++ b/ultraplot/axes/__init__.py @@ -14,13 +14,16 @@ _BasemapAxes, _CartopyAxes, ) -from .astro import ASTROPY_WCS_TYPES, AstroAxes from .plot import PlotAxes # noqa: F401 from .polar import PolarAxes from .shared import _SharedAxes # noqa: F401 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", @@ -32,32 +35,75 @@ "ThreeAxes", "ExternalAxesContainer", ] -if AstroAxes is not None: - __all__.append("AstroAxes") # Register projections with package prefix to avoid conflicts # 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 -_projection_classes = [ + + +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, TaylorAxes, _CartopyAxes, _BasemapAxes, ThreeAxes, -] -if AstroAxes is not None: - _projection_classes.append(AstroAxes) -for _cls in _projection_classes: - 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 + from .astro import ASTROPY_WCS_TYPES as _types, AstroAxes as _astro_axes + + _ASTRO_LOADED = True + _ASTROPY_WCS_TYPES = _types + _ASTRO_AXES_CLASS = _astro_axes + if _ASTRO_AXES_CLASS is not None: + 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}") diff --git a/ultraplot/axes/astro.py b/ultraplot/axes/astro.py index 45edf69eb..8f124d102 100644 --- a/ultraplot/axes/astro.py +++ b/ultraplot/axes/astro.py @@ -145,7 +145,8 @@ def _update_axis_label( if coord is None: return if label is None and not any( - value is not None for value in (labelpad, labelcolor, labelsize, labelweight) + value is not None + for value in (labelpad, labelcolor, labelsize, labelweight) ): return setter = getattr(self, f"set_{axis}label") @@ -248,8 +249,7 @@ def _is_ticklabel_on(self, side: str) -> bool: def _get_ticklabel_state(self, axis: str) -> dict[str, bool]: sides = ("top", "bottom") if axis == "x" else ("left", "right") return { - f"label{side}": self._is_ticklabel_on(f"label{side}") - for side in sides + f"label{side}": self._is_ticklabel_on(f"label{side}") for side in sides } def _set_ticklabel_state(self, axis: str, state: dict): @@ -376,7 +376,6 @@ def format( ) return base.Axes.format(self, **kwargs) - AstroAxes._format_signatures[AstroAxes] = inspect.signature(CartesianAxes.format) else: # pragma: no cover AstroAxes = None diff --git a/ultraplot/axes/base.py b/ultraplot/axes/base.py index 64a4b9750..a6daf882c 100644 --- a/ultraplot/axes/base.py +++ b/ultraplot/axes/base.py @@ -3600,10 +3600,7 @@ def _get_ticklabel_state(self, axis: str) -> dict[str, bool]: Return visible ticklabel sides for one logical axis. """ sides = ("top", "bottom") if axis == "x" else ("left", "right") - return { - f"label{side}": self._is_ticklabel_on(f"label{side}") - for side in sides - } + return {f"label{side}": self._is_ticklabel_on(f"label{side}") for side in sides} def _set_ticklabel_state(self, axis: str, state: dict) -> None: """ diff --git a/ultraplot/axes/geo.py b/ultraplot/axes/geo.py index 07238833f..61741b937 100644 --- a/ultraplot/axes/geo.py +++ b/ultraplot/axes/geo.py @@ -2369,10 +2369,7 @@ def _is_ticklabel_on(self, side: str) -> bool: def _get_ticklabel_state(self, axis: str) -> dict[str, bool]: sides = ("top", "bottom") if axis == "x" else ("left", "right") - return { - f"label{side}": self._is_ticklabel_on(f"label{side}") - for side in sides - } + return {f"label{side}": self._is_ticklabel_on(f"label{side}") for side in sides} def _set_ticklabel_state(self, axis: str, state: dict) -> None: sides = ("top", "bottom") if axis == "x" else ("left", "right") diff --git a/ultraplot/figure.py b/ultraplot/figure.py index 7a70821e7..7ecbf30e6 100644 --- a/ultraplot/figure.py +++ b/ultraplot/figure.py @@ -1176,8 +1176,9 @@ def _share_axes_compatible(self, ref, other, which: str): ): return False, "different Geo projection classes" - ref_astro = paxes.AstroAxes is not None and isinstance(ref, paxes.AstroAxes) - other_astro = paxes.AstroAxes is not None and isinstance(other, paxes.AstroAxes) + astro_cls = paxes.get_astro_axes_class() + ref_astro = astro_cls is not None and isinstance(ref, astro_cls) + other_astro = astro_cls is not None and isinstance(other, astro_cls) if ref_astro or other_astro: if not (ref_astro and other_astro): return False, "astro and non-astro axes cannot be shared" @@ -1590,8 +1591,9 @@ def _compute_baseline_tick_state(self, group_axes, axis: str): # Supported axes types supported = (paxes.CartesianAxes, paxes._CartopyAxes, paxes._BasemapAxes) - if paxes.AstroAxes is not None: - supported = (*supported, paxes.AstroAxes) + astro_cls = paxes.get_astro_axes_class() + if astro_cls is not None: + supported = (*supported, astro_cls) if not isinstance(axi, supported): warnings._warn_ultraplot( f"Tick label sharing not implemented for {type(axi)} subplots." diff --git a/ultraplot/internals/projections.py b/ultraplot/internals/projections.py index 8b3116f95..61eb45f01 100644 --- a/ultraplot/internals/projections.py +++ b/ultraplot/internals/projections.py @@ -81,6 +81,11 @@ def _get_axes_module(): return paxes +def _looks_like_astropy_projection(proj): + module = getattr(type(proj), "__module__", "") + return module.startswith("astropy.") + + def _prefixed_projection_name(name): if name.startswith("ultraplot_"): return name if name in mproj.get_projection_names() else None @@ -147,6 +152,17 @@ def _wrap_external_projection(figure, projection): ) +@register_projection_binding( + "astropy_wcs_string", + lambda proj, context: isinstance(proj, str) + and proj in ("astro", "astropy", "wcs", "ultraplot_astro"), +) +def _resolve_astropy_wcs_string(proj, context): + if _get_axes_module().get_astro_axes_class(load=True) is None: + return ProjectionResolution() + return ProjectionResolution(projection="ultraplot_astro") + + @register_projection_binding( "native_ultraplot_string", lambda proj, context: isinstance(proj, str) @@ -160,8 +176,9 @@ def _resolve_native_ultraplot_string(proj, context): "astropy_wcs_object", lambda proj, context: ( not isinstance(proj, str) - and bool(_get_axes_module().ASTROPY_WCS_TYPES) - and isinstance(proj, _get_axes_module().ASTROPY_WCS_TYPES) + and _looks_like_astropy_projection(proj) + and bool(_get_axes_module().get_astropy_wcs_types(load=True)) + and isinstance(proj, _get_axes_module().get_astropy_wcs_types()) ), ) def _resolve_astropy_wcs_object(proj, context): @@ -221,7 +238,8 @@ def _resolve_geographic_projection_name(proj, context): @register_projection_binding( "registered_matplotlib_string", - lambda proj, context: isinstance(proj, str) and proj in mproj.get_projection_names(), + lambda proj, context: isinstance(proj, str) + and proj in mproj.get_projection_names(), ) def _resolve_registered_matplotlib_string(proj, context): return ProjectionResolution(projection=proj) @@ -257,7 +275,9 @@ def resolve_projection(proj, *, figure, proj_kw=None, backend=None): final = _wrap_external_projection(figure, resolution.projection) merged_kwargs = dict(resolution.kwargs) merged_kwargs.update(final.kwargs) - projection = final.projection if final.projection is not None else resolution.projection + projection = ( + final.projection if final.projection is not None else resolution.projection + ) return ProjectionResolution(projection=projection, kwargs=merged_kwargs) diff --git a/ultraplot/tests/test_imports.py b/ultraplot/tests/test_imports.py index ff253a8d7..168b2e627 100644 --- a/ultraplot/tests/test_imports.py +++ b/ultraplot/tests/test_imports.py @@ -36,6 +36,42 @@ def test_import_is_lightweight(): assert out == "[]" +def test_loading_axes_does_not_import_astropy(): + code = """ +import json +import sys +import ultraplot as uplt +uplt.subplots() +mods = [name for name in sys.modules if name == "astropy" or name.startswith("astropy.")] +print(json.dumps(sorted(mods))) +""" + out = _run(code) + assert out == "[]" + + +def test_axes_astro_attr_is_lazy_optional(): + code = """ +import importlib.util +import json +import sys +import ultraplot.axes as paxes +spec = importlib.util.find_spec("astropy.visualization.wcsaxes") +astro = paxes.AstroAxes +mods = [name for name in sys.modules if name == "astropy" or name.startswith("astropy.")] +print(json.dumps({ + "available": bool(spec), + "astro_is_none": astro is None, + "loaded": bool(mods), +})) +""" + out = json.loads(_run(code)) + if out["available"]: + assert not out["astro_is_none"] + assert out["loaded"] + else: + assert out["astro_is_none"] + + def test_star_import_exposes_public_api(): code = """ from ultraplot import * # noqa: F403 From b3edc4f678a5ff94792a496bc486e19b75d27713 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Thu, 12 Mar 2026 09:28:15 +1000 Subject: [PATCH 06/12] Refine optional Astro integration packaging Make Astropy a true opt-in runtime extra by adding dedicated package extras for Astro, geographic, and ternary integrations, along with an all extra that installs the optional runtime stack without changing the default UltraPlot dependency set. At the same time, tighten the Astro axes import path so explicit WCS usage fails with a clear installation message instead of silently returning None. astro.py is now treated as an Astropy-backed module again, the lazy loader in axes.__init__ raises a targeted ImportError with the recommended install command, and WCS string resolution in projections.py now relies on that explicit failure path. This also expands inline documentation in AstroAxes to make the WCS override points clearer, and updates the import tests to match the new optional-dependency behavior. --- README.rst | 12 +- docs/contributing.rst | 4 + pyproject.toml | 23 + ultraplot/axes/__init__.py | 15 +- ultraplot/axes/astro.py | 834 ++++++++++++++++------------- ultraplot/internals/projections.py | 3 +- ultraplot/tests/test_imports.py | 15 +- 7 files changed, 532 insertions(+), 374 deletions(-) diff --git a/README.rst b/README.rst index de4ee8185..70148ea59 100644 --- a/README.rst +++ b/README.rst @@ -104,7 +104,17 @@ UltraPlot is published on `PyPi `__ and pip install ultraplot conda install -c conda-forge ultraplot -The default install includes optional features (for example, pyCirclize-based plots). +The default install keeps a broad core runtime set (for example, pyCirclize-based +plots) without pulling in every optional integration. 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 + For a minimal install, use ``--no-deps`` and install the core requirements: .. code-block:: bash diff --git a/docs/contributing.rst b/docs/contributing.rst index 6ccc4cc9c..7a6550b57 100644 --- a/docs/contributing.rst +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index d03a0c617..50c5d2d6e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,6 +71,29 @@ filterwarnings = [ ] mpl-default-style = { axes.prop_cycle = "cycler('color', ['#4c72b0ff', '#55a868ff', '#c44e52ff', '#8172b2ff', '#ccb974ff', '#64b5cdff'])" } [project.optional-dependencies] +astro = [ + "astropy", +] +astropy = [ + "astropy", +] +geo = [ + "cartopy", +] +cartopy = [ + "cartopy", +] +ternary = [ + "mpltern", +] +mpltern = [ + "mpltern", +] +all = [ + "astropy", + "cartopy", + "mpltern", +] docs = [ "jupyter", "jupytext", diff --git a/ultraplot/axes/__init__.py b/ultraplot/axes/__init__.py index dfa1e252f..67bfc0ad7 100644 --- a/ultraplot/axes/__init__.py +++ b/ultraplot/axes/__init__.py @@ -77,15 +77,20 @@ def _load_astro_axes(): global _ASTROPY_WCS_TYPES, _ASTRO_AXES_CLASS, _ASTRO_LOADED if _ASTRO_LOADED: return _ASTRO_AXES_CLASS - from .astro import ASTROPY_WCS_TYPES as _types, AstroAxes as _astro_axes + 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 _ASTRO_AXES_CLASS is not None: - if "AstroAxes" not in __all__: - __all__.append("AstroAxes") - _register_projection_class(_ASTRO_AXES_CLASS) + if "AstroAxes" not in __all__: + __all__.append("AstroAxes") + _register_projection_class(_ASTRO_AXES_CLASS) return _ASTRO_AXES_CLASS diff --git a/ultraplot/axes/astro.py b/ultraplot/axes/astro.py index 8f124d102..3118cf78c 100644 --- a/ultraplot/axes/astro.py +++ b/ultraplot/axes/astro.py @@ -12,370 +12,478 @@ from . import base from .cartesian import CartesianAxes +from astropy.visualization.wcsaxes.core import WCSAxes +from astropy.wcs.wcsapi import BaseHighLevelWCS, BaseLowLevelWCS + try: - from astropy.visualization.wcsaxes.core import WCSAxes - from astropy.wcs.wcsapi import BaseHighLevelWCS, BaseLowLevelWCS -except ImportError: # pragma: no cover - WCSAxes = None - ASTROPY_WCS_TYPES = () -else: - ASTROPY_WCS_TYPES = (BaseLowLevelWCS, BaseHighLevelWCS) - - -if WCSAxes is not None: - - class AstroAxes(base.Axes, WCSAxes): - """ - Native UltraPlot wrapper for Astropy WCS axes. - """ - - _name = "astro" - _name_aliases = ("astropy", "wcs") - - def _update_background(self, **kwargs): - kw_face, kw_edge = rc._get_background_props(**kwargs) - self.patch.update(kw_face) - self.patch.update(kw_edge) - - def _get_coord_helper(self, axis): - index = {"x": 0, "y": 1}[axis] - try: - return self.coords[index] - except IndexError: - return None - - def _share_coord_signature(self, axis): - coord = self._get_coord_helper(axis) - if coord is None: - return None - unit = getattr(coord, "coord_unit", None) - if unit is not None and hasattr(unit, "to_string"): - unit = unit.to_string() - return ( - getattr(coord, "coord_type", None), - unit, - getattr(coord, "default_label", None), - ) - - def _update_coord_locator(self, axis, locator): - coord = self._get_coord_helper(axis) - if coord is None or locator is None: - return - if isinstance(locator, numbers.Real) and not isinstance(locator, bool): - coord.set_ticks(number=locator) - return - if isinstance(locator, Iterable) and not isinstance(locator, (str, bytes)): - coord.set_ticks(values=locator) - return - raise TypeError( - "AstroAxes.format only supports numeric or iterable tick locators. " - f"Received {locator!r}. Use ax.coords[...] for advanced locator setup." - ) - - def _update_coord_formatter(self, axis, formatter): - coord = self._get_coord_helper(axis) - if coord is None or formatter is None: - return - if isinstance(formatter, str) or callable(formatter): - coord.set_major_formatter(formatter) - return - raise TypeError( - "AstroAxes.format only supports string or callable tick formatters. " - f"Received {formatter!r}. Use ax.coords[...] for advanced formatter setup." - ) - - def _update_coord_ticks( - self, - axis, - *, - grid=None, - gridcolor=None, - tickcolor=None, - ticklen=None, - tickwidth=None, - tickdir=None, - ticklabelpad=None, - ticklabelcolor=None, - ticklabelsize=None, - ticklabelweight=None, - tickminor=None, - ): - coord = self._get_coord_helper(axis) - if coord is None: - return - if tickminor is not None: - coord.display_minor_ticks(bool(tickminor)) - major = {} - if ticklen is not None: - major["length"] = ticklen - if tickwidth is not None: - major["width"] = tickwidth - if tickcolor is not None: - major["color"] = tickcolor - if tickdir is not None: - major["direction"] = tickdir - if ticklabelpad is not None: - major["pad"] = ticklabelpad - if ticklabelcolor is not None: - major["labelcolor"] = ticklabelcolor - if ticklabelsize is not None: - major["labelsize"] = ticklabelsize - if major: - coord.tick_params(**major) - if ticklabelweight is not None: - coord.set_ticklabel(weight=ticklabelweight) - if grid is not None or gridcolor is not None: - kw = {} - if gridcolor is not None: - kw["color"] = gridcolor - coord.grid(draw_grid=grid, **kw) - - def _update_axis_label( - self, - axis, - *, - label=None, - labelpad=None, - labelcolor=None, - labelsize=None, - labelweight=None, - label_kw=None, + from typing import override +except ImportError: + from typing_extensions import override + +ASTROPY_WCS_TYPES = (BaseLowLevelWCS, BaseHighLevelWCS) + + +class AstroAxes(base.Axes, WCSAxes): + """ + Native UltraPlot wrapper for Astropy WCS axes. + + This class keeps Astropy's `WCSAxes` drawing/transform machinery intact + while overriding the small subset of UltraPlot hooks needed for + formatting, sharing, and shared-label layout. + """ + + _name = "astro" + _name_aliases = ("astropy", "wcs") + + @override + def _update_background(self, **kwargs): + """ + Override `shared._SharedAxes._update_background` for WCS axes. + + WCSAxes owns its own patch artist, so the shared 2D helper can be + reused as long as we apply the resolved face/edge props directly to + `self.patch`. + """ + kw_face, kw_edge = rc._get_background_props(**kwargs) + self.patch.update(kw_face) + self.patch.update(kw_edge) + + def _get_coord_helper(self, axis): + """ + Return the Astropy coordinate helper backing logical ``x`` or ``y``. + + UltraPlot's formatting code talks in Cartesian ``x``/``y`` terms, + while WCSAxes exposes coordinate state through `self.coords[...]`. + This helper is the translation point between those APIs. + """ + index = {"x": 0, "y": 1}[axis] + try: + return self.coords[index] + except IndexError: + return None + + def _share_coord_signature(self, axis): + """ + Build a lightweight share-compatibility signature for one axis. + + Two Astro axes should only share if their coordinate family matches + in the ways that affect label/tick semantics. + """ + coord = self._get_coord_helper(axis) + if coord is None: + return None + unit = getattr(coord, "coord_unit", None) + if unit is not None and hasattr(unit, "to_string"): + unit = unit.to_string() + return ( + getattr(coord, "coord_type", None), + unit, + getattr(coord, "default_label", None), + ) + + def _update_coord_locator(self, axis, locator): + """ + Apply UltraPlot locator-style inputs to a WCS coordinate helper. + + This intentionally supports only the small subset that maps cleanly + onto Astropy's API. More advanced WCS locator setup should go + through `ax.coords[...]` directly. + """ + coord = self._get_coord_helper(axis) + if coord is None or locator is None: + return + if isinstance(locator, numbers.Real) and not isinstance(locator, bool): + coord.set_ticks(number=locator) + return + if isinstance(locator, Iterable) and not isinstance(locator, (str, bytes)): + coord.set_ticks(values=locator) + return + raise TypeError( + "AstroAxes.format only supports numeric or iterable tick locators. " + f"Received {locator!r}. Use ax.coords[...] for advanced locator setup." + ) + + def _update_coord_formatter(self, axis, formatter): + """ + Apply UltraPlot formatter inputs to a WCS coordinate helper. + + WCSAxes formatter configuration differs from Matplotlib's ordinary + axis formatter API, so this bridge keeps the supported surface + intentionally small and explicit. + """ + coord = self._get_coord_helper(axis) + if coord is None or formatter is None: + return + if isinstance(formatter, str) or callable(formatter): + coord.set_major_formatter(formatter) + return + raise TypeError( + "AstroAxes.format only supports string or callable tick formatters. " + f"Received {formatter!r}. Use ax.coords[...] for advanced formatter setup." + ) + + def _update_coord_ticks( + self, + axis, + *, + grid=None, + gridcolor=None, + tickcolor=None, + ticklen=None, + tickwidth=None, + tickdir=None, + ticklabelpad=None, + ticklabelcolor=None, + ticklabelsize=None, + ticklabelweight=None, + tickminor=None, + ): + """ + Translate UltraPlot tick/grid kwargs to Astropy coordinate helpers. + + This is the WCS equivalent of the shared Cartesian tick-update path: + collect the supported styling inputs and forward them to + `CoordinateHelper.tick_params()` / `grid()`. + """ + coord = self._get_coord_helper(axis) + if coord is None: + return + if tickminor is not None: + coord.display_minor_ticks(bool(tickminor)) + major = {} + if ticklen is not None: + major["length"] = ticklen + if tickwidth is not None: + major["width"] = tickwidth + if tickcolor is not None: + major["color"] = tickcolor + if tickdir is not None: + major["direction"] = tickdir + if ticklabelpad is not None: + major["pad"] = ticklabelpad + if ticklabelcolor is not None: + major["labelcolor"] = ticklabelcolor + if ticklabelsize is not None: + major["labelsize"] = ticklabelsize + if major: + coord.tick_params(**major) + if ticklabelweight is not None: + coord.set_ticklabel(weight=ticklabelweight) + if grid is not None or gridcolor is not None: + kw = {} + if gridcolor is not None: + kw["color"] = gridcolor + coord.grid(draw_grid=grid, **kw) + + def _update_axis_label( + self, + axis, + *, + label=None, + labelpad=None, + labelcolor=None, + labelsize=None, + labelweight=None, + label_kw=None, + ): + """ + Update WCS axis labels through Astropy's label API. + + This mirrors the behavior of `Axes.format` for Cartesian axes, but + delegates to `set_xlabel` / `set_ylabel` so Astropy can place and + style labels on the active coordinate helpers. + """ + coord = self._get_coord_helper(axis) + if coord is None: + return + if label is None and not any( + value is not None + for value in (labelpad, labelcolor, labelsize, labelweight) ): - coord = self._get_coord_helper(axis) - if coord is None: - return - if label is None and not any( - value is not None - for value in (labelpad, labelcolor, labelsize, labelweight) - ): - return - setter = getattr(self, f"set_{axis}label") - getter = getattr(self, f"get_{axis}label") - kw = dict(label_kw or {}) - if labelcolor is not None: - kw["color"] = labelcolor - if labelsize is not None: - kw["size"] = labelsize - if labelweight is not None: - kw["weight"] = labelweight - if labelpad is not None: - kw["labelpad"] = labelpad - setter(getter() if label is None else label, **kw) - - def _update_limits(self, axis, *, lim=None, min_=None, max_=None, reverse=None): - lo = hi = None - if lim is not None: - lo, hi = lim - lo = _not_none(min_=min_, lim_0=lo) - hi = _not_none(max_=max_, lim_1=hi) - if lo is not None or hi is not None: - get_lim = getattr(self, f"get_{axis}lim") - set_lim = getattr(self, f"set_{axis}lim") - cur_lo, cur_hi = get_lim() - set_lim((_not_none(lo, cur_lo), _not_none(hi, cur_hi))) - if reverse is not None: - inverted = getattr(self, f"{axis}axis_inverted")() - if bool(reverse) != bool(inverted): - getattr(self, f"invert_{axis}axis")() - - def _share_axis_limits(self, other, which): - self._shared_axes[which].join(self, other) - axis = getattr(self, f"{which}axis") - other_axis = getattr(other, f"{which}axis") - setattr(self, f"_share{which}", other) - axis.major = other_axis.major - axis.minor = other_axis.minor - get_lim = getattr(other, f"get_{which}lim") - set_lim = getattr(self, f"set_{which}lim") - get_auto = getattr(other, f"get_autoscale{which}_on") - set_lim(*get_lim(), emit=False, auto=get_auto()) - axis._scale = other_axis._scale - - def _sharex_setup(self, sharex, *, labels=True, limits=True): - super()._sharex_setup(sharex) - level = ( - 3 - if self._panel_sharex_group and self._is_panel_group_member(sharex) - else self.figure._sharex - ) - if level not in range(5): - raise ValueError(f"Invalid sharing level sharex={level!r}.") - if sharex in (None, self) or not isinstance(sharex, AstroAxes): - return - if level > 0 and labels: - self._sharex = sharex - if level > 1 and limits: - self._share_axis_limits(sharex, "x") - - def _sharey_setup(self, sharey, *, labels=True, limits=True): - super()._sharey_setup(sharey) - level = ( - 3 - if self._panel_sharey_group and self._is_panel_group_member(sharey) - else self.figure._sharey - ) - if level not in range(5): - raise ValueError(f"Invalid sharing level sharey={level!r}.") - if sharey in (None, self) or not isinstance(sharey, AstroAxes): - return - if level > 0 and labels: - self._sharey = sharey - if level > 1 and limits: - self._share_axis_limits(sharey, "y") - - def _is_ticklabel_on(self, side: str) -> bool: - axis = "x" if side in ("labelbottom", "labeltop") else "y" - coord = self._get_coord_helper(axis) - if coord is None or not coord.get_ticklabel_visible(): - return False - positions = coord.get_ticklabel_position() - tokens = { - "labelbottom": "b", - "labeltop": "t", - "labelleft": "l", - "labelright": "r", - "bottom": "b", - "top": "t", - "left": "l", - "right": "r", - } - token = tokens.get(side, side) - if token in positions: - return True - if "#" in positions: - return token == ("b" if axis == "x" else "l") + return + setter = getattr(self, f"set_{axis}label") + getter = getattr(self, f"get_{axis}label") + kw = dict(label_kw or {}) + if labelcolor is not None: + kw["color"] = labelcolor + if labelsize is not None: + kw["size"] = labelsize + if labelweight is not None: + kw["weight"] = labelweight + if labelpad is not None: + kw["labelpad"] = labelpad + setter(getter() if label is None else label, **kw) + + def _update_limits(self, axis, *, lim=None, min_=None, max_=None, reverse=None): + """ + Apply x/y limit and inversion requests using the normal Matplotlib API. + + WCSAxes still exposes `get_xlim`, `set_xlim`, etc., so we keep limit + handling on the axes object itself rather than trying to route it + through `coords[...]`. + """ + lo = hi = None + if lim is not None: + lo, hi = lim + lo = _not_none(min_=min_, lim_0=lo) + hi = _not_none(max_=max_, lim_1=hi) + if lo is not None or hi is not None: + get_lim = getattr(self, f"get_{axis}lim") + set_lim = getattr(self, f"set_{axis}lim") + cur_lo, cur_hi = get_lim() + set_lim((_not_none(lo, cur_lo), _not_none(hi, cur_hi))) + if reverse is not None: + inverted = getattr(self, f"{axis}axis_inverted")() + if bool(reverse) != bool(inverted): + getattr(self, f"invert_{axis}axis")() + + def _share_axis_limits(self, other, which): + """ + Share axis limit state with another Astro axes instance. + + This mirrors the relevant parts of Matplotlib's shared-axis setup + while staying inside UltraPlot's higher-level share policy. + """ + self._shared_axes[which].join(self, other) + axis = getattr(self, f"{which}axis") + other_axis = getattr(other, f"{which}axis") + setattr(self, f"_share{which}", other) + axis.major = other_axis.major + axis.minor = other_axis.minor + get_lim = getattr(other, f"get_{which}lim") + set_lim = getattr(self, f"set_{which}lim") + get_auto = getattr(other, f"get_autoscale{which}_on") + set_lim(*get_lim(), emit=False, auto=get_auto()) + axis._scale = other_axis._scale + + @override + def _sharex_setup(self, sharex, *, labels=True, limits=True): + """ + Override `base.Axes._sharex_setup` for Astro-aware share policy. + + Astro axes can share labels and limits with other Astro axes, but we + keep the compatibility check narrow so incompatible WCS coordinate + families do not silently enter the same share group. + """ + super()._sharex_setup(sharex) + level = ( + 3 + if self._panel_sharex_group and self._is_panel_group_member(sharex) + else self.figure._sharex + ) + if level not in range(5): + raise ValueError(f"Invalid sharing level sharex={level!r}.") + if sharex in (None, self) or not isinstance(sharex, AstroAxes): + return + if level > 0 and labels: + self._sharex = sharex + if level > 1 and limits: + self._share_axis_limits(sharex, "x") + + @override + def _sharey_setup(self, sharey, *, labels=True, limits=True): + """ + Override `base.Axes._sharey_setup` for Astro-aware share policy. + """ + super()._sharey_setup(sharey) + level = ( + 3 + if self._panel_sharey_group and self._is_panel_group_member(sharey) + else self.figure._sharey + ) + if level not in range(5): + raise ValueError(f"Invalid sharing level sharey={level!r}.") + if sharey in (None, self) or not isinstance(sharey, AstroAxes): + return + if level > 0 and labels: + self._sharey = sharey + if level > 1 and limits: + self._share_axis_limits(sharey, "y") + + def _is_ticklabel_on(self, side: str) -> bool: + """ + Interpret Astropy ticklabel position tokens as UltraPlot booleans. + + Astropy can return explicit side tokens (``'t'``, ``'b'``, ``'l'``, + ``'r'``) or the special ``'#'`` default token. Figure-level sharing + logic wants plain on/off state per side, so we normalize that here. + """ + axis = "x" if side in ("labelbottom", "labeltop") else "y" + coord = self._get_coord_helper(axis) + if coord is None or not coord.get_ticklabel_visible(): return False + positions = coord.get_ticklabel_position() + tokens = { + "labelbottom": "b", + "labeltop": "t", + "labelleft": "l", + "labelright": "r", + "bottom": "b", + "top": "t", + "left": "l", + "right": "r", + } + token = tokens.get(side, side) + if token in positions: + return True + # These are default tokens used by Astropy to indicate sides. + if "#" in positions: + return token == ("b" if axis == "x" else "l") + return False - def _get_ticklabel_state(self, axis: str) -> dict[str, bool]: - sides = ("top", "bottom") if axis == "x" else ("left", "right") - return { - f"label{side}": self._is_ticklabel_on(f"label{side}") for side in sides - } - - def _set_ticklabel_state(self, axis: str, state: dict): - coord = self._get_coord_helper(axis) - if coord is None: - return - positions = [] - for side in ("bottom", "top") if axis == "x" else ("left", "right"): - if state.get(f"label{side}", False): - positions.append(side[0]) - position = "".join(positions) - coord.set_ticklabel_position(position) - coord.set_axislabel_position(position) - coord.set_ticklabel_visible(bool(positions)) - - def _apply_ticklabel_state(self, axis: str, state: dict): - self._set_ticklabel_state(axis, state) - - def format( - self, - *, - aspect=None, - xreverse=None, - yreverse=None, - xlim=None, - ylim=None, - xmin=None, - ymin=None, - xmax=None, - ymax=None, - xformatter=None, - yformatter=None, - xlocator=None, - ylocator=None, - xtickminor=None, - ytickminor=None, - xtickcolor=None, - ytickcolor=None, - xticklen=None, - yticklen=None, - xtickwidth=None, - ytickwidth=None, - xtickdir=None, - ytickdir=None, - xticklabelpad=None, - yticklabelpad=None, - xticklabelcolor=None, - yticklabelcolor=None, - xticklabelsize=None, - yticklabelsize=None, - xticklabelweight=None, - yticklabelweight=None, - xlabel=None, - ylabel=None, - xlabelpad=None, - ylabelpad=None, - xlabelcolor=None, - ylabelcolor=None, - xlabelsize=None, - ylabelsize=None, - xlabelweight=None, - ylabelweight=None, - xgrid=None, - ygrid=None, - xgridcolor=None, - ygridcolor=None, - xlabel_kw=None, - ylabel_kw=None, - **kwargs, - ): - if aspect is not None: - self.set_aspect(aspect) - self._update_limits("x", lim=xlim, min_=xmin, max_=xmax, reverse=xreverse) - self._update_limits("y", lim=ylim, min_=ymin, max_=ymax, reverse=yreverse) - self._update_coord_locator("x", xlocator) - self._update_coord_locator("y", ylocator) - self._update_coord_formatter("x", xformatter) - self._update_coord_formatter("y", yformatter) - self._update_coord_ticks( - "x", - grid=xgrid, - gridcolor=xgridcolor, - tickcolor=xtickcolor, - ticklen=xticklen, - tickwidth=xtickwidth, - tickdir=xtickdir, - ticklabelpad=xticklabelpad, - ticklabelcolor=xticklabelcolor, - ticklabelsize=xticklabelsize, - ticklabelweight=xticklabelweight, - tickminor=xtickminor, - ) - self._update_coord_ticks( - "y", - grid=ygrid, - gridcolor=ygridcolor, - tickcolor=ytickcolor, - ticklen=yticklen, - tickwidth=ytickwidth, - tickdir=ytickdir, - ticklabelpad=yticklabelpad, - ticklabelcolor=yticklabelcolor, - ticklabelsize=yticklabelsize, - ticklabelweight=yticklabelweight, - tickminor=ytickminor, - ) - self._update_axis_label( - "x", - label=xlabel, - labelpad=xlabelpad, - labelcolor=xlabelcolor, - labelsize=xlabelsize, - labelweight=xlabelweight, - label_kw=xlabel_kw, - ) - self._update_axis_label( - "y", - label=ylabel, - labelpad=ylabelpad, - labelcolor=ylabelcolor, - labelsize=ylabelsize, - labelweight=ylabelweight, - label_kw=ylabel_kw, - ) - return base.Axes.format(self, **kwargs) - - AstroAxes._format_signatures[AstroAxes] = inspect.signature(CartesianAxes.format) -else: # pragma: no cover - AstroAxes = None + @override + def _get_ticklabel_state(self, axis: str) -> dict[str, bool]: + """ + Override `base.Axes._get_ticklabel_state` for WCS ticklabel sides. + """ + sides = ("top", "bottom") if axis == "x" else ("left", "right") + return {f"label{side}": self._is_ticklabel_on(f"label{side}") for side in sides} + + @override + def _set_ticklabel_state(self, axis: str, state: dict): + """ + Override `base.Axes._set_ticklabel_state` using Astropy side tokens. + + Figure-level sharing/panel code passes the same state dictionary used + by Cartesian axes, and this method converts it into the position + string expected by `CoordinateHelper`. + """ + coord = self._get_coord_helper(axis) + if coord is None: + return + positions = [] + for side in ("bottom", "top") if axis == "x" else ("left", "right"): + if state.get(f"label{side}", False): + positions.append(side[0]) + position = "".join(positions) + coord.set_ticklabel_position(position) + coord.set_axislabel_position(position) + coord.set_ticklabel_visible(bool(positions)) + + def _apply_ticklabel_state(self, axis: str, state: dict): + """ + Local helper used by the figure-sharing bridge. + + This is not overriding a base method; it just gives the figure-side + label-sharing logic a single entry point for Astro ticklabel state. + """ + self._set_ticklabel_state(axis, state) + + @override + def format( + self, + *, + aspect=None, + xreverse=None, + yreverse=None, + xlim=None, + ylim=None, + xmin=None, + ymin=None, + xmax=None, + ymax=None, + xformatter=None, + yformatter=None, + xlocator=None, + ylocator=None, + xtickminor=None, + ytickminor=None, + xtickcolor=None, + ytickcolor=None, + xticklen=None, + yticklen=None, + xtickwidth=None, + ytickwidth=None, + xtickdir=None, + ytickdir=None, + xticklabelpad=None, + yticklabelpad=None, + xticklabelcolor=None, + yticklabelcolor=None, + xticklabelsize=None, + yticklabelsize=None, + xticklabelweight=None, + yticklabelweight=None, + xlabel=None, + ylabel=None, + xlabelpad=None, + ylabelpad=None, + xlabelcolor=None, + ylabelcolor=None, + xlabelsize=None, + ylabelsize=None, + xlabelweight=None, + ylabelweight=None, + xgrid=None, + ygrid=None, + xgridcolor=None, + ygridcolor=None, + xlabel_kw=None, + ylabel_kw=None, + **kwargs, + ): + """ + Override `base.Axes.format` with a narrow WCS-aware front-end. + + The Astro-specific pieces are applied first through the coordinate + helpers, then the remaining generic UltraPlot formatting is delegated + back to `base.Axes.format`. + """ + if aspect is not None: + self.set_aspect(aspect) + self._update_limits("x", lim=xlim, min_=xmin, max_=xmax, reverse=xreverse) + self._update_limits("y", lim=ylim, min_=ymin, max_=ymax, reverse=yreverse) + self._update_coord_locator("x", xlocator) + self._update_coord_locator("y", ylocator) + self._update_coord_formatter("x", xformatter) + self._update_coord_formatter("y", yformatter) + self._update_coord_ticks( + "x", + grid=xgrid, + gridcolor=xgridcolor, + tickcolor=xtickcolor, + ticklen=xticklen, + tickwidth=xtickwidth, + tickdir=xtickdir, + ticklabelpad=xticklabelpad, + ticklabelcolor=xticklabelcolor, + ticklabelsize=xticklabelsize, + ticklabelweight=xticklabelweight, + tickminor=xtickminor, + ) + self._update_coord_ticks( + "y", + grid=ygrid, + gridcolor=ygridcolor, + tickcolor=ytickcolor, + ticklen=yticklen, + tickwidth=ytickwidth, + tickdir=ytickdir, + ticklabelpad=yticklabelpad, + ticklabelcolor=yticklabelcolor, + ticklabelsize=yticklabelsize, + ticklabelweight=yticklabelweight, + tickminor=ytickminor, + ) + self._update_axis_label( + "x", + label=xlabel, + labelpad=xlabelpad, + labelcolor=xlabelcolor, + labelsize=xlabelsize, + labelweight=xlabelweight, + label_kw=xlabel_kw, + ) + self._update_axis_label( + "y", + label=ylabel, + labelpad=ylabelpad, + labelcolor=ylabelcolor, + labelsize=ylabelsize, + labelweight=ylabelweight, + label_kw=ylabel_kw, + ) + return base.Axes.format(self, **kwargs) + + +AstroAxes._format_signatures[AstroAxes] = inspect.signature(CartesianAxes.format) diff --git a/ultraplot/internals/projections.py b/ultraplot/internals/projections.py index 61eb45f01..fea9828fb 100644 --- a/ultraplot/internals/projections.py +++ b/ultraplot/internals/projections.py @@ -158,8 +158,7 @@ def _wrap_external_projection(figure, projection): and proj in ("astro", "astropy", "wcs", "ultraplot_astro"), ) def _resolve_astropy_wcs_string(proj, context): - if _get_axes_module().get_astro_axes_class(load=True) is None: - return ProjectionResolution() + _get_axes_module().get_astro_axes_class(load=True) return ProjectionResolution(projection="ultraplot_astro") diff --git a/ultraplot/tests/test_imports.py b/ultraplot/tests/test_imports.py index 168b2e627..da706ef2a 100644 --- a/ultraplot/tests/test_imports.py +++ b/ultraplot/tests/test_imports.py @@ -56,11 +56,19 @@ def test_axes_astro_attr_is_lazy_optional(): import sys import ultraplot.axes as paxes spec = importlib.util.find_spec("astropy.visualization.wcsaxes") -astro = paxes.AstroAxes +error = None +astro_is_none = None +try: + astro = paxes.AstroAxes +except ImportError as exc: + error = str(exc) +else: + astro_is_none = astro is None mods = [name for name in sys.modules if name == "astropy" or name.startswith("astropy.")] print(json.dumps({ "available": bool(spec), - "astro_is_none": astro is None, + "astro_is_none": astro_is_none, + "error": error, "loaded": bool(mods), })) """ @@ -69,7 +77,8 @@ def test_axes_astro_attr_is_lazy_optional(): assert not out["astro_is_none"] assert out["loaded"] else: - assert out["astro_is_none"] + assert out["error"] is not None + assert "requires astropy" in out["error"] def test_star_import_exposes_public_api(): From 625aa12e0e4767ef00be6d7852bbf331e9108571 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Thu, 12 Mar 2026 09:29:32 +1000 Subject: [PATCH 07/12] Document Astropy WCS support in the user guide Add a dedicated Astropy WCS section to docs/usage.rst so the new AstroAxes integration is discoverable from the main usage guide. The section explains the optional install path, shows the two supported creation patterns (direct WCS objects and the "wcs" alias), and clarifies that UltraPlot handles figure-level sharing/layout while advanced coordinate customization still lives on ax.coords[...]. --- docs/usage.rst | 42 +++++++++++++++++++++++++++++++++++++++--- 1 file changed, 39 insertions(+), 3 deletions(-) diff --git a/docs/usage.rst b/docs/usage.rst index 298c12132..2df4977d3 100644 --- a/docs/usage.rst +++ b/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 `, @@ -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,38 @@ 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]" + +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) From 899ef533b2182b60cd2d537012136bc2667cb042 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Thu, 12 Mar 2026 09:30:51 +1000 Subject: [PATCH 08/12] Add basemap to Geo --- pyproject.toml | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 50c5d2d6e..f21f07f01 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -79,9 +79,7 @@ astropy = [ ] geo = [ "cartopy", -] -cartopy = [ - "cartopy", + "basemap", ] ternary = [ "mpltern", @@ -92,6 +90,7 @@ mpltern = [ all = [ "astropy", "cartopy", + "basemap", "mpltern", ] docs = [ From fb10fa86cf40eb2e7f583e5e2e612e53cec33aed Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Thu, 12 Mar 2026 09:59:11 +1000 Subject: [PATCH 09/12] Restore legacy non-sharing panel layout Revert the non-sharing left and bottom panel relabeling that shifted shared subplot labels and broke existing image baselines. Also harden the Astro lazy-import test so missing astropy does not crash the subprocess before the optional import path is exercised. --- ultraplot/figure.py | 35 --------------------- ultraplot/tests/test_imports.py | 5 ++- ultraplot/tests/test_subplots.py | 53 -------------------------------- 3 files changed, 4 insertions(+), 89 deletions(-) diff --git a/ultraplot/figure.py b/ultraplot/figure.py index 7ecbf30e6..44dac55c6 100644 --- a/ultraplot/figure.py +++ b/ultraplot/figure.py @@ -2108,41 +2108,6 @@ def _add_axes_panel( pax._label_key("labelleft"): False, } ) - elif side == "left" and not share and not filled: - ax.yaxis.tick_right() - ax.yaxis.set_label_position("right") - ax.yaxis.set_tick_params( - **{ - ax._label_key("labelleft"): False, - ax._label_key("labelright"): True, - } - ) - pax.yaxis.tick_left() - pax.yaxis.set_label_position("left") - pax.yaxis.set_tick_params( - **{ - pax._label_key("labelleft"): True, - pax._label_key("labelright"): False, - } - ) - elif side == "bottom" and not share and not filled: - ax.xaxis.tick_top() - ax.xaxis.set_label_position("top") - ax.xaxis.set_tick_params( - **{ - ax._label_key("labelbottom"): False, - ax._label_key("labeltop"): True, - } - ) - pax.xaxis.tick_bottom() - pax.xaxis.set_label_position("bottom") - pax.xaxis.set_tick_params( - **{ - pax._label_key("labelbottom"): True, - pax._label_key("labeltop"): False, - } - ) - return pax @_clear_border_cache diff --git a/ultraplot/tests/test_imports.py b/ultraplot/tests/test_imports.py index da706ef2a..2f728b1bf 100644 --- a/ultraplot/tests/test_imports.py +++ b/ultraplot/tests/test_imports.py @@ -55,7 +55,10 @@ def test_axes_astro_attr_is_lazy_optional(): import json import sys import ultraplot.axes as paxes -spec = importlib.util.find_spec("astropy.visualization.wcsaxes") +try: + spec = importlib.util.find_spec("astropy.visualization.wcsaxes") +except ModuleNotFoundError: + spec = None error = None astro_is_none = None try: diff --git a/ultraplot/tests/test_subplots.py b/ultraplot/tests/test_subplots.py index 08acfecf4..c9922bfe7 100644 --- a/ultraplot/tests/test_subplots.py +++ b/ultraplot/tests/test_subplots.py @@ -830,59 +830,6 @@ def test_panel_share_flag_controls_group_membership(): assert ax2[0]._panel_sharex_group is False -def test_nonsharing_left_panel_moves_main_labels_outside(): - fig, axs = uplt.subplots() - ax = axs[0] - ax.format(ylabel="main ylabel") - pax = ax.panel("left", share=False) - pax.format(ylabel="panel ylabel") - - fig.canvas.draw() - - assert not ax._is_ticklabel_on("labelleft") - assert ax._is_ticklabel_on("labelright") - assert pax._is_ticklabel_on("labelleft") - assert not pax._is_ticklabel_on("labelright") - assert ax.yaxis.get_label_position() == "right" - assert pax.yaxis.get_label_position() == "left" - - -def test_nonsharing_bottom_panel_moves_main_labels_outside(): - fig, axs = uplt.subplots() - ax = axs[0] - ax.format(xlabel="main xlabel") - pax = ax.panel("bottom", share=False) - pax.format(xlabel="panel xlabel") - - fig.canvas.draw() - - assert not ax._is_ticklabel_on("labelbottom") - assert ax._is_ticklabel_on("labeltop") - assert pax._is_ticklabel_on("labelbottom") - assert not pax._is_ticklabel_on("labeltop") - assert ax.xaxis.get_label_position() == "top" - assert pax.xaxis.get_label_position() == "bottom" - - -def test_nonsharing_left_panel_gap_matches_right_panel(): - def _panel_gap(side): - fig, axs = uplt.subplots() - ax = axs[0] - ax.format(ylabel="main ylabel") - pax = ax.panel(side, share=False) - pax.format(xlabel="panel xlabel", ylabel="panel ylabel") - fig.canvas.draw() - main = ax.get_position().bounds - panel = pax.get_position().bounds - if side == "left": - return main[0] - (panel[0] + panel[2]) - return panel[0] - (main[0] + main[2]) - - gap_left = _panel_gap("left") - gap_right = _panel_gap("right") - assert abs(gap_left - gap_right) < 1e-3 - - def test_ticklabels_with_guides_share_true_cartesian(): """ With share=True, tick labels should only appear on bottom row and left column From a7d77f768472a12c8141c16fca53b1c50384f393 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Tue, 14 Jul 2026 11:58:18 +1000 Subject: [PATCH 10/12] Fix WCS sharing after main merge --- ultraplot/figure.py | 22 ++++++++++------------ ultraplot/tests/test_projections.py | 4 +--- 2 files changed, 11 insertions(+), 15 deletions(-) diff --git a/ultraplot/figure.py b/ultraplot/figure.py index 44dac55c6..c66e0730a 100644 --- a/ultraplot/figure.py +++ b/ultraplot/figure.py @@ -1507,12 +1507,11 @@ def _share_ticklabels(self, *, axis: str) -> None: main_axes = [ axi for axi in group_axes if not getattr(axi, "_panel_side", None) ] - supported_main_axes = any( - isinstance( - axi, (paxes.CartesianAxes, paxes._CartopyAxes, paxes._BasemapAxes) - ) - for axi in main_axes - ) + supported = (paxes.CartesianAxes, paxes._CartopyAxes, paxes._BasemapAxes) + astro_cls = paxes.get_astro_axes_class() + if astro_cls is not None: + supported = (*supported, astro_cls) + supported_main_axes = any(isinstance(axi, supported) for axi in main_axes) if len(group_axes) < 2 and not supported_main_axes: continue if all( @@ -1571,12 +1570,11 @@ def _compute_baseline_tick_state(self, group_axes, axis: str): sides = ("top", "bottom") if axis == "x" else ("left", "right") main_axes = [axi for axi in group_axes if not getattr(axi, "_panel_side", None)] if len(main_axes) < 2: - supported = all( - isinstance( - axi, (paxes.CartesianAxes, paxes._CartopyAxes, paxes._BasemapAxes) - ) - for axi in main_axes - ) + supported = (paxes.CartesianAxes, paxes._CartopyAxes, paxes._BasemapAxes) + astro_cls = paxes.get_astro_axes_class() + if astro_cls is not None: + supported = (*supported, astro_cls) + supported = all(isinstance(axi, supported) for axi in main_axes) if not supported: return {}, True diff --git a/ultraplot/tests/test_projections.py b/ultraplot/tests/test_projections.py index 1c5a5c54a..9dac038d9 100644 --- a/ultraplot/tests/test_projections.py +++ b/ultraplot/tests/test_projections.py @@ -320,9 +320,7 @@ def test_taylor_projection_validation_errors(): def test_taylor_single_axes_skips_shared_ticklabel_baseline(): fig, axs = uplt.subplots(proj="taylor") - baseline, skip = fig._compute_baseline_tick_state( - [axs[0]], "x", ("labelbottom", "labeltop") - ) + baseline, skip = fig._compute_baseline_tick_state([axs[0]], "x") assert baseline == {} assert skip From 4775a12eb10af6486129d7feca8eabc2317a7df4 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Thu, 16 Jul 2026 15:52:52 +1000 Subject: [PATCH 11/12] Fix projection resolution edge cases --- ultraplot/_subplots.py | 10 ++----- ultraplot/figure.py | 4 +++ ultraplot/internals/projections.py | 13 ++++---- ultraplot/tests/test_astro_axes.py | 18 +++++++++++ ultraplot/tests/test_subplot_manager.py | 40 +++++++++++++++++++++++++ 5 files changed, 72 insertions(+), 13 deletions(-) diff --git a/ultraplot/_subplots.py b/ultraplot/_subplots.py index 5911b70a1..9eb88b06d 100644 --- a/ultraplot/_subplots.py +++ b/ultraplot/_subplots.py @@ -182,13 +182,9 @@ def add_subplot(self, *args, **kwargs): 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 diff --git a/ultraplot/figure.py b/ultraplot/figure.py index c66e0730a..1caa09c7e 100644 --- a/ultraplot/figure.py +++ b/ultraplot/figure.py @@ -2156,6 +2156,10 @@ def _add_subplot(self, *args, **kwargs): """Delegate to SubplotManager.""" return self._subplots.add_subplot(*args, **kwargs) + def _add_subplot_mpl(self, *args, **kwargs): + """Dispatch subplot creation to matplotlib without manager recursion.""" + return super().add_subplot(*args, **kwargs) + def _unshare_axes(self): for which in "xyz": diff --git a/ultraplot/internals/projections.py b/ultraplot/internals/projections.py index fea9828fb..699de937c 100644 --- a/ultraplot/internals/projections.py +++ b/ultraplot/internals/projections.py @@ -82,11 +82,14 @@ def _get_axes_module(): def _looks_like_astropy_projection(proj): - module = getattr(type(proj), "__module__", "") - return module.startswith("astropy.") + return any( + getattr(cls, "__module__", "").startswith("astropy.") + for cls in type(proj).__mro__ + ) def _prefixed_projection_name(name): + name = name.lower() if name.startswith("ultraplot_"): return name if name in mproj.get_projection_names() else None prefixed = "ultraplot_" + name @@ -155,7 +158,7 @@ def _wrap_external_projection(figure, projection): @register_projection_binding( "astropy_wcs_string", lambda proj, context: isinstance(proj, str) - and proj in ("astro", "astropy", "wcs", "ultraplot_astro"), + and proj.lower() in ("astro", "astropy", "wcs", "ultraplot_astro"), ) def _resolve_astropy_wcs_string(proj, context): _get_axes_module().get_astro_axes_class(load=True) @@ -222,7 +225,7 @@ def _resolve_basemap_projection_object(proj, context): def _resolve_geographic_projection_name(proj, context): try: proj_obj = constructor.Proj( - proj, + proj.lower(), backend=context.backend, include_axes=True, **context.proj_kw, @@ -249,8 +252,6 @@ def resolve_projection(proj, *, figure, proj_kw=None, backend=None): Resolve a user projection spec to a final projection and kwargs. """ proj_kw = proj_kw or {} - if isinstance(proj, str): - proj = proj.lower() context = ProjectionContext(figure=figure, proj_kw=proj_kw, backend=backend) resolution = None diff --git a/ultraplot/tests/test_astro_axes.py b/ultraplot/tests/test_astro_axes.py index 8403263ff..e88b6de5b 100644 --- a/ultraplot/tests/test_astro_axes.py +++ b/ultraplot/tests/test_astro_axes.py @@ -20,6 +20,24 @@ def _make_test_wcs(): return wcs +class CustomWCS(WCS): + """Application-defined WCS subclass outside the astropy namespace.""" + + +@pytest.mark.parametrize("method", ["add_axes", "add_subplot"]) +def test_custom_wcs_subclass_uses_native_astro_axes(method): + wcs = CustomWCS(_make_test_wcs().to_header()) + fig = uplt.figure() + if method == "add_axes": + ax = fig.add_axes([0.1, 0.1, 0.8, 0.8], projection=wcs) + else: + ax = fig.add_subplot(111, projection=wcs) + + assert isinstance(ax, paxes.AstroAxes) + assert ax.wcs is wcs + uplt.close(fig) + + def test_add_subplot_with_wcs_projection_returns_native_astro_axes(): fig = uplt.figure() ax = fig.add_subplot(111, projection=_make_test_wcs()) diff --git a/ultraplot/tests/test_subplot_manager.py b/ultraplot/tests/test_subplot_manager.py index 46e6bb06d..d82421b44 100644 --- a/ultraplot/tests/test_subplot_manager.py +++ b/ultraplot/tests/test_subplot_manager.py @@ -4,6 +4,7 @@ import inspect +import matplotlib.axes as maxes import matplotlib.projections as mproj import numpy as np import pytest @@ -359,6 +360,45 @@ def test_add_subplot_external_projection_reuses_container(): uplt.close(fig) +def test_case_sensitive_matplotlib_projection_name(): + """Registered matplotlib projection names retain their exact spelling.""" + + class CaseSensitiveAxes(maxes.Axes): + name = "UltraPlot.CaseSensitive.Test" + + mproj.register_projection(CaseSensitiveAxes) + fig = uplt.figure() + ax = fig.add_axes([0.1, 0.1, 0.35, 0.8], projection=CaseSensitiveAxes.name) + sax = fig.add_subplot(122, projection=CaseSensitiveAxes.name) + + assert isinstance(ax, ExternalAxesContainer) + assert isinstance(sax, ExternalAxesContainer) + assert isinstance(ax.get_external_child(), CaseSensitiveAxes) + assert isinstance(sax.get_external_child(), CaseSensitiveAxes) + uplt.close(fig) + + +def test_add_subplot_mpl_extension_hook(): + """Figure subclasses can intercept matplotlib subplot creation safely.""" + + class SubplotHookMixin: + def _add_subplot_mpl(self, *args, **kwargs): + self.subplot_hook_calls += 1 + return super()._add_subplot_mpl(*args, **kwargs) + + class HookFigure(SubplotHookMixin, pfigure.Figure): + def __init__(self, *args, **kwargs): + self.subplot_hook_calls = 0 + super().__init__(*args, **kwargs) + + fig = HookFigure() + ax = fig.add_subplot(111) + + assert ax.figure is fig + assert fig.subplot_hook_calls == 1 + uplt.close(fig) + + @pytest.mark.parametrize("key", ["proj", "projection"]) def test_ui_subplot_routes_projection_kwargs(key): """ From f7fa13fb10308c3d61dbbf1c49642d8dfaf4314f Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Thu, 13 Aug 2026 14:51:25 +1000 Subject: [PATCH 12/12] fix: preserve named projection container compatibility --- ultraplot/_subplots.py | 2 -- ultraplot/internals/projections.py | 8 ++++++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/ultraplot/_subplots.py b/ultraplot/_subplots.py index 9eb88b06d..0cd65a278 100644 --- a/ultraplot/_subplots.py +++ b/ultraplot/_subplots.py @@ -7,10 +7,8 @@ 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 diff --git a/ultraplot/internals/projections.py b/ultraplot/internals/projections.py index 699de937c..c19f14ec6 100644 --- a/ultraplot/internals/projections.py +++ b/ultraplot/internals/projections.py @@ -111,7 +111,8 @@ def _wrap_external_projection(figure, projection): external_axes_class = None external_axes_kwargs = {} - if isinstance(projection, str): + projection_name = projection if isinstance(projection, str) else None + if projection_name is not None: if projection.startswith("ultraplot_") or projection.startswith( "_ultraplot_container_" ): @@ -139,7 +140,10 @@ def _wrap_external_projection(figure, projection): from ..axes.container import create_external_axes_container - container_name = _container_projection_name(external_axes_class) + if projection_name is None: + container_name = _container_projection_name(external_axes_class) + else: + container_name = "_ultraplot_container_" + projection_name if container_name not in mproj.get_projection_names(): container_class = create_external_axes_container( external_axes_class, projection_name=container_name