Skip to content

Commit 8c87330

Browse files
committed
Add measure_obs: persist per-cell centroid/area/equivalent diameter into the annotating table
`measure_obs(sdata, element=None, ...)` computes one centroid, area and equivalent diameter per instance of a shapes or 2D-labels element and writes them, squidpy-style, into the annotating AnnData table: centroids to `obsm["spatial"]` (the canonical (n_obs, 2) array), area and equivalent diameter to `obs`. Values are stored in the element's intrinsic coordinates/units; equivalent diameter is `2*sqrt(area/pi)`. Labels use a streaming bincount aggregator that processes the raster block by block (one chunk plus O(n_labels) accumulators), so it stays out-of-core and scales to Xenium-size masks where a whole-array regionprops table would run out of memory; area (the per-label pixel count) is a free by-product. Shapes use shapely's vectorized centroid/area. The function is idempotent: outputs already present and current are not recomputed, a pre-existing `obsm["spatial"]` is trusted and never overwritten, and an instance-count change invalidates the cache. `inplace` follows the scanpy convention (mutate and return None, or operate on a deep copy and return it). Per-cell measurements require an annotating table to write into. Render-side wiring (routing `as_points` through these measurements for footprint dot sizing) is intentionally deferred to a follow-up PR.
1 parent 78d9233 commit 8c87330

3 files changed

Lines changed: 482 additions & 2 deletions

File tree

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
from ._palette import make_palette, make_palette_from_data
22
from .basic import PlotAccessor
3+
from .utils import measure_obs
34

45
__all__ = [
56
"PlotAccessor",
67
"make_palette",
78
"make_palette_from_data",
9+
"measure_obs",
810
]

‎src/spatialdata_plot/pl/utils.py‎

Lines changed: 363 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import spatialdata as sd
2727
from anndata import AnnData
2828
from cycler import Cycler, cycler
29+
from dask.array.core import slices_from_chunks
2930
from datashader.core import Canvas
3031
from geopandas import GeoDataFrame
3132
from matplotlib import colors, patheffects, rcParams
@@ -59,6 +60,7 @@
5960
from skimage.util import map_array
6061
from spatialdata import (
6162
SpatialData,
63+
deepcopy as sd_deepcopy,
6264
get_element_annotators,
6365
get_extent,
6466
get_values,
@@ -67,7 +69,14 @@
6769
)
6870
from spatialdata._core.query.relational_query import _locate_value
6971
from spatialdata._types import ArrayLike
70-
from spatialdata.models import Image2DModel, Labels2DModel, SpatialElement, get_table_keys
72+
from spatialdata.models import (
73+
Image2DModel,
74+
Labels2DModel,
75+
ShapesModel,
76+
SpatialElement,
77+
get_model,
78+
get_table_keys,
79+
)
7180
from spatialdata.transformations.operations import get_transformation
7281
from spatialdata.transformations.transformations import Scale, Translation
7382
from spatialdata.transformations.transformations import Sequence as TransformSequence
@@ -4417,3 +4426,356 @@ def _convert_alpha_to_datashader_range(alpha: float) -> float:
44174426
"""Convert alpha from the range [0, 1] to the range [0, 255] used in datashader."""
44184427
# prevent a value of 255, bc that led to fully colored test plots instead of just colored points/shapes
44194428
return min([254, alpha * 255])
4429+
4430+
4431+
# --- Per-cell measurements into the annotating table (centroid / area / equivalent diameter) ---
4432+
4433+
# squidpy/scanpy canonical key for per-cell coordinates: an N x 2 array, row-aligned to obs.
4434+
_CENTROID_OBSM_KEY = "spatial"
4435+
# Our provenance namespace in ``table.uns``. Measurements are stored in the element's *intrinsic*
4436+
# coordinates/units (coordinate-system-independent), so the marker records the resolution level
4437+
# and the instance count (to invalidate when cells are added), not a coordinate system.
4438+
_CENTROID_PROVENANCE_UNS_KEY = "spatialdata_plot"
4439+
_CENTROID_SCALE_LEVEL = "scale0" # rasters are reduced at full resolution (decision: exact)
4440+
4441+
4442+
def _transform_carrier(element: SpatialElement) -> Any:
4443+
"""Return the object carrying ``element``'s data (the ``scale0`` level for multiscale rasters)."""
4444+
if isinstance(element, DataTree):
4445+
return next(iter(element[_CENTROID_SCALE_LEVEL].values()))
4446+
return element
4447+
4448+
4449+
def _stream_label_centroid_stats(data: Any) -> tuple[ArrayLike, ArrayLike, ArrayLike, ArrayLike]:
4450+
"""Per-label ``(labels, mean_x_index, mean_y_index, area)`` via a streaming bincount aggregator.
4451+
4452+
Streams the raster block by block — one chunk in memory at a time for a dask array, a
4453+
bounded row-block at a time for a numpy array — accumulating per-label ``count`` (= area),
4454+
``sum_x`` and ``sum_y``. The reduction is additive, so it is exact across block boundaries
4455+
and uses only O(n_labels) memory regardless of raster size. Background label 0 is excluded.
4456+
"""
4457+
n_rows, n_cols = data.shape
4458+
is_dask = hasattr(data, "chunks") and not isinstance(data, np.ndarray)
4459+
if is_dask:
4460+
n_labels = int(data.max().compute()) + 1
4461+
block_slices = slices_from_chunks(data.chunks)
4462+
else:
4463+
data = np.asarray(data)
4464+
n_labels = int(data.max()) + 1
4465+
# bound the per-block coordinate-weight arrays to ~8M pixels
4466+
step = max(1, min(n_rows, (8 << 20) // max(1, n_cols)))
4467+
block_slices = [(slice(r0, min(r0 + step, n_rows)), slice(0, n_cols)) for r0 in range(0, n_rows, step)]
4468+
4469+
count = np.zeros(n_labels, dtype=np.float64)
4470+
sum_x = np.zeros(n_labels, dtype=np.float64)
4471+
sum_y = np.zeros(n_labels, dtype=np.float64)
4472+
for row_sl, col_sl in block_slices:
4473+
block = data[row_sl, col_sl]
4474+
block = np.asarray(block.compute() if hasattr(block, "compute") else block)
4475+
block_rows, block_cols = block.shape
4476+
flat = block.reshape(-1)
4477+
cols = np.tile(np.arange(col_sl.start, col_sl.start + block_cols, dtype=np.float64), block_rows)
4478+
rows = np.repeat(np.arange(row_sl.start, row_sl.start + block_rows, dtype=np.float64), block_cols)
4479+
count += np.bincount(flat, minlength=n_labels)
4480+
sum_x += np.bincount(flat, weights=cols, minlength=n_labels)
4481+
sum_y += np.bincount(flat, weights=rows, minlength=n_labels)
4482+
4483+
labels = np.flatnonzero(count)
4484+
labels = labels[labels != 0] # drop background and absent labels
4485+
return labels, sum_x[labels] / count[labels], sum_y[labels] / count[labels], count[labels]
4486+
4487+
4488+
def _compute_label_measurements(element: DataArray | DataTree) -> pd.DataFrame:
4489+
"""Per-label intrinsic centroid + pixel area via the streaming bincount aggregator.
4490+
4491+
Returns a DataFrame indexed by label value with columns ``["x", "y", "area"]`` in the
4492+
element's *intrinsic* coordinates; ``area`` is the pixel count, a free by-product of the
4493+
aggregator. The aggregator streams the (possibly huge, dask-backed) raster block by block,
4494+
holding only one chunk plus O(n_labels) accumulators, so it scales to Xenium-size masks
4495+
where ``regionprops`` (whole-array, per-label table) would run out of memory.
4496+
"""
4497+
raster = _transform_carrier(element)
4498+
labels, x_idx, y_idx, area = _stream_label_centroid_stats(raster.data)
4499+
4500+
# bincount gives mean 0-based pixel indices; map them onto the raster's intrinsic coordinate
4501+
# arrays (spatialdata uses pixel-center coords 0.5, 1.5, ... and possibly non-unit spacing).
4502+
def _index_to_coord(idx: ArrayLike, coord: ArrayLike) -> ArrayLike:
4503+
spacing = (coord[1] - coord[0]) if len(coord) > 1 else 1.0
4504+
return coord[0] + idx * spacing
4505+
4506+
xcoord = np.asarray(raster.coords["x"].values)
4507+
ycoord = np.asarray(raster.coords["y"].values)
4508+
return pd.DataFrame(
4509+
{
4510+
"x": _index_to_coord(x_idx, xcoord),
4511+
"y": _index_to_coord(y_idx, ycoord),
4512+
"area": np.asarray(area, dtype=float),
4513+
},
4514+
index=labels,
4515+
)
4516+
4517+
4518+
def _compute_element_measurements(sdata: SpatialData, element_name: str) -> pd.DataFrame:
4519+
"""One row per instance with intrinsic ``["x", "y", "area"]``, indexed by instance id.
4520+
4521+
Shapes use shapely's vectorized ``.centroid`` / ``.area`` (area in intrinsic squared units);
4522+
2D labels use the streaming bincount aggregator (area = pixel count). Other element types are
4523+
rejected explicitly. Area meaning therefore differs across element types but is consistent
4524+
within one element.
4525+
"""
4526+
element = sdata[element_name]
4527+
model = get_model(element)
4528+
if model is ShapesModel:
4529+
centroids = element.geometry.centroid
4530+
return pd.DataFrame(
4531+
{
4532+
"x": centroids.x.to_numpy(),
4533+
"y": centroids.y.to_numpy(),
4534+
"area": element.geometry.area.to_numpy(),
4535+
},
4536+
index=element.index,
4537+
)
4538+
if model is Labels2DModel:
4539+
return _compute_label_measurements(element)
4540+
raise NotImplementedError(
4541+
f"Measurement is only supported for shapes and 2D labels; "
4542+
f"element {element_name!r} is a {model.__name__}."
4543+
)
4544+
4545+
4546+
def _region_mask_and_keys(table: AnnData, element_name: str) -> tuple[str, ArrayLike]:
4547+
"""Return ``(instance_key, mask)`` for the rows of ``table`` that annotate ``element_name``."""
4548+
_, region_key, instance_key = get_table_keys(table)
4549+
mask = (table.obs[region_key].astype(str) == str(element_name)).to_numpy()
4550+
return instance_key, mask
4551+
4552+
4553+
def _valid_spatial_obsm(arr: ArrayLike, n_obs: int) -> bool:
4554+
"""Whether ``arr`` is a usable ``obsm["spatial"]``: a 2D ``(n_obs, 2)`` coordinate grid."""
4555+
return bool(arr.ndim == 2 and arr.shape == (n_obs, 2))
4556+
4557+
4558+
def _measurement_provenance(table: AnnData, element_name: str) -> dict[str, Any] | None:
4559+
"""Our measurement provenance marker for ``element_name``, or None if absent/malformed."""
4560+
root = table.uns.get(_CENTROID_PROVENANCE_UNS_KEY)
4561+
if not isinstance(root, dict):
4562+
return None
4563+
marker = root.get("centroids", {}).get(element_name)
4564+
return marker if isinstance(marker, dict) else None
4565+
4566+
4567+
def _record_measurement_provenance(table: AnnData, element_name: str, n: int) -> None:
4568+
"""Record ``(n, scale_level)`` for ``element_name`` so a later instance-count change invalidates."""
4569+
root = table.uns.get(_CENTROID_PROVENANCE_UNS_KEY)
4570+
if not isinstance(root, dict):
4571+
root = {}
4572+
table.uns[_CENTROID_PROVENANCE_UNS_KEY] = root
4573+
root.setdefault("centroids", {})[element_name] = {"n": int(n), "scale_level": _CENTROID_SCALE_LEVEL}
4574+
4575+
4576+
def _obs_region_finite(table: AnnData, key: str, mask: ArrayLike) -> bool:
4577+
"""Whether ``obs[key]`` exists and is finite for every row in ``mask`` (already populated)."""
4578+
if key not in table.obs:
4579+
return False
4580+
vals = np.asarray(table.obs[key].to_numpy(), dtype=float)[mask]
4581+
return bool(vals.size and np.isfinite(vals).all())
4582+
4583+
4584+
def _obsm_region_finite(table: AnnData, key: str, mask: ArrayLike) -> bool:
4585+
"""Whether ``obsm[key]`` is a valid ``(n_obs, 2)`` grid and finite for every row in ``mask``."""
4586+
if key not in table.obsm:
4587+
return False
4588+
arr = np.asarray(table.obsm[key])
4589+
if not _valid_spatial_obsm(arr, table.n_obs):
4590+
return False
4591+
region = arr[mask].astype(float)
4592+
return bool(region.size and np.isfinite(region).all())
4593+
4594+
4595+
def _write_obs_region(table: AnnData, mask: ArrayLike, key: str, values: ArrayLike) -> None:
4596+
"""Write ``values`` into ``obs[key]`` at ``mask`` rows; other rows stay/are NaN."""
4597+
if key in table.obs:
4598+
col = np.asarray(table.obs[key].to_numpy(), dtype=float)
4599+
else:
4600+
col = np.full(table.n_obs, np.nan, dtype=float)
4601+
col[mask] = np.asarray(values, dtype=float)
4602+
table.obs[key] = col
4603+
4604+
4605+
def _write_obsm_region(table: AnnData, mask: ArrayLike, key: str, xy: ArrayLike) -> None:
4606+
"""Write intrinsic centroids ``xy`` into ``obsm[key]`` at ``mask`` rows; other rows stay/are NaN.
4607+
4608+
Never reshapes an incompatible existing ``obsm[key]`` (e.g. a 3-column xyz array): that would
4609+
silently drop data, so it raises instead.
4610+
"""
4611+
if key in table.obsm:
4612+
existing = np.asarray(table.obsm[key])
4613+
if not _valid_spatial_obsm(existing, table.n_obs):
4614+
raise ValueError(
4615+
f"Refusing to overwrite obsm[{key!r}] with shape {existing.shape}; expected "
4616+
f"({table.n_obs}, 2). Remove it first if you want it replaced."
4617+
)
4618+
arr = existing.astype(float, copy=True)
4619+
else:
4620+
arr = np.full((table.n_obs, 2), np.nan, dtype=float)
4621+
arr[mask] = np.asarray(xy, dtype=float)
4622+
table.obsm[key] = arr
4623+
4624+
4625+
def _measure_into_table(
4626+
sdata: SpatialData,
4627+
element_name: str,
4628+
table_name: str,
4629+
*,
4630+
centroids: bool,
4631+
area: bool,
4632+
diameter: bool,
4633+
obsm_key: str,
4634+
area_key: str,
4635+
diameter_key: str,
4636+
force: bool,
4637+
) -> None:
4638+
"""Compute and persist the requested measurements for one element into its annotating table."""
4639+
table = sdata.tables[table_name]
4640+
instance_key, mask = _region_mask_and_keys(table, element_name)
4641+
if not mask.any():
4642+
raise ValueError(f"Table {table_name!r} does not annotate element {element_name!r} (no matching rows).")
4643+
n = int(mask.sum())
4644+
4645+
prov = _measurement_provenance(table, element_name)
4646+
stale = prov is not None and prov.get("n") != n # cells added/removed since we last wrote
4647+
want_centroids = centroids and (force or stale or not _obsm_region_finite(table, obsm_key, mask))
4648+
want_area = area and (force or stale or not _obs_region_finite(table, area_key, mask))
4649+
want_diameter = diameter and (force or stale or not _obs_region_finite(table, diameter_key, mask))
4650+
if not (want_centroids or want_area or want_diameter):
4651+
return # everything requested is already present and current
4652+
4653+
meas = _compute_element_measurements(sdata, element_name) # intrinsic [x, y, area]
4654+
aligned = meas.reindex(table.obs[instance_key].to_numpy()[mask])
4655+
if want_centroids:
4656+
_write_obsm_region(table, mask, obsm_key, aligned[["x", "y"]].to_numpy())
4657+
if want_area:
4658+
_write_obs_region(table, mask, area_key, aligned["area"].to_numpy())
4659+
if want_diameter:
4660+
_write_obs_region(table, mask, diameter_key, 2.0 * np.sqrt(aligned["area"].to_numpy() / np.pi))
4661+
_record_measurement_provenance(table, element_name, n)
4662+
4663+
4664+
def _measurable_elements(sdata: SpatialData) -> list[str]:
4665+
"""Shapes / 2D-labels elements with exactly one annotating table (where measurements persist)."""
4666+
names: list[str] = []
4667+
for name in list(sdata.shapes.keys()) + list(sdata.labels.keys()):
4668+
if len(get_element_annotators(sdata, name)) != 1:
4669+
continue
4670+
if get_model(sdata[name]) in (ShapesModel, Labels2DModel):
4671+
names.append(name)
4672+
return names
4673+
4674+
4675+
def _resolve_measure_table(sdata: SpatialData, element_name: str, table_name: str | None) -> str:
4676+
"""Resolve the single annotating table for ``element_name`` (where measurements are written)."""
4677+
if table_name is not None:
4678+
if table_name not in sdata.tables:
4679+
raise KeyError(f"Table {table_name!r} not found in `sdata.tables`.")
4680+
return table_name
4681+
annotators = sorted(get_element_annotators(sdata, element_name))
4682+
if not annotators:
4683+
raise ValueError(
4684+
f"Element {element_name!r} has no annotating table; per-cell measurements need a table "
4685+
f"to write into. Pass `table_name=` or annotate the element first."
4686+
)
4687+
if len(annotators) > 1:
4688+
raise ValueError(
4689+
f"Element {element_name!r} is annotated by multiple tables ({', '.join(annotators)}); "
4690+
f"pass `table_name=` to pick one."
4691+
)
4692+
return annotators[0]
4693+
4694+
4695+
def measure_obs(
4696+
sdata: SpatialData,
4697+
element: str | None = None,
4698+
*,
4699+
table_name: str | None = None,
4700+
centroids: bool = True,
4701+
area: bool = True,
4702+
diameter: bool = True,
4703+
obsm_key: str = _CENTROID_OBSM_KEY,
4704+
area_key: str = "area",
4705+
diameter_key: str = "equivalent_diameter",
4706+
force: bool = False,
4707+
inplace: bool = True,
4708+
) -> SpatialData | None:
4709+
"""Measure per-cell centroids, area and equivalent diameter into an element's annotating table.
4710+
4711+
Computes one centroid, area and equivalent diameter per instance of a shapes or 2D-labels
4712+
element and writes them, squidpy-style, into the annotating :class:`~anndata.AnnData` table:
4713+
centroids go to ``table.obsm[obsm_key]`` (an ``(n_obs, 2)`` array, the canonical
4714+
``obsm["spatial"]``), area and diameter to ``table.obs``. Values are stored in the element's
4715+
*intrinsic* coordinates/units (coordinate-system independent). Labels area is the pixel count;
4716+
shapes area is ``geometry.area``; equivalent diameter is ``2 * sqrt(area / pi)``. Persisting
4717+
them once lets later renders (and downstream tools such as squidpy) reuse them instead of
4718+
recomputing.
4719+
4720+
The function is idempotent: outputs already present and current for an element are not
4721+
recomputed. A pre-existing ``obsm[obsm_key]`` (e.g. from a reader) is trusted and never
4722+
overwritten. Adding cells (which leaves new rows unmeasured) triggers a recompute; pass
4723+
``force=True`` to recompute unconditionally.
4724+
4725+
Per-cell measurements need a table to write into, so an element without an annotating table
4726+
cannot be measured (this raises). The function never densifies the raster: the label path
4727+
streams it block by block, scaling to Xenium-size masks.
4728+
4729+
Parameters
4730+
----------
4731+
sdata
4732+
The ``SpatialData`` object holding the element and its annotating table.
4733+
element
4734+
Name of the shapes/2D-labels element to measure. If ``None``, every shapes/2D-labels
4735+
element that has exactly one annotating table is measured.
4736+
table_name
4737+
Name of the annotating table to write into. If ``None``, it is inferred from the element's
4738+
annotators (an error is raised when there are zero or several).
4739+
centroids, area, diameter
4740+
Which measurements to compute/persist. At least one must be ``True``.
4741+
obsm_key
4742+
``obsm`` key for the centroids (default ``"spatial"``, the squidpy convention).
4743+
area_key, diameter_key
4744+
``obs`` column names for area and equivalent diameter.
4745+
force
4746+
Recompute and overwrite even when current values are already present.
4747+
inplace
4748+
If ``True`` (default), mutate ``sdata``'s table in place and return ``None``. If ``False``,
4749+
operate on a deep copy and return the modified ``SpatialData``.
4750+
4751+
Returns
4752+
-------
4753+
``None`` if ``inplace`` is ``True``, otherwise the modified deep-copied ``SpatialData``.
4754+
"""
4755+
if not (centroids or area or diameter):
4756+
raise ValueError("Nothing to measure: set at least one of `centroids`, `area`, `diameter` to True.")
4757+
target = sdata if inplace else sd_deepcopy(sdata)
4758+
if element is None:
4759+
names = _measurable_elements(target)
4760+
if not names:
4761+
raise ValueError(
4762+
"No shapes/2D-labels element with a single annotating table was found; nothing to "
4763+
"measure. Pass an explicit `element=` (and `table_name=` if ambiguous)."
4764+
)
4765+
else:
4766+
names = [element]
4767+
for name in names:
4768+
resolved_table = _resolve_measure_table(target, name, table_name)
4769+
_measure_into_table(
4770+
target,
4771+
name,
4772+
resolved_table,
4773+
centroids=centroids,
4774+
area=area,
4775+
diameter=diameter,
4776+
obsm_key=obsm_key,
4777+
area_key=area_key,
4778+
diameter_key=diameter_key,
4779+
force=force,
4780+
)
4781+
return None if inplace else target

0 commit comments

Comments
 (0)