From 1e95beeb68c8b7f6c9c91faf0231eb2041d21ef7 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Wed, 29 Jul 2026 20:41:28 +1000 Subject: [PATCH 1/6] perf: avoid redundant interactive layout work --- ultraplot/_animation.py | 231 ++++++++++++++++++++++++++++++ ultraplot/axes/_formatting.py | 43 ++++++ ultraplot/axes/base.py | 2 +- ultraplot/axes/cartesian.py | 21 ++- ultraplot/figure.py | 67 ++++++++- ultraplot/tests/test_animation.py | 178 +++++++++++++++++++++++ 6 files changed, 533 insertions(+), 9 deletions(-) create mode 100644 ultraplot/_animation.py diff --git a/ultraplot/_animation.py b/ultraplot/_animation.py new file mode 100644 index 000000000..026365cad --- /dev/null +++ b/ultraplot/_animation.py @@ -0,0 +1,231 @@ +#!/usr/bin/env python3 +""" +Helpers for responsive interactive and animated UltraPlot figures. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from contextlib import contextmanager +from weakref import WeakSet + +import matplotlib.artist as martist +import matplotlib.transforms as mtransforms + + +class _BlitManager: + """ + Manage efficient updates of a small set of changing artists. + + The manager caches the static canvas background, restores it for each + update, redraws only the managed artists, and blits the affected region. + Backends without blitting support safely fall back to ``draw_idle()``. + + Parameters + ---------- + canvas : `~matplotlib.backend_bases.FigureCanvasBase` + Canvas containing the artists. + artists : iterable of `~matplotlib.artist.Artist`, optional + Artists that will change between updates. + bbox : `~matplotlib.transforms.Bbox` or object with a ``bbox`` attribute, optional + Region to cache and blit. By default, the union of the managed artists' + axes bounding boxes is used. Figure-level artists fall back to the full + figure bounding box. + + Notes + ----- + Managed artists are drawn above the cached static background, matching + Matplotlib's standard blitting behavior. + """ + + def __init__(self, canvas, artists: Iterable[martist.Artist] = (), bbox=None): + if canvas.figure is None: + raise RuntimeError("Cannot manage blitting for a canvas without a figure.") + self.canvas = canvas + self.figure = canvas.figure + self._artists = [] + self._animated = {} + self._bbox = bbox + self._background = None + self._closed = False + self._suspended = False + self._supports_blit = bool(canvas.supports_blit) + self._draw_cid = canvas.mpl_connect("draw_event", self._on_draw) + self._resize_cid = canvas.mpl_connect("resize_event", self._on_resize) + managers = getattr(self.figure, "_blit_managers", None) + if managers is None: + managers = self.figure._blit_managers = WeakSet() + managers.add(self) + for artist in artists: + self.add_artist(artist) + + @property + def artists(self): + """Managed artists as an immutable tuple.""" + return tuple(self._artists) + + @property + def supports_blit(self): + """Whether the associated canvas supports blitting.""" + return self._supports_blit + + def _resolve_bbox(self): + bbox = self._bbox + if bbox is not None: + return getattr(bbox, "bbox", bbox) + + axes = [] + for artist in self._artists: + axis = getattr(artist, "axes", None) + if axis is None: + return self.figure.bbox + if axis not in axes: + axes.append(axis) + if not axes: + return self.figure.bbox + return mtransforms.Bbox.union([axis.bbox for axis in axes]) + + def _draw_artists(self): + for artist in sorted(self._artists, key=lambda item: item.get_zorder()): + self.figure.draw_artist(artist) + + def _on_draw(self, event): + if self._closed or self._suspended or not self._supports_blit: + return + if event is not None and event.canvas is not self.canvas: + return + self._background = self.canvas.copy_from_bbox(self._resolve_bbox()) + self._draw_artists() + + def _on_resize(self, event): + if event is None or event.canvas is self.canvas: + self.invalidate() + + def add_artist(self, artist: martist.Artist): + """ + Add an artist to the managed update set. + + Returns + ------- + _BlitManager + This manager, to permit chained calls. + """ + if self._closed: + raise RuntimeError("Cannot add artists to a closed _BlitManager.") + if not isinstance(artist, martist.Artist): + raise TypeError( + f"Expected a matplotlib Artist, got {type(artist).__name__}." + ) + if artist.figure is not self.figure: + raise RuntimeError("The artist must belong to the manager's figure.") + if artist in self._artists: + return self + self._artists.append(artist) + self._animated[artist] = artist.get_animated() + if self._supports_blit: + artist.set_animated(True) + self.invalidate() + return self + + def remove_artist(self, artist: martist.Artist): + """ + Stop managing an artist and restore its original animated state. + + Returns + ------- + _BlitManager + This manager, to permit chained calls. + """ + if artist not in self._artists: + return self + self._artists.remove(artist) + artist.set_animated(self._animated.pop(artist)) + self.invalidate() + return self + + def invalidate(self): + """Discard the cached background before the next update.""" + self._background = None + + @contextmanager + def _save_context(self): + """Temporarily restore original artist states for a complete export.""" + if self._closed: + yield + return + self._suspended = True + current = {artist: artist.get_animated() for artist in self._artists} + try: + for artist in self._artists: + artist.set_animated(self._animated[artist]) + yield + finally: + for artist, animated in current.items(): + artist.set_animated(animated) + self._suspended = False + self.invalidate() + + def update(self, *, flush=False): + """ + Redraw the managed artists. + + Parameters + ---------- + flush : bool, default: False + Whether to immediately process pending GUI events after blitting. + + Returns + ------- + bool + ``True`` when the blitting fast path was used, otherwise ``False``. + """ + if self._closed: + raise RuntimeError("Cannot update a closed _BlitManager.") + if not self._supports_blit: + self.canvas.draw_idle() + if flush: + self.canvas.flush_events() + return False + + if self._background is None: + # The draw_event callback captures the static background and draws + # the animated artists before the backend presents the frame. + self.canvas.draw() + else: + bbox = self._resolve_bbox() + self.canvas.restore_region(self._background) + self._draw_artists() + self.canvas.blit(bbox) + if flush: + self.canvas.flush_events() + return True + + def close(self, *, redraw=True): + """ + Disconnect callbacks and restore the artists' animated states. + + Parameters + ---------- + redraw : bool, default: True + Whether to schedule a normal full redraw after restoring the artists. + """ + if self._closed: + return + self.canvas.mpl_disconnect(self._draw_cid) + self.canvas.mpl_disconnect(self._resize_cid) + for artist in tuple(self._artists): + artist.set_animated(self._animated[artist]) + self._artists.clear() + self._animated.clear() + self._background = None + self._closed = True + managers = getattr(self.figure, "_blit_managers", ()) + managers.discard(self) + if redraw: + self.canvas.draw_idle() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, traceback): + self.close() diff --git a/ultraplot/axes/_formatting.py b/ultraplot/axes/_formatting.py index 655fee713..f0489c1da 100644 --- a/ultraplot/axes/_formatting.py +++ b/ultraplot/axes/_formatting.py @@ -50,6 +50,22 @@ "labelweight": ("{axis}labelweight", "labelweight"), } +# These fields change only how existing geometry is painted. They do not change +# tick locations, text metrics, padding, or another layout input. Keep this list +# deliberately conservative: unknown fields must continue to invalidate layout. +_PAINT_ONLY_AXIS_STYLE_FIELDS = { + "color", + "linewidth", + "grid", + "gridminor", + "gridcolor", + "tickcolor", + "tickwidth", + "tickwidthratio", + "ticklabelcolor", + "labelcolor", +} + def _dedupe(items): return tuple(dict.fromkeys(items)) @@ -62,6 +78,13 @@ def _dedupe(items): if "{axis}" not in name ) +PAINT_ONLY_AXIS_FORMAT_KEYS = frozenset( + name.format(axis=axis) + for field in _PAINT_ONLY_AXIS_STYLE_FIELDS + for name in _AXIS_STYLE_FIELD_TEMPLATES[field] + for axis in ("x", "y") +) + CARTESIAN_PARENT_FILTER_KEYS = GENERIC_AXIS_FORMAT_KEYS + ( "label_kw", @@ -72,6 +95,26 @@ def _dedupe(items): ) +def axis_format_requires_layout(keys): + """ + Return whether explicit Cartesian formatting keys can affect layout. + + Unknown keys are treated as layout-affecting so new formatting options + remain correct until they are deliberately classified. + """ + keys = set(keys) + keys.difference_update( + { + "_explicit_format_keys", + "rc_kw", + "rc_mode", + "skip_axes", + "skip_figure", + } + ) + return bool(keys - PAINT_ONLY_AXIS_FORMAT_KEYS) + + def get_axis_style_fields(axis): """ Return the parameter names used to store explicit style overrides. diff --git a/ultraplot/axes/base.py b/ultraplot/axes/base.py index 735e99bee..fbae9408b 100644 --- a/ultraplot/axes/base.py +++ b/ultraplot/axes/base.py @@ -3335,7 +3335,7 @@ def format( ultraplot.gridspec.SubplotGrid.format ultraplot.config.Configurator.context """ - if self.figure is not None: + if self.figure is not None and getattr(self, "_format_layout_required", True): self.figure._layout_dirty = True skip_figure = kwargs.pop("skip_figure", False) # internal keyword arg params = _pop_params(kwargs, self.figure._format_signature) diff --git a/ultraplot/axes/cartesian.py b/ultraplot/axes/cartesian.py index 677f81989..bbe6a541c 100644 --- a/ultraplot/axes/cartesian.py +++ b/ultraplot/axes/cartesian.py @@ -32,6 +32,7 @@ from ..utils import units from ._formatting import ( CARTESIAN_PARENT_FILTER_KEYS, + axis_format_requires_layout, get_axis_style_fields, pop_axis_format_kwargs, ) @@ -1697,8 +1698,7 @@ def format( or `datetime.datetime` array as the x or y axis coordinate, the axis ticks and tick labels will be automatically formatted as dates. """ - explicit_format_keys = set(kwargs) - explicit_format_keys.update(kwargs.pop("_explicit_format_keys", ())) + explicit_format_keys = set(kwargs.pop("_explicit_format_keys", ())) signature_axis_kwargs, generic_axis_kwargs = pop_axis_format_kwargs( kwargs, self._format_signatures[CartesianAxes] ) @@ -1750,7 +1750,22 @@ def format( if aspect is not None: self.set_aspect(aspect) - super().format(rc_kw=rc_kw, rc_mode=rc_mode, **base_kwargs) + + # Base Axes.format() historically invalidated layout on every call. Let + # clearly paint-only Cartesian updates bypass that invalidation while + # remaining conservative for rc changes and unknown formatting keys. + sentinel = object() + previous = getattr(self, "_format_layout_required", sentinel) + self._format_layout_required = bool(rc_kw) or axis_format_requires_layout( + explicit_format_keys + ) + try: + super().format(rc_kw=rc_kw, rc_mode=rc_mode, **base_kwargs) + finally: + if previous is sentinel: + del self._format_layout_required + else: + self._format_layout_required = previous @docstring._snippet_manager def altx(self, **kwargs): diff --git a/ultraplot/figure.py b/ultraplot/figure.py index ed2955e99..1e862ca00 100644 --- a/ultraplot/figure.py +++ b/ultraplot/figure.py @@ -6,6 +6,7 @@ import functools import inspect import os +from contextlib import ExitStack try: from typing import Any, Iterable, List, Optional, Tuple, Union @@ -24,7 +25,7 @@ from typing_extensions import override from . import axes as paxes -from .axes._formatting import pop_axis_format_kwargs +from .axes._formatting import axis_format_requires_layout, pop_axis_format_kwargs from . import constructor from . import gridspec as pgridspec from . import legend as plegend @@ -1116,6 +1117,28 @@ def draw_without_rendering(self): if self.dpi != dpi: mfigure.Figure.set_dpi(self, dpi) + def _blit_manager(self, *artists, bbox=None): + """ + Return a manager for efficient updates of changing artists. + + Parameters + ---------- + *artists : `~matplotlib.artist.Artist` + Artists that will change between updates. + bbox : `~matplotlib.transforms.Bbox` or object with a ``bbox`` attribute, optional + Region to cache and blit. By default, the union of the artists' + axes bounding boxes is used. + + Returns + ------- + `~ultraplot._animation._BlitManager` + Manager that restores the cached static background and redraws only + the supplied artists. + """ + from ._animation import _BlitManager + + return _BlitManager(self.canvas, artists, bbox=bbox) + def _is_auto_share_mode(self, which: str) -> bool: """Return whether a given axis uses auto-share mode.""" if which not in ("x", "y"): @@ -3419,7 +3442,6 @@ def format( ultraplot.gridspec.SubplotGrid.format ultraplot.config.Configurator.context """ - self._layout_dirty = True # Initiate context block axs = axs or self._iter_subplots() skip_axes = kwargs.pop("skip_axes", False) # internal keyword arg @@ -3430,6 +3452,32 @@ def format( explicit_format_keys.update(signature_axis_kwargs) explicit_format_keys.update(generic_axis_kwargs) rc_kw, rc_mode = _pop_rc(kwargs) + figure_layout_values = ( + figtitle, + suptitle, + suptitle_kw, + llabels, + leftlabels, + leftlabels_kw, + rlabels, + rightlabels, + rightlabels_kw, + blabels, + bottomlabels, + bottomlabels_kw, + tlabels, + toplabels, + toplabels_kw, + rowlabels, + collabels, + includepanels, + ) + if ( + any(value is not None for value in figure_layout_values) + or bool(rc_kw) + or axis_format_requires_layout(explicit_format_keys) + ): + self._layout_dirty = True kwargs.update(signature_axis_kwargs) with rc.context(rc_kw, mode=rc_mode): # Update background patch @@ -3996,9 +4044,18 @@ def savefig(self, filename, **kwargs): # do not want to overwrite the matplotlib docstring. if isinstance(filename, str): filename = os.path.expanduser(filename) - # NOTE: this draw ensures that we are applying ultraplots layout adjustment. It is unclear what changed with ultraplot's history that makes this necessary, but it seems to cause no issues. Future devs, if unnecessary remove this line and test. - self.canvas.draw() - super().savefig(filename, **kwargs) + # Blitting marks managed artists as animated so full interactive draws can + # cache the static background. Temporarily restore their original states + # so savefig includes them in normal z-order. + managers = tuple(getattr(self, "_blit_managers", ())) + with ExitStack() as stack: + for manager in managers: + stack.enter_context(manager._save_context()) + # NOTE: this draw ensures that we are applying ultraplots layout + # adjustment. It is unclear what changed with ultraplot's history that + # makes this necessary, but it seems to cause no issues. + self.canvas.draw() + super().savefig(filename, **kwargs) @docstring._concatenate_inherited def set_canvas(self, canvas): diff --git a/ultraplot/tests/test_animation.py b/ultraplot/tests/test_animation.py index 00ee7c007..344dcec66 100644 --- a/ultraplot/tests/test_animation.py +++ b/ultraplot/tests/test_animation.py @@ -1,11 +1,15 @@ from unittest.mock import MagicMock +import io import matplotlib import numpy as np import pytest from matplotlib.animation import FuncAnimation +from matplotlib.backend_bases import FigureCanvasBase +from PIL import Image import ultraplot as uplt +from ultraplot._animation import _BlitManager def test_auto_layout_not_called_on_every_frame(): @@ -44,6 +48,180 @@ def test_draw_idle_skips_auto_layout_after_first_draw(): assert fig.auto_layout.call_count == 1 +@pytest.mark.parametrize( + "kwargs", + ( + {"xgrid": True}, + {"gridcolor": "red"}, + {"xtickcolor": "red"}, + {"xticklabelcolor": "red"}, + {"xlabelcolor": "red"}, + {"xlinewidth": 1.5}, + ), +) +def test_paint_only_format_does_not_invalidate_layout(kwargs): + """ + Paint-only Cartesian formatting should reuse the initialized layout. + """ + fig, ax = uplt.subplots() + fig.canvas.draw() + assert not fig._layout_dirty + + ax.format(**kwargs) + + assert not fig._layout_dirty + + +@pytest.mark.parametrize( + "kwargs", + ( + {"title": "Changed"}, + {"xlabel": "Changed"}, + {"xticklabelsize": 14}, + {"xlim": (0, 2)}, + {"xlocator": 0.5}, + ), +) +def test_geometry_format_invalidates_layout(kwargs): + """ + Text, tick geometry, limits, and locators must still invalidate layout. + """ + fig, ax = uplt.subplots() + fig.canvas.draw() + assert not fig._layout_dirty + + ax.format(**kwargs) + + assert fig._layout_dirty + + +def test_figure_paint_only_format_does_not_invalidate_layout(): + """ + Figure-wide routing should preserve paint-only invalidation decisions. + """ + fig, _ = uplt.subplots(ncols=2) + fig.canvas.draw() + + fig.format(xgrid=True, ygrid=True) + + assert not fig._layout_dirty + + +def test_blit_manager_updates_without_full_redraw(): + """ + After capturing a background, updates should use draw_artist and blit. + """ + fig, ax = uplt.subplots() + (line,) = ax.plot([0, 1, 2], [0, 1, 0]) + manager = fig._blit_manager(line) + assert isinstance(manager, _BlitManager) + assert manager.supports_blit + assert line.get_animated() + + manager.update() + assert manager._background is not None + + original_draw = fig.canvas.draw + original_blit = fig.canvas.blit + fig.canvas.draw = MagicMock(wraps=original_draw) + fig.canvas.blit = MagicMock(wraps=original_blit) + line.set_ydata([1, 0, 1]) + assert manager.update() + + fig.canvas.draw.assert_not_called() + fig.canvas.blit.assert_called_once() + manager.close(redraw=False) + assert not line.get_animated() + + +def test_blit_manager_matches_full_draw(): + """ + A blitted artist update should match a subsequent complete Agg draw. + """ + fig, ax = uplt.subplots() + (line,) = ax.plot([0.2, 1, 1.8], [0.2, 0.8, 0.2]) + ax.format(xlim=(0, 2), ylim=(0, 1)) + manager = fig._blit_manager(line) + manager.update() + + line.set_ydata([0.8, 0.2, 0.8]) + manager.update() + blitted = np.asarray(fig.canvas.buffer_rgba()).copy() + + manager.close(redraw=False) + fig.canvas.draw() + complete = np.asarray(fig.canvas.buffer_rgba()).copy() + + assert np.array_equal(blitted, complete) + + +def test_blit_manager_savefig_includes_managed_artist(): + """ + Saving should temporarily restore normal artist drawing and z-order. + """ + fig, ax = uplt.subplots() + (line,) = ax.plot([0.2, 1, 1.8], [0.2, 0.8, 0.2], color="red", lw=4) + ax.format(xlim=(0, 2), ylim=(0, 1)) + + baseline_buffer = io.BytesIO() + fig.savefig(baseline_buffer, format="png") + baseline_buffer.seek(0) + baseline = np.asarray(Image.open(baseline_buffer)).copy() + + manager = fig._blit_manager(line) + manager.update() + managed_buffer = io.BytesIO() + fig.savefig(managed_buffer, format="png") + managed_buffer.seek(0) + managed = np.asarray(Image.open(managed_buffer)).copy() + + assert np.array_equal(managed, baseline) + assert line.get_animated() + manager.close(redraw=False) + + +def test_blit_manager_invalidates_on_resize(): + """ + Resize events must discard pixel-sized background caches. + """ + fig, ax = uplt.subplots() + (line,) = ax.plot([0, 1], [0, 1]) + manager = fig._blit_manager(line) + manager.update() + assert manager._background is not None + + manager._on_resize(None) + + assert manager._background is None + manager.close(redraw=False) + + +def test_blit_manager_falls_back_for_unsupported_canvas(): + """ + Non-blitting canvases should retain normal artists and request idle draws. + """ + fig, ax = uplt.subplots() + (line,) = ax.plot([0, 1], [0, 1]) + canvas = FigureCanvasBase(fig) + canvas.draw_idle = MagicMock() + manager = _BlitManager(canvas, [line]) + + assert not manager.supports_blit + assert not line.get_animated() + assert not manager.update() + canvas.draw_idle.assert_called_once() + manager.close(redraw=False) + + +def test_blit_manager_rejects_artist_from_another_figure(): + fig1, _ = uplt.subplots() + _fig2, ax2 = uplt.subplots() + (line,) = ax2.plot([0, 1], [0, 1]) + + with pytest.raises(RuntimeError, match="must belong"): + fig1._blit_manager(line) + + def test_layout_array_no_crash(): """ Test that using layout_array with FuncAnimation does not crash. From 95278f5ea10de556123f8577b56da14599592c97 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Wed, 29 Jul 2026 21:59:03 +1000 Subject: [PATCH 2/6] perf: cache repeated layout measurements --- ultraplot/_layout.py | 601 ++++++++++++++++++++++++++++++ ultraplot/figure.py | 110 +++++- ultraplot/gridspec.py | 10 +- ultraplot/tests/test_animation.py | 217 +++++++++++ 4 files changed, 932 insertions(+), 6 deletions(-) create mode 100644 ultraplot/_layout.py diff --git a/ultraplot/_layout.py b/ultraplot/_layout.py new file mode 100644 index 000000000..f10154710 --- /dev/null +++ b/ultraplot/_layout.py @@ -0,0 +1,601 @@ +""" +Private helpers for reducing repeated layout work. + +The objects in this module are deliberately scoped to a single canvas draw. +They do not change matplotlib's persistent axis state or public API. +""" + +from __future__ import annotations + +from collections import OrderedDict +from dataclasses import dataclass + +import matplotlib.transforms as mtransforms +import numpy as np + +_MISSING = object() + + +def _is_internal_ticker(obj): + """Return whether a locator or formatter is safe for draw-local reuse.""" + module = type(obj).__module__ + return ( + module == "matplotlib" + or module.startswith("matplotlib.") + or module == "ultraplot.ticker" + ) + + +def _interval_key(values): + """Convert a numerical interval to an immutable exact cache key.""" + return tuple(np.asarray(values).reshape(-1).tolist()) + + +@dataclass(frozen=True) +class _AxisTickState: + """State that can affect ``Axis._update_ticks`` within one canvas draw.""" + + view_interval: tuple + data_interval: tuple + axes_size: tuple + dpi: float + scale: str + major_locator: int + major_formatter: int + minor_locator: int + minor_formatter: int + converter: int + units: int + + +@dataclass +class _AxisTickResult: + """Cached ticks and formatter locations for one axis state.""" + + ticks: list + major_locs: np.ndarray | tuple + minor_locs: np.ndarray | tuple + + +class _AxisTickCache: + """ + Cache repeated tick updates during one layout-and-render transaction. + + Tight bounding-box calculation and the final axes draw repeatedly call + ``Axis._update_ticks`` with identical state. The method runs locators, + formatters, tick positioning, and visibility filtering each time. This + manager replaces the method on individual axes for the duration of a + canvas draw and restores the original instance state afterwards. + + Custom third-party locators and formatters conservatively bypass the + cache because they may rely on repeated side effects. + """ + + _MAX_STATES_PER_AXIS = 4 + + def __init__(self, figure): + self.figure = figure + self._cache = {} + self._patched = {} + self.hits = 0 + self.misses = 0 + self.bypasses = 0 + self.evictions = 0 + + def __enter__(self): + self.figure._axis_tick_cache = self + self.refresh() + return self + + def __exit__(self, *args): + for axis, previous in self._patched.items(): + if previous is _MISSING: + axis.__dict__.pop("_update_ticks", None) + else: + axis.__dict__["_update_ticks"] = previous + self._patched.clear() + self.figure.__dict__.pop("_axis_tick_cache", None) + self.figure._last_axis_tick_cache_stats = { + "hits": self.hits, + "misses": self.misses, + "bypasses": self.bypasses, + "evictions": self.evictions, + } + + def refresh(self): + """Patch axes added while queued guides and panels are materialized.""" + for axes in self.figure.axes: + axis_map = getattr(axes, "_axis_map", {}) + for axis in axis_map.values(): + if axis is not None and axis not in self._patched: + self._patch(axis) + + def _patch(self, axis): + previous = axis.__dict__.get("_update_ticks", _MISSING) + original = axis._update_ticks + self._patched[axis] = previous + + def _cached_update_ticks(): + if not self._is_cacheable(axis): + self.bypasses += 1 + return original() + state = self._get_state(axis) + states = self._cache.setdefault(axis, OrderedDict()) + result = states.pop(state, None) + if result is not None: + self.hits += 1 + states[state] = result + self._restore_formatter_locs(axis, result) + return result.ticks + self.misses += 1 + ticks = original() + result = _AxisTickResult( + ticks=ticks, + major_locs=self._copy_formatter_locs(axis.major.formatter), + minor_locs=self._copy_formatter_locs(axis.minor.formatter), + ) + states[state] = result + if len(states) > self._MAX_STATES_PER_AXIS: + states.popitem(last=False) + self.evictions += 1 + return ticks + + axis._update_ticks = _cached_update_ticks + + @staticmethod + def _is_cacheable(axis): + if getattr(axis.axes, "_name", None) != "cartesian": + return False + ticker_objects = ( + axis.major.locator, + axis.major.formatter, + axis.minor.locator, + axis.minor.formatter, + ) + return all(_is_internal_ticker(obj) for obj in ticker_objects) + + @staticmethod + def _get_state(axis): + axes = axis.axes + figure = axes.get_figure(root=True) + get_converter = getattr(axis, "get_converter", None) + converter = get_converter() if get_converter is not None else None + get_units = getattr(axis, "get_units", None) + units = get_units() if get_units is not None else None + return _AxisTickState( + view_interval=_interval_key(axis.get_view_interval()), + data_interval=_interval_key(axis.get_data_interval()), + axes_size=(axes.bbox.width, axes.bbox.height), + dpi=float(figure.dpi), + scale=axis.get_scale(), + major_locator=id(axis.major.locator), + major_formatter=id(axis.major.formatter), + minor_locator=id(axis.minor.locator), + minor_formatter=id(axis.minor.formatter), + converter=id(converter), + units=id(units), + ) + + @staticmethod + def _copy_formatter_locs(formatter): + locs = getattr(formatter, "locs", ()) + try: + return np.array(locs, copy=True) + except Exception: + return tuple(locs) + + @staticmethod + def _restore_formatter_locs(axis, result): + for formatter, locs in ( + (axis.major.formatter, result.major_locs), + (axis.minor.formatter, result.minor_locs), + ): + setter = getattr(formatter, "set_locs", None) + if setter is not None: + setter(locs) + + +@dataclass(frozen=True) +class _AxesExtentState: + """Geometry that can alter outsets relative to an axes rectangle.""" + + bbox_size: tuple + bbox_position: tuple + dpi: float + axis_states: tuple + decorations: tuple + subset_titles: tuple + + +@dataclass +class _AxesExtentRecord: + """One relative tight-bbox measurement.""" + + version: int + state: _AxesExtentState + outsets: tuple + + +class _ExtentReductionTree: + """Small mutable min/max tree for the union of axes extents.""" + + def __init__(self, axes): + self.axes = tuple(axes) + self._indices = {axis: index for index, axis in enumerate(self.axes)} + capacity = 1 + while capacity < max(1, len(self.axes)): + capacity *= 2 + self._capacity = capacity + self._values = np.empty((2 * capacity, 4), dtype=float) + self.clear() + + def clear(self): + self._values[:, :2] = np.inf + self._values[:, 2:] = -np.inf + + def update(self, axis, bbox): + index = self._indices.get(axis) + if index is None: + return + node = self._capacity + index + if bbox is None: + self._values[node] = (np.inf, np.inf, -np.inf, -np.inf) + else: + self._values[node] = (bbox.xmin, bbox.ymin, bbox.xmax, bbox.ymax) + node //= 2 + while node: + left = self._values[2 * node] + right = self._values[2 * node + 1] + self._values[node, :2] = np.minimum(left[:2], right[:2]) + self._values[node, 2:] = np.maximum(left[2:], right[2:]) + node //= 2 + + def get_bbox(self): + xmin, ymin, xmax, ymax = self._values[1] + if not np.all(np.isfinite((xmin, ymin, xmax, ymax))): + return None + return mtransforms.Bbox.from_extents(xmin, ymin, xmax, ymax) + + +class _LayoutExtentStore: + """ + Persist relative axes outsets and dependency versions between layouts. + + Absolute axes positions are solver outputs. Tick labels, axis labels, and + titles are better represented as four overhangs around those positions. + Standard Cartesian axes can therefore move without repeating renderer text + measurements. Position-sensitive axes and extra artists automatically add + the absolute origin to their state key. + """ + + def __init__(self, figure): + self.figure = figure + self._records = {} + self._versions = {} + self._tree = None + self._topology = {} + self._active = False + self.hits = 0 + self.misses = 0 + + def __enter__(self): + self._active = True + self.hits = 0 + self.misses = 0 + axes = self.refresh() + for axis in axes: + if axis.stale: + self._versions[axis] += 1 + return self + + def refresh(self): + """Synchronize axes added by queued guide and panel creation.""" + axes = tuple(self.figure.axes) + if self._tree is None or self._tree.axes != axes: + self._tree = _ExtentReductionTree(axes) + self._topology.clear() + current = set(axes) + self._records = { + key: value for key, value in self._records.items() if key[0] in current + } + self._versions = { + axis: version + for axis, version in self._versions.items() + if axis in current + } + for axis in axes: + if axis not in self._versions: + self._versions[axis] = 1 + return axes + + def __exit__(self, *args): + self._rebase_records() + self._active = False + self.figure._last_layout_extent_stats = { + "hits": self.hits, + "misses": self.misses, + } + + def get_tightbbox( + self, + axes, + renderer, + *, + include_subset_titles=True, + use_cache=True, + ): + """Return an exact or reconstructed tight bbox in display units.""" + if not self._active or not use_cache or not self._is_cacheable_axes(axes): + bbox = self._measure_tightbbox(axes, renderer, include_subset_titles) + self.update_union(axes, bbox) + return bbox + + version = self._versions.setdefault(axes, 1) + state = self._get_state(axes, include_subset_titles) + cache_key = (axes, bool(include_subset_titles)) + record = self._records.get(cache_key) + if record is not None and record.version == version and record.state == state: + self.hits += 1 + bbox = self._bbox_from_outsets(axes.bbox, record.outsets) + axes._tight_bbox = bbox + else: + self.misses += 1 + bbox = self._measure_tightbbox(axes, renderer, include_subset_titles) + if bbox is not None: + self._records[cache_key] = _AxesExtentRecord( + version=version, + state=self._get_state(axes, include_subset_titles), + outsets=self._get_outsets(axes.bbox, bbox), + ) + self.update_union(axes, bbox) + return bbox + + def clear_union(self): + if self._tree is not None: + self._tree.clear() + + def update_union(self, axes, bbox): + if self._tree is not None: + self._tree.update(axes, bbox) + + def get_union(self): + return None if self._tree is None else self._tree.get_bbox() + + def get_subplot_ranges(self, axes, along, across): + """Return compact topology arrays used by boundary reductions.""" + key = (tuple(axes), along, across) + cached = self._topology.get(key) + if cached is not None: + return cached + along_ranges = np.asarray( + [axis._range_subplotspec(along) for axis in axes], dtype=int + ) + across_ranges = np.asarray( + [axis._range_subplotspec(across) for axis in axes], dtype=int + ) + result = (along_ranges, across_ranges) + self._topology[key] = result + return result + + def invalidate_topology(self): + """Discard structural row/column range lookups.""" + self._topology.clear() + + def _get_state(self, axes, include_subset_titles=True): + bbox = axes.bbox + position_sensitive = self._is_position_sensitive(axes) + axis_states = [] + for axis in getattr(axes, "_axis_map", {}).values(): + if axis is None: + continue + get_converter = getattr(axis, "get_converter", None) + converter = get_converter() if get_converter is not None else None + get_units = getattr(axis, "get_units", None) + units = get_units() if get_units is not None else None + axis_states.append( + ( + _interval_key(axis.get_view_interval()), + _interval_key(axis.get_data_interval()), + axis.get_scale(), + id(axis.major.locator), + id(axis.major.formatter), + id(axis.minor.locator), + id(axis.minor.formatter), + id(converter), + id(units), + _interval_key(getattr(axis.major.formatter, "locs", ())), + _interval_key(getattr(axis.minor.formatter, "locs", ())), + ) + ) + return _AxesExtentState( + bbox_size=(bbox.width, bbox.height), + bbox_position=(bbox.x0, bbox.y0) if position_sensitive else (), + dpi=float(axes.get_figure(root=True).dpi), + axis_states=tuple(axis_states), + decorations=self._get_decoration_state(axes), + subset_titles=self._get_subset_title_state(axes, include_subset_titles), + ) + + def _rebase_records(self): + """ + Rebase reusable outsets onto final post-render axes dimensions. + + UltraLayout may make a small solver adjustment after measuring an axes. + The final render updates locator/formatter locations for that geometry. + If those locations and every non-size dependency are unchanged, the + relative outsets remain valid for the next layout transaction. + """ + for (axes, include_subset_titles), record in self._records.items(): + if axes not in self._versions: + continue + if record.version != self._versions[axes]: + continue + if self._is_position_sensitive(axes): + continue + state = self._get_state(axes, include_subset_titles) + if ( + record.state.dpi == state.dpi + and record.state.axis_states == state.axis_states + and record.state.decorations == state.decorations + and record.state.subset_titles == state.subset_titles + ): + record.state = state + + @staticmethod + def _get_decoration_state(axes): + texts = [] + core_texts = ( + getattr(axes, "title", None), + getattr(axes, "_left_title", None), + getattr(axes, "_right_title", None), + getattr(getattr(axes, "xaxis", None), "label", None), + getattr(getattr(axes, "yaxis", None), "label", None), + ) + extra_texts = tuple( + text for text in getattr(axes, "texts", ()) if text.get_text() + ) + for text in (*core_texts, *extra_texts): + if text is None: + continue + texts.append( + ( + id(text), + text.get_visible(), + text.get_in_layout(), + text.get_text(), + text.get_rotation(), + hash(text.get_fontproperties()), + ) + ) + legend = getattr(axes, "legend_", None) + if legend is None: + legend_state = () + else: + legend_state = ( + id(legend), + legend.get_visible(), + legend.get_in_layout(), + id(getattr(legend, "_bbox_to_anchor", None)), + tuple(text.get_text() for text in legend.get_texts()), + ) + extra_ids = tuple( + id(artist) + for artists in ( + getattr(axes, "artists", ()), + getattr(axes, "tables", ()), + ) + for artist in artists + if artist.get_visible() and artist.get_in_layout() + ) + return tuple(texts), legend_state, extra_ids + + def _get_subset_title_state(self, axes, include_subset_titles): + if not include_subset_titles: + return () + figure = axes.get_figure(root=True) + groups = getattr(figure, "_subset_title_dict", {}) + state = [] + parent = getattr(axes, "_panel_parent", None) or axes + for group in groups.values(): + group_axes = tuple( + getattr(item, "_panel_parent", None) or item + for item in group["axes"] + if item is not None and item.figure is figure + ) + if parent not in group_axes: + continue + artist = group["artist"] + state.append( + ( + tuple(id(item) for item in group_axes), + artist.get_visible(), + artist.get_text(), + artist.get_position(), + artist.get_ha(), + artist.get_va(), + artist.get_rotation(), + hash(artist.get_fontproperties()), + group.get("pad"), + group.get("y"), + ) + ) + return tuple(state) + + @staticmethod + def _is_position_sensitive(axes): + if getattr(axes, "_name", None) != "cartesian": + return True + legend = getattr(axes, "legend_", None) + if legend is not None and getattr(legend, "_bbox_to_anchor", None) is not None: + return True + collections = ( + getattr(axes, "artists", ()), + getattr(axes, "tables", ()), + ) + if any( + artist.get_visible() and artist.get_in_layout() + for artists in collections + for artist in artists + ): + return True + relative_texts = set(getattr(axes, "_title_dict", {}).values()) + for text in getattr(axes, "texts", ()): + if not ( + text.get_visible() and text.get_in_layout() and bool(text.get_text()) + ): + continue + if text in relative_texts: + continue + transform = text.get_transform() + if not ( + transform.contains_branch(axes.transAxes) + or transform.contains_branch(axes.transData) + ): + return True + return False + + @staticmethod + def _is_cacheable_axes(axes): + if getattr(axes, "_name", None) != "cartesian": + return False + return all( + axis is None + or all( + _is_internal_ticker(obj) + for obj in ( + axis.major.locator, + axis.major.formatter, + axis.minor.locator, + axis.minor.formatter, + ) + ) + for axis in getattr(axes, "_axis_map", {}).values() + ) + + @staticmethod + def _get_outsets(axes_bbox, tight_bbox): + return ( + axes_bbox.xmin - tight_bbox.xmin, + tight_bbox.xmax - axes_bbox.xmax, + axes_bbox.ymin - tight_bbox.ymin, + tight_bbox.ymax - axes_bbox.ymax, + ) + + @staticmethod + def _bbox_from_outsets(axes_bbox, outsets): + left, right, bottom, top = outsets + return mtransforms.Bbox.from_extents( + axes_bbox.xmin - left, + axes_bbox.ymin - bottom, + axes_bbox.xmax + right, + axes_bbox.ymax + top, + ) + + @staticmethod + def _measure_tightbbox(axes, renderer, include_subset_titles): + try: + return axes.get_tightbbox( + renderer, include_subset_titles=include_subset_titles + ) + except TypeError: + return axes.get_tightbbox(renderer) diff --git a/ultraplot/figure.py b/ultraplot/figure.py index 1e862ca00..b12bdf3e9 100644 --- a/ultraplot/figure.py +++ b/ultraplot/figure.py @@ -42,6 +42,7 @@ labels, warnings, ) +from ._layout import _AxisTickCache, _LayoutExtentStore from ._subplots import SubplotManager from .utils import _Crawler, units @@ -618,6 +619,7 @@ def _canvas_preprocess(self, *args, **kwargs): skip_autolayout = getattr(fig, "_skip_autolayout", False) layout_dirty = getattr(fig, "_layout_dirty", False) + needs_layout = not getattr(fig, "_layout_initialized", False) or layout_dirty saving_frame_count = getattr(fig, "_saving_frame_count", 0) lock_tight_during_save = ( getattr(self, "_is_saving", False) @@ -639,7 +641,19 @@ def _canvas_preprocess(self, *args, **kwargs): ctx1 = fig._context_adjusting(cache=cache) ctx2 = fig._context_authorized() # skip backend set_constrained_layout() ctx3 = rc.context(fig._render_context) # draw with figure-specific setting - with ctx1, ctx2, ctx3: + ctx4 = ( + _AxisTickCache(fig) + if needs_layout and not getattr(fig, "_disable_axis_tick_cache", False) + else context._empty_context() + ) + if needs_layout and not getattr(fig, "_disable_layout_extent_cache", False): + store = getattr(fig, "_layout_extent_store", None) + if store is None: + store = fig._layout_extent_store = _LayoutExtentStore(fig) + ctx5 = store + else: + ctx5 = context._empty_context() + with ctx1, ctx2, ctx3, ctx4, ctx5: needs_post_layout = False if not fig._layout_initialized or layout_dirty: fig.auto_layout(tight=False if lock_tight_during_save else None) @@ -673,6 +687,9 @@ def wrapper(self, *args, **kwargs): result = func(self, *args, **kwargs) if hasattr(self, "_cached_border_axes"): delattr(self, "_cached_border_axes") + store = getattr(self, "_layout_extent_store", None) + if store is not None: + store.invalidate_topology() return result return wrapper @@ -1089,6 +1106,7 @@ def clear(self, keep_observers=False): self._panel_dict = {"left": [], "right": [], "bottom": [], "top": []} self._layout_initialized = False self._layout_dirty = True + self.__dict__.pop("_layout_extent_store", None) self._init_super_labels() @override @@ -1971,8 +1989,11 @@ def _get_offset_coord( if label.get_text() and not label.get_text().strip(): label.set_visible(False) if isinstance(obj, paxes.Axes): - bbox = obj.get_tightbbox( - renderer, include_subset_titles=include_subset_titles + bbox = self._get_layout_axes_bbox( + obj, + renderer, + include_subset_titles=include_subset_titles, + use_cache=not exclude_spanning_axis_labels, ) else: bbox = obj.get_tightbbox(renderer) # cannot use cached bbox @@ -1990,6 +2011,83 @@ def _get_offset_coord( pad = pad / width if side in ("left", "right") else pad / height return min(cs) - pad if side in ("left", "bottom") else max(cs) + pad + def _get_layout_axes_bbox( + self, + axes, + renderer, + *, + include_subset_titles=True, + use_cache=True, + ): + """Return an axes bbox using the active relative-outset store.""" + store = getattr(self, "_layout_extent_store", None) + if store is not None and store._active: + return store.get_tightbbox( + axes, + renderer, + include_subset_titles=include_subset_titles, + use_cache=use_cache, + ) + try: + return axes.get_tightbbox( + renderer, include_subset_titles=include_subset_titles + ) + except TypeError: + return axes.get_tightbbox(renderer) + + def _get_layout_tightbbox(self, renderer): + """ + Return the figure tight bbox while reusing relative axes outsets. + + This mirrors matplotlib's ``Figure.get_tightbbox`` but routes axes + measurements through the active dependency store and reduction tree. + """ + artists = [ + artist + for artist in self.get_children() + if ( + artist not in self.axes + and artist.get_visible() + and artist.get_in_layout() + ) + ] + bboxes = [] + for artist in artists: + bbox = artist.get_tightbbox(renderer) + if bbox is not None: + bboxes.append(bbox) + + store = getattr(self, "_layout_extent_store", None) + if store is not None and store._active: + store.clear_union() + for axes in self.axes: + if not axes.get_visible(): + if store is not None and store._active: + store.update_union(axes, None) + continue + bbox = self._get_layout_axes_bbox(axes, renderer) + if store is None or not store._active: + if bbox is not None: + bboxes.append(bbox) + if store is not None and store._active: + bbox = store.get_union() + if bbox is not None: + bboxes.append(bbox) + + bboxes = [ + bbox + for bbox in bboxes + if ( + np.isfinite(bbox.width) + and np.isfinite(bbox.height) + and (bbox.width != 0 or bbox.height != 0) + ) + ] + if not bboxes: + return self.bbox_inches + bbox = mtransforms.Bbox.union(bboxes) + return mtransforms.TransformedBbox(bbox, self.dpi_scale_trans.inverted()) + def _get_renderer(self): """ Get a renderer at all costs. See matplotlib's tight_layout.py. @@ -3366,6 +3464,12 @@ def _align_content(): # noqa: E306 # WARNING: Tried to avoid two figure resizes but made # subsequent tight layout really weird. Have to resize twice. _draw_content() + manager = getattr(self, "_axis_tick_cache", None) + if manager is not None: + manager.refresh() + store = getattr(self, "_layout_extent_store", None) + if store is not None and store._active: + store.refresh() if not gs: return if aspect: diff --git a/ultraplot/gridspec.py b/ultraplot/gridspec.py index eb7fdd59f..eed1e481b 100644 --- a/ultraplot/gridspec.py +++ b/ultraplot/gridspec.py @@ -1117,8 +1117,12 @@ def _get_tight_space(self, w): # Iterate along each row or column space axs = tuple(fig._iter_axes(hidden=True, children=False)) space = list(space) # a copy - ralong = np.array([ax._range_subplotspec(x) for ax in axs]) - racross = np.array([ax._range_subplotspec(y) for ax in axs]) + store = getattr(fig, "_layout_extent_store", None) + if store is not None and store._active: + ralong, racross = store.get_subplot_ranges(axs, x, y) + else: + ralong = np.array([ax._range_subplotspec(x) for ax in axs]) + racross = np.array([ax._range_subplotspec(y) for ax in axs]) for i, (s, p) in enumerate(zip(space, pad)): # Find axes that abutt aginst this row or column space groups = [] @@ -1249,7 +1253,7 @@ def _auto_layout_tight(self, renderer): # computed bounding boxes used by _range_tightbbox below. pad = self._outerpad obox = fig.bbox_inches # original bbox - bbox = fig.get_tightbbox(renderer) + bbox = fig._get_layout_tightbbox(renderer) # Calculate new figure margins # NOTE: Negative spaces are common where entire rows/columns of gridspec diff --git a/ultraplot/tests/test_animation.py b/ultraplot/tests/test_animation.py index 344dcec66..7491fa820 100644 --- a/ultraplot/tests/test_animation.py +++ b/ultraplot/tests/test_animation.py @@ -1,15 +1,18 @@ from unittest.mock import MagicMock +import datetime import io import matplotlib import numpy as np import pytest +from matplotlib import ticker as mticker from matplotlib.animation import FuncAnimation from matplotlib.backend_bases import FigureCanvasBase from PIL import Image import ultraplot as uplt from ultraplot._animation import _BlitManager +from ultraplot._layout import _AxisTickCache def test_auto_layout_not_called_on_every_frame(): @@ -48,6 +51,220 @@ def test_draw_idle_skips_auto_layout_after_first_draw(): assert fig.auto_layout.call_count == 1 +def test_initial_draw_reuses_tick_updates(): + """ + Layout and render phases should share identical tick computations. + """ + fig, axs = uplt.subplots(nrows=2, ncols=2, share=False) + for ax in axs: + ax.plot([0, 1, 2], [0, 1, 0]) + axs.format(xlabel="Coordinate", ylabel="Response", suptitle="Tick cache") + + fig.canvas.draw() + + stats = fig._last_axis_tick_cache_stats + assert stats["hits"] > 0 + assert stats["misses"] > 0 + assert stats["bypasses"] == 0 + assert stats["evictions"] == 0 + assert "_axis_tick_cache" not in fig.__dict__ + for ax in fig.axes: + for axis in ax._axis_map.values(): + assert "_update_ticks" not in axis.__dict__ + + +def test_tick_cache_preserves_rendered_pixels(): + """ + Caching tick updates must produce the exact uncached raster output. + """ + + def _draw(disable): + fig, axs = uplt.subplots(nrows=2, ncols=2, share=False) + x = np.linspace(0, 2 * np.pi, 100) + for index, ax in enumerate(axs): + ax.plot(x, np.sin(x + index)) + axs.format( + xlabel="Coordinate", + ylabel="Response", + suptitle="Pixel comparison", + grid=True, + ) + fig._disable_axis_tick_cache = disable + fig._disable_layout_extent_cache = disable + fig.canvas.draw() + return np.asarray(fig.canvas.buffer_rgba()).copy() + + expected = _draw(True) + actual = _draw(False) + + assert np.array_equal(actual, expected) + + +@pytest.mark.parametrize("kind", ("log", "date", "categorical", "shared")) +def test_layout_caches_preserve_specialized_tick_pixels(kind): + """ + Built-in non-linear, unit-aware, categorical, and shared ticks stay exact. + """ + + def _draw(disable): + if kind == "shared": + fig, axs = uplt.subplots(ncols=2, share=True) + else: + fig, ax = uplt.subplots() + axs = [ax] + if kind == "log": + axs[0].semilogx([1, 10, 100, 1000], [0, 1, 0, 1]) + elif kind == "date": + dates = [datetime.date(2025, 1, day) for day in range(1, 8)] + axs[0].plot(dates, np.arange(len(dates))) + elif kind == "categorical": + axs[0].bar(["alpha", "beta", "gamma"], [1, 3, 2]) + else: + for index, ax in enumerate(axs): + ax.plot([0, 1, 2], np.asarray([0, 1, 0]) + index) + fig._disable_axis_tick_cache = disable + fig._disable_layout_extent_cache = disable + fig.canvas.draw() + return np.asarray(fig.canvas.buffer_rgba()).copy() + + expected = _draw(True) + actual = _draw(False) + + assert np.array_equal(actual, expected) + + +def test_layout_caches_bypass_non_cartesian_axes(): + """ + Projection-specific tick and bbox implementations use their native paths. + """ + fig, ax = uplt.subplots(proj="polar") + ax.plot(np.linspace(0, 2 * np.pi, 20), np.linspace(0, 1, 20)) + + fig.canvas.draw() + + stats = fig._last_axis_tick_cache_stats + assert stats["bypasses"] > 0 + assert fig._last_layout_extent_stats == {"hits": 0, "misses": 0} + + +def test_tick_cache_invalidates_changed_axis_geometry(): + """ + Limits and pixel geometry are part of the draw-local cache key. + """ + fig, ax = uplt.subplots() + axis = ax.xaxis + with _AxisTickCache(fig) as cache: + axis._update_ticks() + axis._update_ticks() + assert (cache.hits, cache.misses) == (1, 1) + + ax.set_xlim(0, 10) + axis._update_ticks() + assert (cache.hits, cache.misses) == (1, 2) + + position = ax.get_position() + ax.set_position( + [position.x0, position.y0, 0.5 * position.width, position.height] + ) + axis._update_ticks() + assert (cache.hits, cache.misses) == (1, 3) + + ax.set_position(position) + axis._update_ticks() + assert (cache.hits, cache.misses) == (2, 3) + + +def test_tick_cache_bypasses_custom_tickers(): + """ + Unknown locator and formatter implementations may have stateful calls. + """ + + class CustomFormatter(mticker.Formatter): + def __call__(self, value, pos=None): + return f"{value:g}" + + fig, ax = uplt.subplots() + ax.xaxis.set_major_formatter(CustomFormatter()) + with _AxisTickCache(fig) as cache: + ax.xaxis._update_ticks() + ax.xaxis._update_ticks() + + assert cache.hits == 0 + assert cache.misses == 0 + assert cache.bypasses == 2 + + +def test_tick_cache_lru_is_bounded(): + """ + Cycling through many geometries must not grow the draw-local cache. + """ + fig, ax = uplt.subplots() + with _AxisTickCache(fig) as cache: + for stop in range(2, 9): + ax.set_xlim(0, stop) + ax.xaxis._update_ticks() + states = cache._cache[ax.xaxis] + + assert len(states) == cache._MAX_STATES_PER_AXIS + assert cache.evictions == 3 + + +def test_layout_extent_store_reuses_unmodified_axes(): + """ + A one-axes title edit should only remeasure that axes. + """ + fig, axs = uplt.subplots(nrows=2, ncols=2, share=False) + axs.format(abc="a.") + axs[1].plot([0, 1], [0, 1], label="Default axes legend") + axs[1].legend() + fig.canvas.draw() + + axs[0].format(title="Changed") + fig.canvas.draw() + + stats = fig._last_layout_extent_stats + assert stats == {"hits": 3, "misses": 1} + + +def test_layout_extent_store_preserves_incremental_pixels(): + """ + Relative outsets must exactly match an uncached geometry update. + """ + + def _draw(disable): + fig, axs = uplt.subplots(nrows=2, ncols=2, share=False) + for ax in axs: + ax.plot([0, 1, 2], [0, 1, 0]) + axs.format(abc="a.") + axs[1].plot([0, 1], [1, 0], label="Legend entry") + axs[1].legend() + fig._disable_layout_extent_cache = disable + fig.canvas.draw() + axs[0].format(title="A longer changed title") + fig.canvas.draw() + return np.asarray(fig.canvas.buffer_rgba()).copy() + + expected = _draw(True) + actual = _draw(False) + + assert np.array_equal(actual, expected) + + +def test_layout_extent_store_tracks_subset_title_changes(): + """ + Shared subset-title state participates in axes extent cache keys. + """ + fig, axs = uplt.subplots(nrows=2, ncols=2, share=False) + axs[0, :].format(title="First") + fig.canvas.draw() + + axs[0, :].format(title="A substantially longer shared title") + fig.canvas.draw() + + stats = fig._last_layout_extent_stats + assert stats["misses"] >= 2 + + @pytest.mark.parametrize( "kwargs", ( From 18d4e4b9bef49e5c7e49ab361de06632f06bde21 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Wed, 29 Jul 2026 22:08:25 +1000 Subject: [PATCH 3/6] perf: cache ticks for layout child axes --- ultraplot/_layout.py | 2 +- ultraplot/tests/test_animation.py | 23 +++++++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/ultraplot/_layout.py b/ultraplot/_layout.py index f10154710..ea1536dd7 100644 --- a/ultraplot/_layout.py +++ b/ultraplot/_layout.py @@ -104,7 +104,7 @@ def __exit__(self, *args): def refresh(self): """Patch axes added while queued guides and panels are materialized.""" - for axes in self.figure.axes: + for axes in self.figure._iter_axes(hidden=True, children=True): axis_map = getattr(axes, "_axis_map", {}) for axis in axis_map.values(): if axis is not None and axis not in self._patched: diff --git a/ultraplot/tests/test_animation.py b/ultraplot/tests/test_animation.py index 7491fa820..ddea1d418 100644 --- a/ultraplot/tests/test_animation.py +++ b/ultraplot/tests/test_animation.py @@ -209,6 +209,29 @@ def test_tick_cache_lru_is_bounded(): assert cache.evictions == 3 +def test_tick_cache_includes_child_axes(): + """ + Colorbar child axes created during layout should share tick computations. + """ + fig, axs = uplt.subplots() + ax = axs[0] + ax.colorbar("magma", loc="r") + ax._add_queued_guides() + child_axes = [ + child + for child in fig._iter_axes(hidden=True, children=True) + if child not in fig.axes + ] + assert child_axes + + with _AxisTickCache(fig) as cache: + child_axes[0].xaxis._update_ticks() + child_axes[0].xaxis._update_ticks() + assert (cache.hits, cache.misses) == (1, 1) + + assert "_update_ticks" not in child_axes[0].xaxis.__dict__ + + def test_layout_extent_store_reuses_unmodified_axes(): """ A one-axes title edit should only remeasure that axes. From 8c03deb98faaed99b5a92e46807dde5bcb9c12f4 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Thu, 30 Jul 2026 06:57:08 +1000 Subject: [PATCH 4/6] refactor: simplify figure layout request check --- ultraplot/figure.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/ultraplot/figure.py b/ultraplot/figure.py index b12bdf3e9..17d0680b3 100644 --- a/ultraplot/figure.py +++ b/ultraplot/figure.py @@ -51,6 +51,11 @@ ] +def _any_not_none(*values): + """Return whether at least one value is not ``None``.""" + return any(value is not None for value in values) + + # Preset figure widths or sizes based on academic journal recommendations # NOTE: Please feel free to add to this! JOURNAL_SIZES = { @@ -3556,7 +3561,7 @@ def format( explicit_format_keys.update(signature_axis_kwargs) explicit_format_keys.update(generic_axis_kwargs) rc_kw, rc_mode = _pop_rc(kwargs) - figure_layout_values = ( + figure_layout_requested = _any_not_none( figtitle, suptitle, suptitle_kw, @@ -3577,7 +3582,7 @@ def format( includepanels, ) if ( - any(value is not None for value in figure_layout_values) + figure_layout_requested or bool(rc_kw) or axis_format_requires_layout(explicit_format_keys) ): From abaf1a779cc4006f5c795ff25ba401ec126b9fab Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Thu, 30 Jul 2026 07:11:38 +1000 Subject: [PATCH 5/6] fix: support layout caches on matplotlib 3.9 --- ultraplot/_layout.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/ultraplot/_layout.py b/ultraplot/_layout.py index ea1536dd7..0ffe300fe 100644 --- a/ultraplot/_layout.py +++ b/ultraplot/_layout.py @@ -154,10 +154,8 @@ def _is_cacheable(axis): ) return all(_is_internal_ticker(obj) for obj in ticker_objects) - @staticmethod - def _get_state(axis): + def _get_state(self, axis): axes = axis.axes - figure = axes.get_figure(root=True) get_converter = getattr(axis, "get_converter", None) converter = get_converter() if get_converter is not None else None get_units = getattr(axis, "get_units", None) @@ -166,7 +164,7 @@ def _get_state(axis): view_interval=_interval_key(axis.get_view_interval()), data_interval=_interval_key(axis.get_data_interval()), axes_size=(axes.bbox.width, axes.bbox.height), - dpi=float(figure.dpi), + dpi=float(self.figure.dpi), scale=axis.get_scale(), major_locator=id(axis.major.locator), major_formatter=id(axis.major.formatter), @@ -410,7 +408,7 @@ def _get_state(self, axes, include_subset_titles=True): return _AxesExtentState( bbox_size=(bbox.width, bbox.height), bbox_position=(bbox.x0, bbox.y0) if position_sensitive else (), - dpi=float(axes.get_figure(root=True).dpi), + dpi=float(self.figure.dpi), axis_states=tuple(axis_states), decorations=self._get_decoration_state(axes), subset_titles=self._get_subset_title_state(axes, include_subset_titles), @@ -492,7 +490,7 @@ def _get_decoration_state(axes): def _get_subset_title_state(self, axes, include_subset_titles): if not include_subset_titles: return () - figure = axes.get_figure(root=True) + figure = self.figure groups = getattr(figure, "_subset_title_dict", {}) state = [] parent = getattr(axes, "_panel_parent", None) or axes From 53299ddb9df553b151c3c3d19ee6a4b8618305c7 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Thu, 30 Jul 2026 08:36:36 +1000 Subject: [PATCH 6/6] refactor: simplify layout cache lifecycle --- ultraplot/_layout.py | 145 +++++++++++++----------------- ultraplot/_subplots.py | 2 +- ultraplot/axes/base.py | 6 +- ultraplot/figure.py | 65 ++++++-------- ultraplot/gridspec.py | 8 +- ultraplot/tests/test_animation.py | 42 ++++++++- 6 files changed, 137 insertions(+), 131 deletions(-) diff --git a/ultraplot/_layout.py b/ultraplot/_layout.py index 0ffe300fe..792dae3be 100644 --- a/ultraplot/_layout.py +++ b/ultraplot/_layout.py @@ -1,13 +1,20 @@ """ Private helpers for reducing repeated layout work. -The objects in this module are deliberately scoped to a single canvas draw. -They do not change matplotlib's persistent axis state or public API. +There are two cache lifetimes: + +- Tick computations are reused only within one layout-and-render transaction. +- Relative axes outsets persist across transactions until their dependencies + change. + +``_LayoutTransaction`` owns both lifecycles. Temporary matplotlib method +overrides are restored when the transaction exits, including after exceptions. """ from __future__ import annotations from collections import OrderedDict +from contextlib import ExitStack from dataclasses import dataclass import matplotlib.transforms as mtransforms @@ -214,47 +221,6 @@ class _AxesExtentRecord: outsets: tuple -class _ExtentReductionTree: - """Small mutable min/max tree for the union of axes extents.""" - - def __init__(self, axes): - self.axes = tuple(axes) - self._indices = {axis: index for index, axis in enumerate(self.axes)} - capacity = 1 - while capacity < max(1, len(self.axes)): - capacity *= 2 - self._capacity = capacity - self._values = np.empty((2 * capacity, 4), dtype=float) - self.clear() - - def clear(self): - self._values[:, :2] = np.inf - self._values[:, 2:] = -np.inf - - def update(self, axis, bbox): - index = self._indices.get(axis) - if index is None: - return - node = self._capacity + index - if bbox is None: - self._values[node] = (np.inf, np.inf, -np.inf, -np.inf) - else: - self._values[node] = (bbox.xmin, bbox.ymin, bbox.xmax, bbox.ymax) - node //= 2 - while node: - left = self._values[2 * node] - right = self._values[2 * node + 1] - self._values[node, :2] = np.minimum(left[:2], right[:2]) - self._values[node, 2:] = np.maximum(left[2:], right[2:]) - node //= 2 - - def get_bbox(self): - xmin, ymin, xmax, ymax = self._values[1] - if not np.all(np.isfinite((xmin, ymin, xmax, ymax))): - return None - return mtransforms.Bbox.from_extents(xmin, ymin, xmax, ymax) - - class _LayoutExtentStore: """ Persist relative axes outsets and dependency versions between layouts. @@ -270,8 +236,7 @@ def __init__(self, figure): self.figure = figure self._records = {} self._versions = {} - self._tree = None - self._topology = {} + self._axes = () self._active = False self.hits = 0 self.misses = 0 @@ -289,9 +254,8 @@ def __enter__(self): def refresh(self): """Synchronize axes added by queued guide and panel creation.""" axes = tuple(self.figure.axes) - if self._tree is None or self._tree.axes != axes: - self._tree = _ExtentReductionTree(axes) - self._topology.clear() + if self._axes != axes: + self._axes = axes current = set(axes) self._records = { key: value for key, value in self._records.items() if key[0] in current @@ -324,9 +288,7 @@ def get_tightbbox( ): """Return an exact or reconstructed tight bbox in display units.""" if not self._active or not use_cache or not self._is_cacheable_axes(axes): - bbox = self._measure_tightbbox(axes, renderer, include_subset_titles) - self.update_union(axes, bbox) - return bbox + return self._measure_tightbbox(axes, renderer, include_subset_titles) version = self._versions.setdefault(axes, 1) state = self._get_state(axes, include_subset_titles) @@ -345,40 +307,8 @@ def get_tightbbox( state=self._get_state(axes, include_subset_titles), outsets=self._get_outsets(axes.bbox, bbox), ) - self.update_union(axes, bbox) return bbox - def clear_union(self): - if self._tree is not None: - self._tree.clear() - - def update_union(self, axes, bbox): - if self._tree is not None: - self._tree.update(axes, bbox) - - def get_union(self): - return None if self._tree is None else self._tree.get_bbox() - - def get_subplot_ranges(self, axes, along, across): - """Return compact topology arrays used by boundary reductions.""" - key = (tuple(axes), along, across) - cached = self._topology.get(key) - if cached is not None: - return cached - along_ranges = np.asarray( - [axis._range_subplotspec(along) for axis in axes], dtype=int - ) - across_ranges = np.asarray( - [axis._range_subplotspec(across) for axis in axes], dtype=int - ) - result = (along_ranges, across_ranges) - self._topology[key] = result - return result - - def invalidate_topology(self): - """Discard structural row/column range lookups.""" - self._topology.clear() - def _get_state(self, axes, include_subset_titles=True): bbox = axes.bbox position_sensitive = self._is_position_sensitive(axes) @@ -597,3 +527,52 @@ def _measure_tightbbox(axes, renderer, include_subset_titles): ) except TypeError: return axes.get_tightbbox(renderer) + + +class _LayoutTransaction: + """ + Own temporary and persistent caches for one dirty canvas draw. + + Figure code only needs to know whether a transaction is active. Cache setup, + dynamic-axes refresh, and exception-safe cleanup stay private to this object. + """ + + def __init__(self, figure, *, cache_ticks=True, cache_extents=True): + self.figure = figure + self.ticks = _AxisTickCache(figure) if cache_ticks else None + if cache_extents: + extents = getattr(figure, "_layout_extent_store", None) + if extents is None: + extents = figure._layout_extent_store = _LayoutExtentStore(figure) + self.extents = extents + else: + self.extents = None + self._stack = None + + def __enter__(self): + stack = self._stack = ExitStack() + self.figure._layout_transaction = self + try: + if self.ticks is not None: + stack.enter_context(self.ticks) + if self.extents is not None: + stack.enter_context(self.extents) + except Exception: + self.figure.__dict__.pop("_layout_transaction", None) + stack.close() + raise + return self + + def __exit__(self, *args): + try: + return self._stack.__exit__(*args) + finally: + self._stack = None + self.figure.__dict__.pop("_layout_transaction", None) + + def refresh(self): + """Synchronize caches after queued guides create axes or panels.""" + if self.ticks is not None: + self.ticks.refresh() + if self.extents is not None: + self.extents.refresh() diff --git a/ultraplot/_subplots.py b/ultraplot/_subplots.py index b0594973b..9e75a4ed3 100644 --- a/ultraplot/_subplots.py +++ b/ultraplot/_subplots.py @@ -153,7 +153,7 @@ def add_subplot(self, *args, **kwargs): The driver function for adding single subplots. """ fig = self.figure - fig._layout_dirty = True + fig._invalidate_layout() kwargs = self.parse_proj(**kwargs) args = args or (1, 1, 1) diff --git a/ultraplot/axes/base.py b/ultraplot/axes/base.py index fbae9408b..d0af2a91a 100644 --- a/ultraplot/axes/base.py +++ b/ultraplot/axes/base.py @@ -3336,7 +3336,11 @@ def format( ultraplot.config.Configurator.context """ if self.figure is not None and getattr(self, "_format_layout_required", True): - self.figure._layout_dirty = True + invalidate = getattr(self.figure, "_invalidate_layout", None) + if invalidate is None: + self.figure._layout_dirty = True + else: + invalidate() skip_figure = kwargs.pop("skip_figure", False) # internal keyword arg params = _pop_params(kwargs, self.figure._format_signature) diff --git a/ultraplot/figure.py b/ultraplot/figure.py index 17d0680b3..b5e98619d 100644 --- a/ultraplot/figure.py +++ b/ultraplot/figure.py @@ -42,7 +42,7 @@ labels, warnings, ) -from ._layout import _AxisTickCache, _LayoutExtentStore +from ._layout import _LayoutTransaction from ._subplots import SubplotManager from .utils import _Crawler, units @@ -647,18 +647,15 @@ def _canvas_preprocess(self, *args, **kwargs): ctx2 = fig._context_authorized() # skip backend set_constrained_layout() ctx3 = rc.context(fig._render_context) # draw with figure-specific setting ctx4 = ( - _AxisTickCache(fig) - if needs_layout and not getattr(fig, "_disable_axis_tick_cache", False) + _LayoutTransaction( + fig, + cache_ticks=not getattr(fig, "_disable_axis_tick_cache", False), + cache_extents=not getattr(fig, "_disable_layout_extent_cache", False), + ) + if needs_layout else context._empty_context() ) - if needs_layout and not getattr(fig, "_disable_layout_extent_cache", False): - store = getattr(fig, "_layout_extent_store", None) - if store is None: - store = fig._layout_extent_store = _LayoutExtentStore(fig) - ctx5 = store - else: - ctx5 = context._empty_context() - with ctx1, ctx2, ctx3, ctx4, ctx5: + with ctx1, ctx2, ctx3, ctx4: needs_post_layout = False if not fig._layout_initialized or layout_dirty: fig.auto_layout(tight=False if lock_tight_during_save else None) @@ -692,9 +689,6 @@ def wrapper(self, *args, **kwargs): result = func(self, *args, **kwargs) if hasattr(self, "_cached_border_axes"): delattr(self, "_cached_border_axes") - store = getattr(self, "_layout_extent_store", None) - if store is not None: - store.invalidate_topology() return result return wrapper @@ -1088,6 +1082,13 @@ def _init_super_labels(self): d["bottom"] = rc["bottomlabel.sharedpad"] d["top"] = rc["toplabel.sharedpad"] + def _invalidate_layout(self, *, reset=False): + """Mark automatic layout stale, optionally discarding persistent state.""" + self._layout_dirty = True + if reset: + self._layout_initialized = False + self.__dict__.pop("_layout_extent_store", None) + @_clear_border_cache def clear(self, keep_observers=False): """ @@ -1109,9 +1110,7 @@ def clear(self, keep_observers=False): super().clear(keep_observers=keep_observers) self._subplots.reset() self._panel_dict = {"left": [], "right": [], "bottom": [], "top": []} - self._layout_initialized = False - self._layout_dirty = True - self.__dict__.pop("_layout_extent_store", None) + self._invalidate_layout(reset=True) self._init_super_labels() @override @@ -2025,8 +2024,9 @@ def _get_layout_axes_bbox( use_cache=True, ): """Return an axes bbox using the active relative-outset store.""" - store = getattr(self, "_layout_extent_store", None) - if store is not None and store._active: + transaction = getattr(self, "_layout_transaction", None) + store = transaction.extents if transaction is not None else None + if store is not None: return store.get_tightbbox( axes, renderer, @@ -2045,7 +2045,7 @@ def _get_layout_tightbbox(self, renderer): Return the figure tight bbox while reusing relative axes outsets. This mirrors matplotlib's ``Figure.get_tightbbox`` but routes axes - measurements through the active dependency store and reduction tree. + measurements through the active relative-outset store. """ artists = [ artist @@ -2062,20 +2062,10 @@ def _get_layout_tightbbox(self, renderer): if bbox is not None: bboxes.append(bbox) - store = getattr(self, "_layout_extent_store", None) - if store is not None and store._active: - store.clear_union() for axes in self.axes: if not axes.get_visible(): - if store is not None and store._active: - store.update_union(axes, None) continue bbox = self._get_layout_axes_bbox(axes, renderer) - if store is None or not store._active: - if bbox is not None: - bboxes.append(bbox) - if store is not None and store._active: - bbox = store.get_union() if bbox is not None: bboxes.append(bbox) @@ -2276,7 +2266,7 @@ def _add_figure_panel( """ Add a figure panel. """ - self._layout_dirty = True + self._invalidate_layout() # Interpret args and enforce sensible keyword args side = _translate_loc(side, "panel", default="right") if side in ("left", "right"): @@ -3469,12 +3459,9 @@ def _align_content(): # noqa: E306 # WARNING: Tried to avoid two figure resizes but made # subsequent tight layout really weird. Have to resize twice. _draw_content() - manager = getattr(self, "_axis_tick_cache", None) - if manager is not None: - manager.refresh() - store = getattr(self, "_layout_extent_store", None) - if store is not None and store._active: - store.refresh() + transaction = getattr(self, "_layout_transaction", None) + if transaction is not None: + transaction.refresh() if not gs: return if aspect: @@ -3586,7 +3573,7 @@ def format( or bool(rc_kw) or axis_format_requires_layout(explicit_format_keys) ): - self._layout_dirty = True + self._invalidate_layout() kwargs.update(signature_axis_kwargs) with rc.context(rc_kw, mode=rc_mode): # Update background patch @@ -4277,7 +4264,7 @@ def set_size_inches(self, w, h=None, *, forward=True, internal=False, eps=None): if not samesize: # gridspec positions will resolve differently self.gridspec.update() if not backend and not internal: - self._layout_dirty = True + self._invalidate_layout() def _iter_axes(self, hidden=False, children=False, panels=True): """ diff --git a/ultraplot/gridspec.py b/ultraplot/gridspec.py index eed1e481b..89915645e 100644 --- a/ultraplot/gridspec.py +++ b/ultraplot/gridspec.py @@ -1117,12 +1117,8 @@ def _get_tight_space(self, w): # Iterate along each row or column space axs = tuple(fig._iter_axes(hidden=True, children=False)) space = list(space) # a copy - store = getattr(fig, "_layout_extent_store", None) - if store is not None and store._active: - ralong, racross = store.get_subplot_ranges(axs, x, y) - else: - ralong = np.array([ax._range_subplotspec(x) for ax in axs]) - racross = np.array([ax._range_subplotspec(y) for ax in axs]) + ralong = np.array([ax._range_subplotspec(x) for ax in axs]) + racross = np.array([ax._range_subplotspec(y) for ax in axs]) for i, (s, p) in enumerate(zip(space, pad)): # Find axes that abutt aginst this row or column space groups = [] diff --git a/ultraplot/tests/test_animation.py b/ultraplot/tests/test_animation.py index ddea1d418..ce95cde73 100644 --- a/ultraplot/tests/test_animation.py +++ b/ultraplot/tests/test_animation.py @@ -12,7 +12,7 @@ import ultraplot as uplt from ultraplot._animation import _BlitManager -from ultraplot._layout import _AxisTickCache +from ultraplot._layout import _AxisTickCache, _LayoutTransaction def test_auto_layout_not_called_on_every_frame(): @@ -68,6 +68,7 @@ def test_initial_draw_reuses_tick_updates(): assert stats["bypasses"] == 0 assert stats["evictions"] == 0 assert "_axis_tick_cache" not in fig.__dict__ + assert "_layout_transaction" not in fig.__dict__ for ax in fig.axes: for axis in ax._axis_map.values(): assert "_update_ticks" not in axis.__dict__ @@ -232,6 +233,45 @@ def test_tick_cache_includes_child_axes(): assert "_update_ticks" not in child_axes[0].xaxis.__dict__ +def test_layout_transaction_restores_after_exception(): + """ + Temporary hooks and active stores must be cleaned up after draw failures. + """ + fig, axs = uplt.subplots() + axis = axs[0].xaxis + + with pytest.raises(RuntimeError, match="draw failed"): + with _LayoutTransaction(fig) as transaction: + assert fig._layout_transaction is transaction + assert "_update_ticks" in axis.__dict__ + assert transaction.extents._active + raise RuntimeError("draw failed") + + assert "_layout_transaction" not in fig.__dict__ + assert "_axis_tick_cache" not in fig.__dict__ + assert "_update_ticks" not in axis.__dict__ + assert not fig._layout_extent_store._active + + +def test_layout_invalidation_has_one_lifecycle(): + """ + Normal invalidation preserves reusable state; reset discards it. + """ + fig, _ = uplt.subplots() + fig.canvas.draw() + store = fig._layout_extent_store + + fig._invalidate_layout() + assert fig._layout_dirty + assert fig._layout_initialized + assert fig._layout_extent_store is store + + fig._invalidate_layout(reset=True) + assert fig._layout_dirty + assert not fig._layout_initialized + assert "_layout_extent_store" not in fig.__dict__ + + def test_layout_extent_store_reuses_unmodified_axes(): """ A one-axes title edit should only remeasure that axes.