From 1e95beeb68c8b7f6c9c91faf0231eb2041d21ef7 Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Wed, 29 Jul 2026 20:41:28 +1000 Subject: [PATCH 1/8] 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/8] 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/8] 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/8] 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/8] 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/8] 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. From a2ab9ff7de320952c10c8d56ca198110a1f0306c Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Tue, 4 Aug 2026 20:50:45 +1000 Subject: [PATCH 7/8] perf: selectively redraw changed data axes --- ultraplot/_animation.py | 260 ++++++++++++++++++++++++++++++ ultraplot/_layout.py | 14 ++ ultraplot/figure.py | 29 +++- ultraplot/tests/test_animation.py | 101 ++++++++++++ 4 files changed, 402 insertions(+), 2 deletions(-) diff --git a/ultraplot/_animation.py b/ultraplot/_animation.py index 026365cad..2eb59d089 100644 --- a/ultraplot/_animation.py +++ b/ultraplot/_animation.py @@ -10,7 +10,267 @@ from weakref import WeakSet import matplotlib.artist as martist +import matplotlib.collections as mcollections +import matplotlib.image as mimage +import matplotlib.lines as mlines import matplotlib.transforms as mtransforms +from matplotlib.backend_bases import DrawEvent + + +class _SelectiveDrawManager: + """Cache axes layers and redraw only axes with paint-only data changes.""" + + _data_artist_types = (mlines.Line2D, mcollections.Collection, mimage.AxesImage) + + def __init__(self, canvas, figure=None): + self.canvas = canvas + self.figure = figure if figure is not None else canvas.figure + self._backgrounds = {} + self._regions = {} + self._signatures = {} + self._axes = () + self._capturing = False + self._selecting = False + self._suspended = False + self._closed = False + self._cache_safe = False + self._selective_draw_count = 0 + self._full_draw_count = 0 + self._supports_blit = bool(canvas.supports_blit) + callbacks = self.figure._canvas_callbacks + self._draw_cid = callbacks.connect("draw_event", self._on_draw) + self._resize_cid = callbacks.connect("resize_event", self._on_resize) + + @staticmethod + def _bbox_signature(bbox): + return tuple(float(value) for value in bbox.bounds) + + def _axes_signature(self, ax): + return ( + tuple(id(child) for child in ax.get_children()), + self._bbox_signature(ax.get_position(original=False)), + self._bbox_signature(ax.viewLim), + ax.get_xscale(), + ax.get_yscale(), + bool(ax.get_visible()), + float(ax.get_zorder()), + ) + + def _visible_axes(self): + return tuple(ax for ax in self.figure.axes if ax.get_visible()) + + def _has_explicit_blit_manager(self): + return bool(tuple(getattr(self.figure, "_blit_managers", ()))) + + def _has_animated_artist(self, axes): + return any(artist.get_animated() for ax in axes for artist in ax.get_children()) + + @staticmethod + def _regions_overlap(regions): + regions = tuple(regions) + for idx, left in enumerate(regions): + for right in regions[idx + 1 :]: + overlap = mtransforms.Bbox.intersection(left, right) + if overlap is not None and overlap.width > 0 and overlap.height > 0: + return True + return False + + def _resolve_region(self, ax, renderer, cached_regions): + # UltraLayout already measured this exact region. Reconstructing it + # from cached outsets avoids a second tick/text traversal. + region = cached_regions.get(ax) + if region is None: + region = ax.get_tightbbox(renderer) + if region is None: + return None + region = region.padded(2) + return mtransforms.Bbox.intersection(region, self.figure.bbox) + + @staticmethod + def _mark_axes_clean(axes): + """Clear placeholder staleness left behind by ``Axes.draw()``.""" + for ax in axes: + for child in ax.get_children(): + child.stale = False + ax.stale = False + + def invalidate(self): + """Discard all retained axes layers.""" + self._backgrounds.clear() + self._regions.clear() + self._signatures.clear() + self._axes = () + self._cache_safe = False + + def _on_resize(self, event): + if event is None or event.canvas is self.canvas: + self.invalidate() + + def _on_draw(self, event): + if ( + self._closed + or self._suspended + or self._selecting + or not self._capturing + or event.canvas is not self.canvas + ): + return + + axes = self._visible_axes() + renderer = event.renderer + store = getattr(self.figure, "_layout_extent_store", None) + cached_regions = {} if store is None else store._get_retained_bboxes(axes) + regions = { + ax: self._resolve_region(ax, renderer, cached_regions) for ax in axes + } + valid = all(region is not None for region in regions.values()) + if valid: + backgrounds = { + ax: self.canvas.copy_from_bbox(region) for ax, region in regions.items() + } + else: + backgrounds = {} + + # Figure.draw() omitted these temporarily animated axes. Draw them now, + # in normal axes z-order, before the backend presents the frame. + for ax in sorted(axes, key=lambda item: item.get_zorder()): + self.figure.draw_artist(ax) + # Axes.draw() intentionally leaves some invisible placeholder text stale. + # Retained drawing needs a clean baseline so a later property mutation can + # be distinguished from that permanent state. + self._mark_axes_clean(axes) + for artist in self.figure.get_children(): + if artist not in axes: + artist.stale = False + + self._full_draw_count += 1 + self._axes = axes + self._regions = regions if valid else {} + self._backgrounds = backgrounds + self._signatures = {ax: self._axes_signature(ax) for ax in axes} + self._cache_safe = bool( + valid + and len(axes) > 1 + and not self._regions_overlap(regions.values()) + and not self._has_explicit_blit_manager() + and not self._has_animated_artist(axes) + ) + for ax in axes: + ax.stale = False + + @contextmanager + def full_draw_context(self): + """Temporarily split a full draw into static and axes layers.""" + if ( + self._closed + or self._suspended + or not self._supports_blit + or self._has_explicit_blit_manager() + ): + self.invalidate() + yield + return + + axes = self._visible_axes() + if len(axes) <= 1: + self.invalidate() + yield + return + animated = {ax: ax.get_animated() for ax in axes} + if any(animated.values()) or self._has_animated_artist(axes): + self.invalidate() + yield + return + + self._capturing = True + try: + for ax in axes: + # This is a transient draw-routing flag, not a user property + # mutation. Avoid set_animated(), which marks every axes stale + # and defeats the persistent layout extent cache. + ax._animated = True + yield + finally: + for ax, state in animated.items(): + ax._animated = state + self._capturing = False + + def _dirty_axes(self): + axes = self._visible_axes() + if axes != self._axes: + return None + if any( + artist.stale for artist in self.figure.get_children() if artist not in axes + ): + return None + + dirty = [] + for ax in axes: + if self._axes_signature(ax) != self._signatures.get(ax): + return None + stale_children = [child for child in ax.get_children() if child.stale] + # Layout/tick cache cleanup can leave the Axes container stale even + # though every drawable child is clean. Child flags carry the useful + # paint-level signal; geometry is guarded by the signature above. + if not stale_children: + continue + if not all( + isinstance(child, self._data_artist_types) for child in stale_children + ): + return None + dirty.append(ax) + return tuple(dirty) + + def draw_if_possible(self): + """Use retained axes layers for paint-only data changes.""" + if ( + self._closed + or self._suspended + or not self._supports_blit + or not self._cache_safe + or self._has_explicit_blit_manager() + or getattr(self.figure, "_layout_dirty", False) + ): + return False + + dirty = self._dirty_axes() + if not dirty: + return False + + self._selecting = True + try: + for ax in dirty: + self.canvas.restore_region(self._backgrounds[ax]) + for ax in sorted(dirty, key=lambda item: item.get_zorder()): + self.figure.draw_artist(ax) + self._mark_axes_clean(dirty) + for ax in dirty: + self.canvas.blit(self._regions[ax]) + self.figure.stale = False + self._selective_draw_count += 1 + DrawEvent("draw_event", self.canvas, self.canvas.get_renderer())._process() + finally: + self._selecting = False + return True + + @contextmanager + def save_context(self): + """Suspend retained drawing while producing external output.""" + self._suspended = True + try: + yield + finally: + self._suspended = False + self.invalidate() + + def close(self): + if self._closed: + return + callbacks = self.figure._canvas_callbacks + callbacks.disconnect(self._draw_cid) + callbacks.disconnect(self._resize_cid) + self.invalidate() + self._closed = True class _BlitManager: diff --git a/ultraplot/_layout.py b/ultraplot/_layout.py index 792dae3be..936354245 100644 --- a/ultraplot/_layout.py +++ b/ultraplot/_layout.py @@ -369,6 +369,20 @@ def _rebase_records(self): ): record.state = state + def _get_retained_bboxes(self, axes): + """Return exact cached display bboxes for retained axes drawing.""" + self._rebase_records() + bboxes = {} + for axis in axes: + record = self._records.get((axis, True)) + if ( + record is not None + and record.version == self._versions.get(axis) + and record.state == self._get_state(axis, True) + ): + bboxes[axis] = self._bbox_from_outsets(axis.bbox, record.outsets) + return bboxes + @staticmethod def _get_decoration_state(axes): texts = [] diff --git a/ultraplot/figure.py b/ultraplot/figure.py index b5e98619d..fe9794f08 100644 --- a/ultraplot/figure.py +++ b/ultraplot/figure.py @@ -664,10 +664,26 @@ def _canvas_preprocess(self, *args, **kwargs): needs_post_layout = ( not lock_tight_during_save and _needs_post_tight_layout(fig) ) - result = func(self, *args, **kwargs) + selective = getattr(fig, "_selective_draw_manager", None) + if ( + method != "print_figure" + and selective is not None + and not needs_layout + and selective.draw_if_possible() + ): + return None + + def _draw_context(): + if method != "print_figure" and selective is not None: + return selective.full_draw_context() + return context._empty_context() + + with _draw_context(): + result = func(self, *args, **kwargs) if needs_post_layout: fig.auto_layout() - result = func(self, *args, **kwargs) + with _draw_context(): + result = func(self, *args, **kwargs) if method == "print_figure" and getattr(self, "_is_saving", False): fig._saving_frame_count = saving_frame_count + 1 elif not getattr(self, "_is_saving", False): @@ -4145,6 +4161,9 @@ def savefig(self, filename, **kwargs): # so savefig includes them in normal z-order. managers = tuple(getattr(self, "_blit_managers", ())) with ExitStack() as stack: + selective = getattr(self, "_selective_draw_manager", None) + if selective is not None: + stack.enter_context(selective.save_context()) for manager in managers: stack.enter_context(manager._save_context()) # NOTE: this draw ensures that we are applying ultraplots layout @@ -4180,6 +4199,9 @@ def set_canvas(self, canvas): # around this by forcing additional draw() call in this function before # proceeding with print_figure). Set the canvas and add monkey patches # to the instance-level draw and print_figure methods. + previous = getattr(self, "_selective_draw_manager", None) + if previous is not None: + previous.close() method = "draw" # if getattr(canvas, "_draw", None): # method = "_draw" @@ -4198,6 +4220,9 @@ def _draw_idle(self, *args, **kwargs): canvas.draw_idle = _draw_idle.__get__(canvas) super().set_canvas(canvas) + from ._animation import _SelectiveDrawManager + + self._selective_draw_manager = _SelectiveDrawManager(canvas, self) def _is_same_size(self, figsize, eps=None): """ diff --git a/ultraplot/tests/test_animation.py b/ultraplot/tests/test_animation.py index ce95cde73..d8266cfc7 100644 --- a/ultraplot/tests/test_animation.py +++ b/ultraplot/tests/test_animation.py @@ -502,6 +502,107 @@ def test_blit_manager_rejects_artist_from_another_figure(): fig1._blit_manager(line) +def test_selective_draw_full_layer_matches_normal_draw(): + """Splitting a full draw into static and axes layers must preserve pixels.""" + fig, axs = uplt.subplots(ncols=2) + for index, ax in enumerate(axs): + ax.plot([0, 0.5, 1], [0.2, 0.8 - 0.1 * index, 0.3]) + ax.format(xlabel=f"x {index}", ylabel=f"y {index}", title=f"axis {index}") + fig.format(suptitle="Layer fidelity") + manager = fig._selective_draw_manager + + with manager.save_context(): + fig.canvas.draw() + normal = np.asarray(fig.canvas.buffer_rgba()).copy() + fig.canvas.draw() + layered = np.asarray(fig.canvas.buffer_rgba()).copy() + + assert np.array_equal(layered, normal) + + +def test_selective_draw_redraws_only_changed_axes(): + """A paint-only line update should bypass every unchanged axes.""" + fig, axs = uplt.subplots(ncols=2) + lines = [ax.plot([0, 0.5, 1], [0.2, 0.8, 0.3])[0] for ax in axs] + fig.canvas.draw() + manager = fig._selective_draw_manager + unchanged_draw = axs[1].draw + axs[1].draw = MagicMock(wraps=unchanged_draw) + + lines[0].set_ydata([0.8, 0.2, 0.7]) + fig.canvas.draw() + + assert manager._selective_draw_count == 1 + axs[1].draw.assert_not_called() + + +def test_selective_draw_stays_fast_across_consecutive_frames(): + """Placeholder staleness must not force alternating full redraws.""" + fig, axs = uplt.subplots(ncols=2) + x = np.linspace(0, 2 * np.pi, 100) + line = axs[0].plot(x, np.sin(x))[0] + axs[1].plot(x, np.cos(x)) + fig.canvas.draw() + manager = fig._selective_draw_manager + + for frame in range(4): + line.set_ydata(np.sin(x + frame / 10)) + fig.canvas.draw() + + assert manager._selective_draw_count == 4 + assert manager._full_draw_count == 1 + + +@pytest.mark.parametrize("kind", ("line", "scatter")) +def test_selective_draw_matches_complete_data_draw(kind): + """Retained line and collection updates must match a complete Agg draw.""" + fig, axs = uplt.subplots(ncols=2) + axs[1].plot([0, 1], [1, 0]) + if kind == "line": + artist = axs[0].plot([0, 0.5, 1], [0.2, 0.8, 0.3])[0] + else: + artist = axs[0].scatter([0, 0.5, 1], [0.2, 0.8, 0.3]) + fig.canvas.draw() + manager = fig._selective_draw_manager + + if kind == "line": + artist.set_ydata([0.8, 0.2, 0.7]) + else: + artist.set_offsets([[0, 0.8], [0.5, 0.2], [1, 0.7]]) + fig.canvas.draw() + selective = np.asarray(fig.canvas.buffer_rgba()).copy() + + with manager.save_context(): + fig.canvas.draw() + complete = np.asarray(fig.canvas.buffer_rgba()).copy() + + assert manager._selective_draw_count == 1 + assert np.array_equal(selective, complete) + + +@pytest.mark.parametrize("change", ("limits", "text", "structure")) +def test_selective_draw_falls_back_for_unsafe_changes(change): + """Geometry, text, and artist-tree changes must retain full-draw semantics.""" + fig, axs = uplt.subplots(ncols=2) + ax = axs[0] + ax.plot([0, 1], [0, 1]) + axs[1].plot([0, 1], [1, 0]) + fig.canvas.draw() + manager = fig._selective_draw_manager + full_count = manager._full_draw_count + + if change == "limits": + ax.set_xlim(-1, 2) + elif change == "text": + ax.set_title("Changed") + else: + ax.plot([0, 1], [1, 0]) + fig.canvas.draw() + + assert manager._selective_draw_count == 0 + assert manager._full_draw_count == full_count + 1 + + def test_layout_array_no_crash(): """ Test that using layout_array with FuncAnimation does not crash. From ed9aa87dc80bff55574967118cb5ff504bd32eff Mon Sep 17 00:00:00 2001 From: cvanelteren Date: Tue, 4 Aug 2026 21:18:20 +1000 Subject: [PATCH 8/8] perf: retain stable single-axes draw layers --- ultraplot/_animation.py | 199 ++++++++++++++++++++++++++---- ultraplot/tests/test_animation.py | 104 ++++++++++++++++ 2 files changed, 280 insertions(+), 23 deletions(-) diff --git a/ultraplot/_animation.py b/ultraplot/_animation.py index 2eb59d089..5163aa978 100644 --- a/ultraplot/_animation.py +++ b/ultraplot/_animation.py @@ -10,15 +10,27 @@ from weakref import WeakSet import matplotlib.artist as martist +import matplotlib.axis as maxis import matplotlib.collections as mcollections import matplotlib.image as mimage import matplotlib.lines as mlines +import matplotlib.text as mtext import matplotlib.transforms as mtransforms +import numpy as np from matplotlib.backend_bases import DrawEvent class _SelectiveDrawManager: - """Cache axes layers and redraw only axes with paint-only data changes.""" + """ + Retain safe draw layers and bypass unchanged Matplotlib traversal. + + Multi-axes figures retain each complete axes as one layer. Single Cartesian + axes retain the stable draw-order prefix below their first clipped numeric + line, then redraw that line and every later artist as an exact z-order suffix. + Unknown stale artists, geometry changes, overlapping layers, unsupported + artist orders, and export draws fall back to a complete draw. The first display + is always untouched; a later complete draw primes the retained layers. + """ _data_artist_types = (mlines.Line2D, mcollections.Collection, mimage.AxesImage) @@ -29,8 +41,13 @@ def __init__(self, canvas, figure=None): self._regions = {} self._signatures = {} self._axes = () + self._suffix = () + self._pending_suffix = () + self._cache_mode = None + self._capture_mode = None self._capturing = False self._selecting = False + self._has_drawn = False self._suspended = False self._closed = False self._cache_safe = False @@ -47,7 +64,19 @@ def _bbox_signature(bbox): def _axes_signature(self, ax): return ( - tuple(id(child) for child in ax.get_children()), + tuple( + (id(child), float(child.get_zorder())) for child in ax.get_children() + ), + tuple( + ( + id(line), + id(line.get_transform()), + bool(line.get_clip_on()), + id(line.get_clip_box()), + id(line.get_clip_path()), + ) + for line in ax.lines + ), self._bbox_signature(ax.get_position(original=False)), self._bbox_signature(ax.viewLim), ax.get_xscale(), @@ -65,6 +94,80 @@ def _has_explicit_blit_manager(self): def _has_animated_artist(self, axes): return any(artist.get_animated() for ax in axes for artist in ax.get_children()) + def _has_overlapping_figure_artist(self, suffix): + """Return whether a later figure artist overlaps the retained suffix.""" + get_renderer = getattr(self.canvas, "get_renderer", None) + if get_renderer is None: + return True + renderer = get_renderer() + for artist in self.figure.get_children(): + if ( + artist is self.figure.patch + or artist in self.figure.axes + or not artist.get_visible() + ): + continue + if isinstance(artist, mtext.Text) and not artist.get_text(): + continue + try: + overlay_bbox = artist.get_window_extent(renderer) + except Exception: + return True + for child in suffix: + if not child.get_visible(): + continue + try: + child_bbox = child.get_window_extent(renderer) + except Exception: + return True + overlap = mtransforms.Bbox.intersection(overlay_bbox, child_bbox) + if overlap is not None and overlap.width > 0 and overlap.height > 0: + return True + return False + + @staticmethod + def _has_numeric_line_data(line): + """Return whether line conversion cannot mutate categorical/date axes.""" + try: + xdata = np.asanyarray(line.get_xdata(orig=True)) + ydata = np.asanyarray(line.get_ydata(orig=True)) + except (TypeError, ValueError): + return False + return xdata.dtype.kind in "biufc" and ydata.dtype.kind in "biufc" + + def _resolve_line_suffix(self, ax): + """Return the exact draw-order suffix starting at the first data line.""" + if ( + getattr(ax, "_name", None) != "cartesian" + or not ax.axison + or not ax.get_frame_on() + or ax.get_rasterization_zorder() is not None + or any(not line.get_clip_on() for line in ax.lines) + or any(not self._has_numeric_line_data(line) for line in ax.lines) + ): + return () + artists = list(ax.get_children()) + if ax.patch not in artists or not ax.lines: + return () + artists.remove(ax.patch) + artists = sorted(artists, key=lambda artist: artist.get_zorder()) + line_ids = {id(line) for line in ax.lines} + positions = [ + index for index, artist in enumerate(artists) if id(artist) in line_ids + ] + if not positions: + return () + suffix = tuple(artists[min(positions) :]) + if any( + isinstance(artist, (maxis.Axis, mimage.AxesImage)) + or (hasattr(artist, "_axis_map") and artist is not ax) + for artist in suffix + ): + return () + if self._has_overlapping_figure_artist(suffix): + return () + return suffix + @staticmethod def _regions_overlap(regions): regions = tuple(regions) @@ -100,6 +203,10 @@ def invalidate(self): self._regions.clear() self._signatures.clear() self._axes = () + self._suffix = () + self._pending_suffix = () + self._cache_mode = None + self._capture_mode = None self._cache_safe = False def _on_resize(self, event): @@ -107,6 +214,8 @@ def _on_resize(self, event): self.invalidate() def _on_draw(self, event): + if event is not None and event.canvas is self.canvas and not self._selecting: + self._has_drawn = True if ( self._closed or self._suspended @@ -131,10 +240,16 @@ def _on_draw(self, event): else: backgrounds = {} - # Figure.draw() omitted these temporarily animated axes. Draw them now, - # in normal axes z-order, before the backend presents the frame. - for ax in sorted(axes, key=lambda item: item.get_zorder()): - self.figure.draw_artist(ax) + if self._capture_mode == "suffix": + suffix = self._pending_suffix + for artist in suffix: + self.figure.draw_artist(artist) + else: + suffix = () + # Figure.draw() omitted these temporarily animated axes. Draw them now, + # in normal axes z-order, before the backend presents the frame. + for ax in sorted(axes, key=lambda item: item.get_zorder()): + self.figure.draw_artist(ax) # Axes.draw() intentionally leaves some invisible placeholder text stale. # Retained drawing needs a clean baseline so a later property mutation can # be distinguished from that permanent state. @@ -145,16 +260,21 @@ def _on_draw(self, event): self._full_draw_count += 1 self._axes = axes + self._suffix = suffix self._regions = regions if valid else {} self._backgrounds = backgrounds self._signatures = {ax: self._axes_signature(ax) for ax in axes} - self._cache_safe = bool( - valid - and len(axes) > 1 - and not self._regions_overlap(regions.values()) - and not self._has_explicit_blit_manager() - and not self._has_animated_artist(axes) - ) + self._cache_mode = self._capture_mode + if self._capture_mode == "suffix": + self._cache_safe = bool(valid and suffix) + else: + self._cache_safe = bool( + valid + and len(axes) > 1 + and not self._regions_overlap(regions.values()) + and not self._has_explicit_blit_manager() + and not self._has_animated_artist(axes) + ) for ax in axes: ax.stale = False @@ -172,28 +292,49 @@ def full_draw_context(self): return axes = self._visible_axes() - if len(axes) <= 1: + if not axes: + self.invalidate() + yield + return + if not self._has_drawn: self.invalidate() yield return - animated = {ax: ax.get_animated() for ax in axes} - if any(animated.values()) or self._has_animated_artist(axes): + if any(ax.get_animated() for ax in axes) or self._has_animated_artist(axes): self.invalidate() yield return + if len(axes) == 1: + suffix = self._resolve_line_suffix(axes[0]) + if not suffix: + self.invalidate() + yield + return + targets = suffix + self._capture_mode = "suffix" + self._pending_suffix = suffix + else: + targets = axes + self._capture_mode = "axes" + self._pending_suffix = () + + animated = {artist: artist.get_animated() for artist in targets} + self._capturing = True try: - for ax in axes: + for artist in targets: # This is a transient draw-routing flag, not a user property # mutation. Avoid set_animated(), which marks every axes stale # and defeats the persistent layout extent cache. - ax._animated = True + artist._animated = True yield finally: - for ax, state in animated.items(): - ax._animated = state + for artist, state in animated.items(): + artist._animated = state self._capturing = False + self._capture_mode = None + self._pending_suffix = () def _dirty_axes(self): axes = self._visible_axes() @@ -214,7 +355,15 @@ def _dirty_axes(self): # paint-level signal; geometry is guarded by the signature above. if not stale_children: continue - if not all( + if self._cache_mode == "suffix": + if not all( + isinstance(child, mlines.Line2D) + and child in self._suffix + and self._has_numeric_line_data(child) + for child in stale_children + ): + return None + elif not all( isinstance(child, self._data_artist_types) for child in stale_children ): return None @@ -241,8 +390,12 @@ def draw_if_possible(self): try: for ax in dirty: self.canvas.restore_region(self._backgrounds[ax]) - for ax in sorted(dirty, key=lambda item: item.get_zorder()): - self.figure.draw_artist(ax) + if self._cache_mode == "suffix": + for artist in self._suffix: + self.figure.draw_artist(artist) + else: + for ax in sorted(dirty, key=lambda item: item.get_zorder()): + self.figure.draw_artist(ax) self._mark_axes_clean(dirty) for ax in dirty: self.canvas.blit(self._regions[ax]) diff --git a/ultraplot/tests/test_animation.py b/ultraplot/tests/test_animation.py index d8266cfc7..7a23d7362 100644 --- a/ultraplot/tests/test_animation.py +++ b/ultraplot/tests/test_animation.py @@ -520,11 +520,29 @@ def test_selective_draw_full_layer_matches_normal_draw(): assert np.array_equal(layered, normal) +@pytest.mark.parametrize("ncols", (1, 2)) +def test_selective_draw_defers_cache_until_after_initial_display(ncols): + """The first display should have no retained-layer setup overhead.""" + fig, _axs = uplt.subplots(ncols=ncols) + for ax in fig.axes: + ax.plot([0, 1], [0, 1]) + manager = fig._selective_draw_manager + + fig.canvas.draw() + assert manager._cache_mode is None + assert manager._full_draw_count == 0 + + fig.canvas.draw() + assert manager._cache_mode == ("suffix" if ncols == 1 else "axes") + assert manager._full_draw_count == 1 + + def test_selective_draw_redraws_only_changed_axes(): """A paint-only line update should bypass every unchanged axes.""" fig, axs = uplt.subplots(ncols=2) lines = [ax.plot([0, 0.5, 1], [0.2, 0.8, 0.3])[0] for ax in axs] fig.canvas.draw() + fig.canvas.draw() # Prime retained axes layers after the initial display. manager = fig._selective_draw_manager unchanged_draw = axs[1].draw axs[1].draw = MagicMock(wraps=unchanged_draw) @@ -543,6 +561,7 @@ def test_selective_draw_stays_fast_across_consecutive_frames(): line = axs[0].plot(x, np.sin(x))[0] axs[1].plot(x, np.cos(x)) fig.canvas.draw() + fig.canvas.draw() # Prime retained axes layers after the initial display. manager = fig._selective_draw_manager for frame in range(4): @@ -553,6 +572,90 @@ def test_selective_draw_stays_fast_across_consecutive_frames(): assert manager._full_draw_count == 1 +def test_selective_draw_single_axes_line_suffix_matches_complete_draw(): + """A retained line suffix should skip axes while preserving exact pixels.""" + fig, ax = uplt.subplots() + x = np.linspace(0, 2 * np.pi, 200) + ax.scatter(x[::8], np.cos(x[::8]), s=8, zorder=1) + ax.plot(x, 0.4 * np.cos(x), color="orange", zorder=2) + line = ax.plot(x, np.sin(x), color="cerulean", zorder=2)[0] + ax.text(0.03, 0.96, "Annotation", transform=ax.transAxes, va="top", zorder=3.2) + ax.format(title="Retained suffix", xlabel="x", ylabel="y", grid=True) + fig.canvas.draw() + manager = fig._selective_draw_manager + axis_draw = ax.xaxis.draw + ax.xaxis.draw = MagicMock(wraps=axis_draw) + + line.set_ydata(np.sin(x + 0.1)) + fig.canvas.draw() # Prime the retained suffix on the first redraw. + ax.xaxis.draw.reset_mock() + for frame in range(4): + line.set_ydata(np.sin(x + frame / 10)) + fig.canvas.draw() + retained = np.asarray(fig.canvas.buffer_rgba()).copy() + + assert manager._cache_mode == "suffix" + assert manager._selective_draw_count == 4 + assert manager._full_draw_count == 1 + ax.xaxis.draw.assert_not_called() + + with manager.save_context(): + fig.canvas.draw() + complete = np.asarray(fig.canvas.buffer_rgba()).copy() + assert np.array_equal(retained, complete) + + +def test_selective_draw_single_axes_savefig_matches_unretained_output(): + """Saving must restore ordinary ordering after retained suffix updates.""" + fig, ax = uplt.subplots() + line = ax.plot([0, 1, 2], [0.8, 0.2, 0.7], color="red", lw=3)[0] + ax.format(xlim=(0, 2), ylim=(0, 1), title="Export fidelity") + + expected_buffer = io.BytesIO() + fig.savefig(expected_buffer, format="png") + expected_buffer.seek(0) + expected = np.asarray(Image.open(expected_buffer)).copy() + + fig.canvas.draw() # Prime the retained suffix after the initial save. + line.set_ydata([0.2, 0.8, 0.3]) + fig.canvas.draw() + line.set_ydata([0.8, 0.2, 0.7]) + fig.canvas.draw() + assert fig._selective_draw_manager._selective_draw_count == 2 + + retained_buffer = io.BytesIO() + fig.savefig(retained_buffer, format="png") + retained_buffer.seek(0) + retained = np.asarray(Image.open(retained_buffer)).copy() + + assert np.array_equal(retained, expected) + + +@pytest.mark.parametrize("kind", ("scatter", "low_line", "unclipped", "overlay")) +def test_selective_draw_single_axes_falls_back_when_suffix_is_unsafe(kind): + """Single-axes retention requires a clipped line above the axes layer.""" + fig, ax = uplt.subplots() + if kind == "scatter": + artist = ax.scatter([0, 1], [0, 1]) + else: + artist = ax.plot([0, 1], [0, 1], zorder=1 if kind == "low_line" else 2)[0] + if kind == "unclipped": + artist.set_clip_on(False) + elif kind == "overlay": + fig.text(0.5, 0.5, "Figure overlay", zorder=10) + fig.canvas.draw() + manager = fig._selective_draw_manager + + if kind == "scatter": + artist.set_offsets([[0, 1], [1, 0]]) + else: + artist.set_ydata([1, 0]) + fig.canvas.draw() + + assert manager._cache_mode is None + assert manager._selective_draw_count == 0 + + @pytest.mark.parametrize("kind", ("line", "scatter")) def test_selective_draw_matches_complete_data_draw(kind): """Retained line and collection updates must match a complete Agg draw.""" @@ -563,6 +666,7 @@ def test_selective_draw_matches_complete_data_draw(kind): else: artist = axs[0].scatter([0, 0.5, 1], [0.2, 0.8, 0.3]) fig.canvas.draw() + fig.canvas.draw() # Prime retained axes layers after the initial display. manager = fig._selective_draw_manager if kind == "line":