Skip transforming all geometries in show()'s get_extent when the transform is axis-aligned (large-shapes perf)

[timtreis]

Summary

pl.show() computes the figure's axis bounds via get_extent(sdata, coordinate_system=cs, exact=True) (basic.py:1831). For shapes/points, spatialdata's exact=True path transforms every geometry into the coordinate system (per-geometry shapely affine, O(N)) just to take a bounding box. For large shape collections this dominates the render and is unrelated to what is being drawn.

Measured on the real Visium HD dataset (render_shapes):

shapes	get_extent exact=True	exact=False	result
351,817	6.9 s	0.97 s	identical
5,479,660	109 s	14.6 s	identical

At 5.5M shapes this get_extent call is ~85% of the whole render (~109 s of ~128 s), independent of method/as_points.

The heuristic

exact=False transforms only the bounding-box corners instead of all geometries. Mathematically:

• exact=True = bbox({T(g) for every geometry g})
• exact=False = bbox(T(corners of the intrinsic bbox))

Since geometries ⊆ intrinsic bbox, exact=False ⊇ exact=True always, and they are equal iff the transform maps axis-aligned boxes to axis-aligned boxes — i.e. the 2×2 linear part is a monomial matrix (exactly one non-zero per row and per column): scale, axis flips, 90°/180°/270° rotations, axis swaps (+ any translation). Only true rotation/shear makes exact=False over-estimate. spatialdata's own get_extent docstring confirms this: "the exact and approximate extent are the same if the transformation does not contain any rotation or shear."

This covers essentially all real Visium/Xenium/MERFISH data (readers produce scale + translation).

Proposed change (spatialdata-plot)

In show(), before calling get_extent, inspect the transforms of the wanted shapes/points elements to the coordinate system; if all are axis-aligned, pass exact=False, else keep exact=True.

def _is_axis_aligned(linear2x2, *, rtol=1e-9):
    """Sends axis-aligned boxes to axis-aligned boxes (scale/flip/90deg/swap) -> exact==approx."""
    m = np.asarray(linear2x2, dtype=float)
    nz = np.abs(m) > rtol * (np.abs(m).max() or 1.0)   # relative tol: ignore float noise
    return bool((nz.sum(0) <= 1).all() and (nz.sum(1) <= 1).all() and nz.sum() == m.shape[0])

• Conservative: any rotated/sheared element among the rendered ones → keep exact=True for the whole call, so output is never wrong.
• Validated end-to-end: render_shapes(as_points=True) on 351k shapes went 8.5 s → 2.6 s (3.2×) by forcing exact=False; 5.5M projected ~128 s → ~33 s (~4×). Extents identical for non-rotation transforms.
• Speeds up the non-as_points shapes render too, since the extent cost is shared.
• Use a relative tolerance in the monomial check so float noise in to_affine_matrix isn't misread as shear.

Acceptance / tests

• exact=False is chosen for a pure scale/translation element; exact=True for a rotated one.
• A render benchmark / assertion that extents are unchanged for scale transforms.

Note: upstream follow-up (spatialdata)

A get_extent(exact=False) still has an O(N)-Python floor: _get_extent_of_shapes filters empties with e["geometry"].apply(lambda g: not g.is_empty) (~14.6 s of the 5.5M number); vectorized ~e.geometry.is_empty is 128× faster. The deeper fix is in spatialdata itself: (1) auto-detect axis-aligned transforms inside get_extent so even exact=True uses the cheap corner-transform path (fixes squidpy/napari too), and (2) vectorize that empty filter. The spatialdata-plot heuristic above is the immediate, self-contained mitigation.

Context: surfaced while profiling render_shapes/render_labels(as_points=True) on Visium HD; the geometry-transform in get_extent, not the scatter, is the bottleneck.

[timtreis]

Implemented + demonstrated (isolated for upstreaming)

Implemented the full fix entirely on the spatialdata-plot side, structured as a single self-contained drop-in so it can be lifted into spatialdata's get_extent verbatim. Commit: 57bcf8d (branch feat/centroid-scatter-helper).

pl/utils.py — three isolated helpers:

• _is_axis_aligned(linear2x2) — monomial-matrix check (scale / flip / 90° / swap), relative tolerance.
• _element_extent_fast(element, cs) — corner-transform extent for one shapes/points element; returns None to signal fallback under rotation/shear and (for circles) anisotropic scale.
• get_extent_fast(sdata, cs, ...) — drop-in for get_extent(sdata, ...); fast path for axis-aligned shapes/points, delegates to spatialdata's get_extent otherwise (rotation/shear, images/labels). Also reads intrinsic bounds vectorised (geometry.bounds / radius arithmetic on numpy), avoiding spatialdata's per-geometry .apply(is_empty) floor.

show() now calls get_extent_fast.

Measured (real Visium HD, render_shapes)

	shapes	old (get_extent)	new (get_extent_fast)	speedup
as_points=True	351,817	8.5 s	1.8 s	4.7×
as_points=True	5,479,660	129 s	24 s	5.5×
full (patches)	351,817	23.8 s	15.2 s	1.6×

Correctness

get_extent_fast is asserted identical to get_extent(exact=True) across scale (iso/aniso), translation, flip, 90° rotation, 45° rotation, and shear, for both circles and polygons (tests/pl/test_utils.py::TestGetExtentFast, 20 cases). Axis-aligned cases use the fast path; rotation/shear and anisotropic circles fall back, so output never differs. 397 non-visual tests pass.

Notes for upstreaming to spatialdata

The element-level pieces (_is_axis_aligned, vectorised intrinsic bounds, corner transform) are the natural optimization for get_extent itself:

1. Inside get_extent, detect axis-aligned transforms and use the corner-transform path even when exact=True — fixes squidpy/napari/etc. with no caller change.
2. Vectorize the empty filter in _get_extent_of_shapes (~e.geometry.is_empty) — ~128× on that step (the residual floor in the table above).

The circle subtlety (a scaled circle is an ellipse; scale_radii stores a single radius) is why the fast path keeps the exact circle handling only for isotropic scales — upstream could compute the true ellipse bbox and be more accurate than today.