Fast cell rendering: render_shapes/labels(as_points=True) + axis-aligned extent#703
Open
timtreis wants to merge 22 commits into
Open
Fast cell rendering: render_shapes/labels(as_points=True) + axis-aligned extent#703timtreis wants to merge 22 commits into
timtreis wants to merge 22 commits into
Conversation
… helper Infrastructure for an upcoming "render cells as centroid points" fast mode (no user-facing render option yet). Phase 0 — shared scatter primitive: - Extract `_scatter_points(ax, x, y, color_vector, ...)` from `_render_points`'s matplotlib branch; `_render_points` now calls it. Byte-identical output (verified vs main on categorical and continuous point renders). This is the reuse seam the fast mode will draw through. Phase 1 — centroid + caching core (headless, fully unit-tested): - `_compute_element_centroids` / `_compute_label_centroids`: per-instance centroids in a coordinate system. Shapes use spatialdata's vectorized `get_centroids`; labels use skimage `regionprops` (the per-label reduction is orders of magnitude faster than `get_centroids` on rasters), mapped onto the raster's intrinsic coordinate arrays so it reproduces `get_centroids` exactly (incl. the pixel-center 0.5 offset) then transformed to the target CS. - `_get_or_compute_centroids`: reuses/persists centroids via the squidpy convention. A pre-existing `obsm["spatial"]` (loader/user-provided) is trusted as the cells' locations; otherwise centroids are computed and written back into the annotating table's `obsm["spatial"]` with a coordinate-system provenance marker in `uns`, so later renders are instant. Reads run before writes, so a valid existing cache is reused rather than clobbered; an incompatible existing `obsm["spatial"]` is never overwritten; the cache is invalidated when the requested coordinate system differs. - Tests: shapes/labels centroids match `get_centroids`; cache round-trip + provenance; CS invalidation; pre-existing obsm trusted; no-table compute path; `cache=False` writes nothing.
… provenance Cleanup from /simplify (no behavioral change): - Extract `_region_mask_and_keys(table, element)` used by both read and write, removing the duplicated `get_table_keys` + O(n_obs) `region_key`-string-cast mask that was computed twice per cold call. - Read path: validate shape on the raw obsm array and cast only the masked subset to float, instead of casting the whole `obsm["spatial"]` on every cache hit (the hot path). - Write path: coerce a non-dict `uns["spatialdata_plot"]` instead of early returning after `obsm` was already mutated, so obsm and the provenance marker are always written together (no half-write). - Drop the dead `"key"` provenance field (constant, never read back). - Rename the misleading `table` local (held a table *name*) in `_get_or_compute_centroids`.
Refactor the centroid cache to store element-*intrinsic* coordinates and
transform to the render coordinate system on demand, instead of caching
coords already mapped into one coordinate system. Decisions from design pass:
- Intrinsic storage: one `obsm["spatial"]` cache serves every coordinate
system (proven equivalent to per-CS computation). `_compute_element_centroids`
returns intrinsic coords (shapes via shapely `.centroid`, labels via
`regionprops`); `_centroids_to_coordinate_system` maps them to the requested
CS via the element's transform; `_get_or_compute_centroids` reads/computes
intrinsic then transforms on return.
- Provenance records `{n, scale_level}` (no coordinate system). Cache is
invalidated when the region's instance count changes (cells added/removed),
not on CS change.
- A pre-existing `obsm["spatial"]` (loader/user-provided) is trusted as the
cells' intrinsic locations and transformed to the render CS.
- Exhaustive model dispatch: shapes and 2D labels supported; other element
types raise NotImplementedError.
- Labels are reduced at full resolution (scale0).
Drops the now-unused `get_centroids` import; adds `ShapesModel`. Tests updated:
parametrized shapes/labels match `get_centroids`; new coordinate-system-
independence test (one cache, two CS); staleness-by-instance-count; trusted
pre-existing obsm; unsupported-type rejection.
…e import, shared obsm gate Cleanup from /simplify: - `_centroids_to_coordinate_system` ran `PointsModel.parse` + a dask `transform(...).compute()` round trip on every call, including cache hits (~19 ms fixed floor, ~100 ms at 1M cells) — defeating the cache. Replace with `to_affine_matrix` + a plain numpy matmul: numerically identical (verified against the dask path for multiple coordinate systems), ~80-140x faster, and it removes the private-API import `spatialdata._core.operations.transform` (a repo non-negotiable) plus the now-unused `PointsModel`. - Widen `_transformable_raster` -> `_transform_carrier` to accept any element (rasters -> scale0, others as-is), dropping the `isinstance` branch in `_centroids_to_coordinate_system`. - Extract `_valid_spatial_obsm(arr, n_obs)` shared by the read and write paths, reconciling their previously divergent obsm-shape checks (read accepted >=2 columns, write required exactly 2) so they cannot drift.
`render_shapes(..., as_points=True)` and `render_labels(..., as_points=True)` draw one dot per cell at its centroid instead of the full geometry / rasterized mask — a large speedup when only cell location matters. New `size=` controls the marker size. - Shared `_render_centroids_as_points` draws the scatter (via `_scatter_points`) and the legend/colorbar. The per-cell color vector is the *same* one the geometry/raster path computes (`_set_color_source_vec`), so colors match the full rendering exactly; only the apply step (scatter vs patches/imshow) differs. - Shapes: centroids from shapely `.centroid` of the (filtered) geometry, positionally aligned to the color vector, drawn in intrinsic coords via the element transform. Labels: centroids from `_get_or_compute_centroids` (regionprops, fast) reindexed to `instance_id`. Positions verified identical to `sd.get_centroids`. - `as_points` short-circuits before the geometry/raster path; outline_*, shape (shapes) and contour_px, outline_* (labels) are ignored with an info log. - Default (`as_points=False`) output is byte-identical to main. Tests: non-visual checks that centroids land exactly on `get_centroids` for both element types and that outline/shape are ignored without error. Note: as_points currently always uses the matplotlib scatter backend; routing through datashader for very large cell counts (and persisting the obsm cache to the user's object rather than show()'s working copy) are follow-ups.
timtreis
changed the title
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #703 +/- ##
==========================================
+ Coverage 76.26% 76.54% +0.27%
==========================================
Files 14 14
Lines 4327 4425 +98
Branches 1006 1029 +23
==========================================
+ Hits 3300 3387 +87
- Misses 667 672 +5
- Partials 360 366 +6
🚀 New features to boost your workflow:
|
`render_labels(element, as_points=True)` with no color crashed: `instance_id` (the raster's unique values) includes the background label `0`, which has no centroid, and the literal/no-color color vector is sized to the raster (not per-instance), so `ax.scatter` got mismatched `c` vs `x`/`y`. Drop the background label from the rendered instances and align the per-cell color: for data-driven color the vector is already per-instance and is subset to match; for the literal/no-color path it is replaced with one na/literal color per centroid. Data-driven (categorical/continuous) renders are unchanged and still land exactly on `get_centroids`. Adds a regression test for the no-color labels case.
…ator Replace the `regionprops` reduction in `_compute_label_centroids` with an additive bincount aggregator that streams the labels raster block by block — one dask chunk (or bounded numpy row-block) in memory at a time — accumulating per-label `count`/`sum_x`/`sum_y`. This is what makes the feature usable at Xenium scale: - Out-of-core: peak memory is one chunk + O(n_labels) accumulators, NOT the whole raster (measured: 9 MB peak streaming a 268 MB mask). `regionprops` needs the full array materialized and OOMs on large morphology masks. - Scales in cell count: 500k+ labels are just array indexing (562k labels in ~1.4 s with 13.5 MB of accumulators); `regionprops`' per-label table does not. - Faster than `regionprops` (~1.3-1.6x) on in-memory rasters. - Exact across chunk boundaries (additive reduction) — verified numpy == dask-chunked, and identical to `sd.get_centroids`. - `count` is the cell area, a free by-product (ready for footprint-based dot sizing). Drops the `regionprops_table` import; adds `slices_from_chunks`. Adds a unit test locking the chunk-exact, out-of-core, area-correct behavior. Note: the chunk loop is currently sequential; parallelizing the per-chunk partials (dask map_blocks + tree-reduce) is a future speedup.
#705 (measure_obs) merged to main, superseding #703's pre-extraction centroid block. This merge takes main's utils.py + test_utils.py wholesale (zero #703 delta there) and rewires the labels as_points path onto main's primitive: - render.py labels branch: drop the deleted `_get_or_compute_centroids` + the whole cache layer; compute full-resolution (scale0) centroids via main's `_compute_element_measurements`, and draw them with a transform built from the *same* scale0 element (`_prepare_transformation(_get_top_data_array(...))`) — so positions are independent of any rasterization applied to the rendered `label`. Coerce point_ids to the label dtype so str/object instance ids (e.g. Xenium readers) align instead of silently reindexing to NaN. - Net: utils.py == main (the 242-line #703 centroid block + cache layer is gone); the surviving diff vs main is just the as_points feature. - Tests: added a non-identity-transform regression test asserting the dots land at the cells' coordinate-system positions in display space (the guard for the transform pairing); existing as_points position tests still pass. Shapes branch unchanged (already intrinsic centroid + trans_data, positionally aligned to its post-filter color vector).
…ked fields The fast-mode draw wrapper read cmap/size/alpha/zorder/colorbar/colorbar_params straight off render_params at both call sites. Pass the render_params dataclass instead and read them internally; only the genuinely per-branch values (x/y/color vectors, norm, na_color, transform, adata, palette, col_for_color) stay explicit. Signature 20->15 params; both as_points call sites lose the restated render_params plumbing.
Collapse the per-centroid color alignment (drop the unconditional asarray pre-init and the nested None-guard into a single if/else + ternary), and drop a redundant np.asarray around the already-ndarray point_ids in the reindex. Behavior unchanged.
Benchmark showed labels as_points was 16-34x SLOWER than the normal render because it recomputed full-resolution scale0 centroids while the normal path downsamples + imshows. Compute centroids on the already-rendered (downsampled) raster and draw with its trans_data instead: 671k cells goes 19.7s -> 0.88s (now ~parity with imshow). Centroid error is sub-pixel at display resolution; position tests updated to display-space within a few-px tolerance.
…gned show()'s get_extent(exact=True) transforms EVERY shapes/points geometry into the coordinate system just to take a bounding box - the dominant cost for large shape collections (~85% of a 5.5M-shape render), unrelated to what is drawn. Add a self-contained get_extent_fast (drop-in for spatialdata.get_extent) that, for axis-aligned transforms (scale/flip/90deg/swap + translation - all real Visium/Xenium data), transforms only the 4 bounding-box corners and reads the intrinsic bounds vectorised (no per-geometry .apply(is_empty)). Proven identical to the exact extent for such transforms (spatialdata's own get_extent docstring notes this); falls back to spatialdata's get_extent for rotation/shear, for anisotropically-scaled circles (radius->ellipse divergence), and for images/labels (whose get_extent is already a cheap corner transform). Measured on real Visium HD render_shapes(as_points=True): 351k 8.5s->1.8s (4.7x), 5.48M 129s->24s (5.5x); the full (non-as_points) render benefits too. The whole block is isolated so it can be lifted into spatialdata's get_extent verbatim (see issue #706). Tests assert it matches get_extent across scale/flip/ rotation/shear for circles and polygons.
timtreis
force-pushed
the
feat/centroid-scatter-helper
branch
from
35c8f1f to
57bcf8d
Compare
…ompute, underscore) Follow-up cleanups from review of the get_extent fast path: - fold _intrinsic_xy_bounds into _element_extent_fast (drop duplicate get_model / geom_type) - geom.bounds -> geom.total_bounds (C-level union, avoids an Nx4 alloc on large collections) - batch the four points min/max into one dask.compute - rename get_extent_fast -> _get_extent_fast (internal helper; matches sibling underscored helpers) No behavior change; output identical for axis-aligned and rotated/sheared elements.
The datashader shapes canvas sized itself via spatialdata's exact get_extent, transforming every geometry (O(N)) just for a bounding box -- the same cost _get_extent_fast already removed from show()'s axis limits, in a second code path. Reuse _element_extent_fast (corner transform for axis-aligned elements; None -> exact get_extent fallback for rotation/shear), so the result is pixel-identical and the per-geometry pass is skipped for the common case. Measured on Curio (69,713 colored shapes): render_shapes 3146 ms -> 1675 ms (1.88x), on top of the get_extent_fast win already in show(). Mirrors _datashader_canvas_from_dataframe, which already avoids get_extent for the points path.
…de review) From the perf-stack review of labels as_points: - No color column collapsed every dot to a single na_color; now each cell gets a distinct random colour, matching the mask path's _map_color_seg Case C. Adds a color assertion to the regression test. - The rasterize drop-filter only ran when a color column was set, so as_points could emit dots at NaN positions (or drop cells) when rasterization removed labels; extend it to as_points so point ids stay within the rendered raster. - ax.scatter autoscales the Normalize in place; copy the shared cmap_params.norm (the shapes path already does). - render_shapes/render_labels(as_points=True) did not validate size; add the points-path numeric/positive check so size=-5 / 'big' raise an actionable error instead of a raw matplotlib failure.
Extract the duplicated as_points size-validation block from render_shapes and render_labels into a shared _validate_as_points_size helper. Defer the np.asarray(color_vector) conversion in the centroid color path to the only branch that uses it.
…to feat/centroid-scatter-helper # Conflicts: # tests/pl/test_utils.py
_get_extent_fast already reads get_transformation(element, get_all=True) for the coordinate-system membership check; pass it through so the fast path does not re-fetch it per element. The optional kwarg leaves the datashader-canvas call site unchanged.
_element_extent_fast returned a NaN extent for an element whose geometries are all empty, which silently poisoned the union in _get_extent_fast (blank axes) instead of raising spatialdata's clear 'empty collection' error. Guard against non-finite bounds and fall back. Also trims the verbose extent-block comment and the scatter/centroid docstrings (net -12 LOC, no behavior change).
timtreis
changed the title
Four PlotTester baselines: render_shapes/labels(as_points=True) at a base size and a larger size (size = scatter marker area). Baselines to be generated from CI.
…l tests Rendered on hatch-test.py3.11-stable; verified dots land at shape/label centroids, colored by instance_id (labels), larger at size=600.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Two related rendering speedups for shapes/labels:
as_points=Truefast mode — draw each cell as a dot at its centroid instead of its full geometry / rasterized mask. A large speedup when only cell location matters (the squidpyspatial_scatteridea, in spatialdata-plot's API).pl.show()no longer transforms every geometry just to size the axes when the element's transform is axis-aligned.What's in this PR
as_points=Truerendering (render.py,basic.py,render_params.py)as_points: bool+size:onrender_shapesandrender_labels._render_centroids_as_points(built on an extracted_scatter_points) draws the dots + legend/colorbar. The per-cell color vector is the same one the geometry/raster path computes, so colors match the full rendering exactly — only the draw step differs (scatter vs patches/imshow).outline_*/shape(shapes) andcontour_px/outline_*(labels) are ignored with an info log.Fast axis-aligned extent (
utils.py)_get_extent_fastis a drop-in forget_extent(sdata, ...): for shapes/points under an axis-aligned transform (scale/flip/90°-rotation/axis-swap + translation) it transforms 4 corners instead of N geometries; rotation/shear, images, labels, and degenerate/empty elements fall back toget_extent. Same result, same union semantics.Verification
sd.get_centroidsfor both shapes and labels (incl. under a non-identity transform)._get_extent_fastmatchesget_extent(exact=True)across scale/flip/rotation/shear/translation for circles and polygons (parametrized test).as_points=False) output is unchanged vsmain._scatter_pointsextraction is byte-identical forrender_points.Known follow-ups (not in this PR)
as_pointsat very large cell counts (currently always matplotlib scatter).test_plot_*baselines foras_points(position/functional tests are included here; baselines to be generated from CI).