Fast cell rendering: render_shapes/labels(as_points=True) + squidpy centroid caching#703
Open
timtreis wants to merge 7 commits into
Open
Fast cell rendering: render_shapes/labels(as_points=True) + squidpy centroid caching#703timtreis wants to merge 7 commits into
timtreis wants to merge 7 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 75.96% 76.45% +0.49%
==========================================
Files 14 14
Lines 4156 4332 +176
Branches 964 996 +32
==========================================
+ Hits 3157 3312 +155
- Misses 647 664 +17
- Partials 352 356 +4
🚀 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.
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
Adds a fast rendering mode that draws each cell as a dot at its centroid instead of its full geometry / rasterized mask — a large speedup when only cell location matters (the squidpy
spatial_scatteridea, integrated into spatialdata-plot's API).What's in this PR
1. Centroid extraction + squidpy
obsm["spatial"]caching (utils.py)_compute_element_centroids: shapes → shapely.centroid; 2D labels →regionprops(≈750–1860× faster thanget_centroidson rasters); other types →NotImplementedError._get_or_compute_centroids: reuses a pre-existingobsm["spatial"](loader/user-provided, trusted as intrinsic) or computes + writes it back with{n, scale_level}provenance inuns; invalidated on instance-count change. Never clobbers an incompatible array.2.
as_points=Truerendering (render.py,basic.py,render_params.py)as_points: bool+size:onrender_shapesandrender_labels._render_centroids_as_points(built on the 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 apply step differs (scatter vs patches/imshow)._get_or_compute_centroidsreindexed toinstance_id. Positions verified identical tosd.get_centroids.outline_*/shape(shapes) andcontour_px/outline_*(labels) are ignored with an info log.Verification
get_centroidsfor shapes and labels.as_points=False) output is byte-identical tomain._scatter_pointsextraction byte-identical; coordinate-system-independence, cache round-trip/provenance, staleness, unsupported-type rejection all unit-tested.Known follow-ups (deliberately not in this PR — flagged for holistic review)
as_pointsat very large cell counts (currently always matplotlib scatter;method=is not yet honored for the dot backend). Needs a_datashader_pointsextraction mirroring_scatter_points.show()operates onself._copy(), so theobsm["spatial"]write doesn't reach the user's original object across separate.show()calls. Making the "magic cache" persist needs writing to the source object — ashow()-architecture decision.test_plot_*baselines foras_points(functional/position tests included here; baselines to be generated from CI).Design + the 12 locked decisions:
plans/fast-render-cells.md.