Skip to content

Commit 398fa05

Browse files
committed
Address review blockers in measure_obs (#1-#5, #7)
- #1 no-clobber: a populated obsm["spatial"] (reader- or prior-call-provided) is no longer overwritten — warn and skip that element's centroids. Coords stay in the element's intrinsic pixel space (documented); area/diameter still overwrite our own columns. Restores the per-element finiteness guard. - #2 unmatched instance ids: instances annotated in the table but absent from the element (e.g. str-vs-int id dtype mismatch) now warn instead of silently writing NaN. - #3 float-dtype labels: the dense relabelling bincounts integer searchsorted indices, never the raw labels, so a float-typed (integer-valued) mask no longer crashes np.bincount. - #4 atomic writes: validate every obs target (non-numeric column collision) before the first mutation, so a bad column never leaves a half-written table. - #5 O(n_labels) memory: relabel labels to a dense 0..k-1 range, so the aggregator's memory scales with the number of distinct labels, not the maximum label id (sparse/global ids no longer blow it up). Single max() pass replaced by a unique() pass; same pass count. - #7 circle area: dispatch on geometry TYPE (all-Point) rather than the presence of a "radius" column, so a polygon element carrying a radius column uses geometry.area, not pi*r**2. Tests updated to the new contract and extended for each fix.
1 parent 556ba6c commit 398fa05

2 files changed

Lines changed: 174 additions & 44 deletions

File tree

‎src/spatialdata_plot/pl/utils.py‎

Lines changed: 107 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -4453,57 +4453,72 @@ def _stream_label_centroid_stats(data: Any) -> tuple[ArrayLike, ArrayLike, Array
44534453
44544454
Streams the raster block by block — one chunk in memory at a time for a dask array, a
44554455
bounded row-block at a time for a numpy array — accumulating per-label ``count`` (= area),
4456-
``sum_x`` and ``sum_y``. The reduction is additive, so it is exact across block boundaries
4457-
and uses only O(n_labels) memory regardless of raster size. Background label 0 is excluded.
4456+
``sum_x`` and ``sum_y``. Labels are relabelled to a dense ``0..k-1`` range, so memory is
4457+
O(number of distinct labels), independent of the raster size *and* of the label-id magnitude
4458+
(sparse/global ids do not blow it up). The reduction is additive, so it is exact across block
4459+
boundaries. Background label 0 is excluded.
44584460
"""
44594461
n_rows, n_cols = data.shape
4460-
is_dask = hasattr(data, "chunks") and not isinstance(data, np.ndarray)
4461-
if is_dask:
4462-
n_labels = int(data.max().compute()) + 1
4462+
if hasattr(data, "chunks"): # dask
44634463
block_slices = slices_from_chunks(data.chunks)
44644464
else:
44654465
data = np.asarray(data)
4466-
n_labels = int(data.max()) + 1
44674466
# bound the per-block coordinate-weight arrays to ~8M pixels
44684467
step = max(1, min(n_rows, (8 << 20) // max(1, n_cols)))
44694468
block_slices = [(slice(r0, min(r0 + step, n_rows)), slice(0, n_cols)) for r0 in range(0, n_rows, step)]
44704469

4471-
count = np.zeros(n_labels, dtype=np.float64)
4472-
sum_x = np.zeros(n_labels, dtype=np.float64)
4473-
sum_y = np.zeros(n_labels, dtype=np.float64)
4474-
for row_sl, col_sl in block_slices:
4470+
def _load(row_sl: slice, col_sl: slice) -> np.ndarray:
44754471
block = data[row_sl, col_sl]
44764472
block = np.asarray(block.compute() if hasattr(block, "compute") else block)
4473+
# label rasters are integer-valued even when stored as float; cast so np.unique/searchsorted
4474+
# stay integer (and the dense bincount indices below are always int).
4475+
return block.astype(np.int64) if block.dtype.kind == "f" else block
4476+
4477+
# Pass 1: the sorted set of present label values -> dense relabelling (keeps memory O(n_labels)).
4478+
uniq = np.zeros(0, dtype=np.int64)
4479+
for row_sl, col_sl in block_slices:
4480+
uniq = np.union1d(uniq, np.unique(_load(row_sl, col_sl)))
4481+
k = uniq.size
4482+
4483+
# Pass 2: additive per-(dense-)label count / sum_x / sum_y.
4484+
count = np.zeros(k)
4485+
sum_x = np.zeros(k)
4486+
sum_y = np.zeros(k)
4487+
for row_sl, col_sl in block_slices:
4488+
block = _load(row_sl, col_sl)
44774489
block_rows, block_cols = block.shape
4478-
flat = block.reshape(-1)
4490+
idx = np.searchsorted(uniq, block.reshape(-1)) # dense 0..k-1 indices (always int)
44794491
cols = np.tile(np.arange(col_sl.start, col_sl.start + block_cols, dtype=np.float64), block_rows)
44804492
rows = np.repeat(np.arange(row_sl.start, row_sl.start + block_rows, dtype=np.float64), block_cols)
4481-
count += np.bincount(flat, minlength=n_labels)
4482-
sum_x += np.bincount(flat, weights=cols, minlength=n_labels)
4483-
sum_y += np.bincount(flat, weights=rows, minlength=n_labels)
4493+
count += np.bincount(idx, minlength=k)
4494+
sum_x += np.bincount(idx, weights=cols, minlength=k)
4495+
sum_y += np.bincount(idx, weights=rows, minlength=k)
44844496

4485-
labels = np.flatnonzero(count)
4486-
labels = labels[labels != 0] # drop background and absent labels
4487-
return labels, sum_x[labels] / count[labels], sum_y[labels] / count[labels], count[labels]
4497+
keep = uniq != 0 # drop background; every kept label has count >= 1
4498+
return uniq[keep], sum_x[keep] / count[keep], sum_y[keep] / count[keep], count[keep]
44884499

44894500

44904501
def _compute_element_measurements(sdata: SpatialData, element_name: str) -> pd.DataFrame:
44914502
"""One row per instance with intrinsic ``["x", "y", "area"]``, indexed by instance id.
44924503
4493-
Shapes use shapely's vectorized centroid; circles (``Point`` + ``radius``) have ``area = pi*r**2``
4494-
(shapely ``.area`` is 0 for them), polygons use the geometric area. 2D labels use the streaming
4495-
bincount aggregator (``area`` = pixel count) — it holds one chunk plus O(n_labels) accumulators,
4496-
so it scales to Xenium-size masks where a whole-array ``regionprops`` table would run out of
4497-
memory. Area meaning differs across element types but is consistent within one element.
4504+
Shapes use shapely's vectorized centroid; circles (``Point`` geometry + ``radius``) have
4505+
``area = pi*r**2`` (shapely ``.area`` is 0 for them), polygons use the geometric area. 2D labels
4506+
use the streaming bincount aggregator (``area`` = pixel count) — it holds one chunk plus
4507+
O(n_labels) accumulators, so it scales to Xenium-size masks where a whole-array ``regionprops``
4508+
table would run out of memory. Area meaning differs across element types but is consistent within
4509+
one element.
44984510
"""
44994511
element = sdata[element_name]
45004512
model = get_model(element)
45014513
if model is ShapesModel:
4502-
centroids = element.geometry.centroid
4503-
if "radius" in element.columns: # circles are Points; shapely .area is 0, so use pi * r**2
4514+
geometry = element.geometry
4515+
centroids = geometry.centroid
4516+
# Dispatch on geometry TYPE, not a column name: circles are Point geometries (shapely .area
4517+
# is 0 for them) with a radius -> pi*r**2; everything else uses the true geometric area.
4518+
if (geometry.geom_type == "Point").all():
45044519
area = np.pi * np.asarray(element["radius"], dtype=float) ** 2
45054520
else:
4506-
area = element.geometry.area.to_numpy()
4521+
area = geometry.area.to_numpy()
45074522
xy = {"x": centroids.x.to_numpy(), "y": centroids.y.to_numpy(), "area": area}
45084523
return pd.DataFrame(xy, index=element.index)
45094524
if model is Labels2DModel:
@@ -4535,6 +4550,26 @@ def _valid_spatial_obsm(arr: ArrayLike, n_obs: int) -> bool:
45354550
return bool(arr.ndim == 2 and arr.shape == (n_obs, 2))
45364551

45374552

4553+
def _obsm_region_finite(table: AnnData, key: str, mask: ArrayLike) -> bool:
4554+
"""Whether ``obsm[key]`` already holds finite coords for every ``mask`` row (already populated)."""
4555+
if key not in table.obsm:
4556+
return False
4557+
arr = np.asarray(table.obsm[key])
4558+
if not _valid_spatial_obsm(arr, table.n_obs):
4559+
return False
4560+
region = arr[mask].astype(float)
4561+
return bool(region.size and np.isfinite(region).all())
4562+
4563+
4564+
def _check_obs_numeric(table: AnnData, key: str) -> None:
4565+
"""Raise if ``obs[key]`` exists but is non-numeric, before any mutation (avoids half-writes)."""
4566+
if key in table.obs and not is_numeric_dtype(table.obs[key]):
4567+
raise ValueError(
4568+
f"Cannot write measurements into obs[{key!r}]: the existing column is "
4569+
f"{table.obs[key].dtype} (not numeric). Pass a different key or drop the column first."
4570+
)
4571+
4572+
45384573
def _write_obs_region(table: AnnData, mask: ArrayLike, key: str, values: ArrayLike) -> None:
45394574
"""Write ``values`` into ``obs[key]`` at ``mask`` rows; other rows stay/are NaN."""
45404575
if key in table.obs:
@@ -4579,14 +4614,46 @@ def _measure_into_table(
45794614
) -> None:
45804615
"""Compute and write the requested measurements for one element into its annotating table.
45814616
4582-
Overwrites existing values for the element's rows. A table may annotate several elements, so
4583-
only the rows belonging to ``element_name`` are touched.
4617+
Only the rows belonging to ``element_name`` are touched (a table may annotate several elements).
4618+
Centroids already present for those rows (e.g. reader-provided ``obsm["spatial"]``) are not
4619+
overwritten; ``area``/``diameter`` overwrite our own columns. All targets are validated before
4620+
the first write, so a bad column never leaves the table half-written.
45844621
"""
45854622
table = sdata.tables[table_name]
45864623
instance_key, mask = _region_mask_and_keys(table, element_name)
45874624
if not mask.any():
45884625
raise ValueError(f"Table {table_name!r} does not annotate element {element_name!r} (no matching rows).")
4589-
meas = _compute_element_measurements(sdata, element_name).reindex(table.obs[instance_key].to_numpy()[mask])
4626+
4627+
# #1: never clobber centroids already populated for this element's rows (reader/prior-call coords).
4628+
if centroids and _obsm_region_finite(table, obsm_key, mask):
4629+
warnings.warn(
4630+
f"obsm[{obsm_key!r}] is already populated for element {element_name!r}; not overwriting "
4631+
f"its centroids. Remove it or use a different `obsm_key` to recompute.",
4632+
UserWarning,
4633+
stacklevel=3,
4634+
)
4635+
centroids = False
4636+
if not (centroids or area or diameter):
4637+
return
4638+
4639+
keys = table.obs[instance_key].to_numpy()[mask]
4640+
meas = _compute_element_measurements(sdata, element_name).reindex(keys)
4641+
# #2: instance ids annotated in the table but absent from the element reindex to NaN -> warn.
4642+
missing = int(meas[["x", "y"]].isna().any(axis=1).sum())
4643+
if missing:
4644+
warnings.warn(
4645+
f"{missing}/{len(keys)} instances annotated for {element_name!r} have no match in the "
4646+
f"element (instance-id dtype mismatch, e.g. str vs int?); writing NaN for them.",
4647+
UserWarning,
4648+
stacklevel=3,
4649+
)
4650+
4651+
# #4: validate every obs target up front so an existing non-numeric column raises before any write.
4652+
if area:
4653+
_check_obs_numeric(table, area_key)
4654+
if diameter:
4655+
_check_obs_numeric(table, diameter_key)
4656+
45904657
if centroids:
45914658
_write_obsm_region(table, mask, obsm_key, meas[["x", "y"]].to_numpy())
45924659
if area:
@@ -4645,16 +4712,18 @@ def measure_obs(
46454712
element and writes them, squidpy-style, into the annotating :class:`~anndata.AnnData` table:
46464713
centroids go to ``table.obsm[obsm_key]`` (an ``(n_obs, 2)`` array, the canonical
46474714
``obsm["spatial"]``), area and diameter to ``table.obs``. Values are stored in the element's
4648-
*intrinsic* coordinates/units (coordinate-system independent). Labels area is the pixel count;
4649-
shapes area is ``geometry.area`` (``pi*r**2`` for circles); equivalent diameter is
4650-
``2 * sqrt(area / pi)``. Persisting them once lets later renders (and downstream tools such as
4651-
squidpy) reuse them instead of recomputing.
4652-
4653-
Existing values for the measured rows are overwritten; pass ``centroids=False`` to keep a
4654-
pre-existing ``obsm[obsm_key]`` (e.g. from a reader). Per-cell measurements need a table to
4655-
write into, so an element without an annotating table cannot be measured (this raises). The
4656-
label path never densifies the raster — it streams it block by block, scaling to Xenium-size
4657-
masks.
4715+
*intrinsic* pixel coordinates/units (which align directly with the element's own raster/geometry).
4716+
Labels area is the pixel count; shapes area is ``geometry.area`` (``pi*r**2`` for circles);
4717+
equivalent diameter is ``2 * sqrt(area / pi)``. Persisting them once lets later renders (and
4718+
downstream tools such as squidpy) reuse them instead of recomputing.
4719+
4720+
Centroids already present for an element's rows (e.g. a reader-provided ``obsm[obsm_key]``) are
4721+
**not** overwritten — a warning is emitted and that element's centroid write is skipped (remove
4722+
the key or use a different ``obsm_key`` to recompute); ``area``/``diameter`` overwrite our own
4723+
columns. Instances annotated in the table but absent from the element are written as NaN with a
4724+
warning. Per-cell measurements need a table to write into, so an element without an annotating
4725+
table cannot be measured (this raises). The label path never densifies the raster — it streams
4726+
it block by block with memory O(n_labels), scaling to Xenium-size masks.
46584727
46594728
Parameters
46604729
----------

‎tests/pl/test_utils.py‎

Lines changed: 67 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from anndata import AnnData
1010
from shapely.geometry import Point
1111
from spatialdata import SpatialData, get_centroids
12-
from spatialdata.models import PointsModel, ShapesModel, TableModel
12+
from spatialdata.models import Labels2DModel, PointsModel, ShapesModel, TableModel
1313

1414
import spatialdata_plot
1515
from spatialdata_plot.pl import measure_obs
@@ -511,6 +511,19 @@ def _add_shapes_table(sdata: SpatialData, element: str = "blobs_polygons", name:
511511
return sdata
512512

513513

514+
def _labels_sdata(arr: np.ndarray, name: str = "lab", table: str = "t") -> SpatialData:
515+
"""Build a SpatialData with a single 2D-labels element annotated by a table."""
516+
ids = np.unique(arr)
517+
ids = ids[ids != 0].astype(int)
518+
adata = AnnData(np.zeros((len(ids), 1), dtype=np.float32))
519+
adata.obs["instance_id"] = ids
520+
adata.obs["region"] = name
521+
return SpatialData(
522+
labels={name: Labels2DModel.parse(arr, dims=("y", "x"))},
523+
tables={table: TableModel.parse(adata, region_key="region", instance_key="instance_id", region=name)},
524+
)
525+
526+
514527
class TestMeasureObs:
515528
"""`measure_obs` materializes centroids/area/equivalent diameter into the annotating table."""
516529

@@ -572,12 +585,15 @@ def test_inplace_false_leaves_original_untouched(self, sdata_blobs: SpatialData)
572585
assert "spatial" in out["table"].obsm
573586
assert "spatial" not in sdata_blobs["table"].obsm # original not mutated
574587

575-
def test_recompute_overwrites(self, sdata_blobs: SpatialData) -> None:
576-
measure_obs(sdata_blobs, "blobs_labels")
588+
def test_existing_centroids_not_clobbered(self, sdata_blobs: SpatialData) -> None:
589+
# #1: a populated obsm["spatial"] (reader- or prior-call-provided) is not overwritten.
577590
table = sdata_blobs["table"]
578-
table.obsm["spatial"][0] = [999.0, 999.0] # corrupt one row
579-
measure_obs(sdata_blobs, "blobs_labels") # second call overwrites with the real centroid
580-
assert tuple(table.obsm["spatial"][0]) != (999.0, 999.0)
591+
sentinel = np.arange(table.n_obs * 2, dtype=float).reshape(table.n_obs, 2)
592+
table.obsm["spatial"] = sentinel.copy()
593+
with pytest.warns(UserWarning, match="already populated"):
594+
measure_obs(sdata_blobs, "blobs_labels")
595+
np.testing.assert_array_equal(table.obsm["spatial"], sentinel) # centroids untouched
596+
assert "area" in table.obs # area/diameter still written
581597

582598
def test_centroids_false_keeps_existing_obsm(self, sdata_blobs: SpatialData) -> None:
583599
table = sdata_blobs["table"]
@@ -593,6 +609,51 @@ def test_incompatible_obsm_shape_raises(self, sdata_blobs: SpatialData) -> None:
593609
with pytest.raises(ValueError, match="Refusing to overwrite"):
594610
measure_obs(sdata_blobs, "blobs_labels")
595611

612+
def test_unmatched_instance_ids_warn_and_write_nan(self, sdata_blobs: SpatialData) -> None:
613+
# #2: instance ids that don't match the element (e.g. str vs int) -> warn + NaN, not silent.
614+
table = sdata_blobs["table"]
615+
table.obs["instance_id"] = table.obs["instance_id"].astype(str)
616+
with pytest.warns(UserWarning, match="no match"):
617+
measure_obs(sdata_blobs, "blobs_labels")
618+
assert np.isnan(table.obsm["spatial"]).all()
619+
620+
def test_float_dtype_labels_supported(self, sdata_blobs: SpatialData) -> None:
621+
# #3: a float-typed (but integer-valued) labels raster must not crash np.bincount.
622+
arr = np.asarray(sdata_blobs["blobs_labels"].data).astype(np.float32)
623+
sd = _labels_sdata(arr)
624+
measure_obs(sd, "lab", table_name="t")
625+
assert np.isfinite(sd["t"].obsm["spatial"]).all()
626+
627+
def test_existing_nonnumeric_column_raises_before_any_write(self, sdata_blobs: SpatialData) -> None:
628+
# #4: a non-numeric collision raises BEFORE obsm is mutated (no half-written table).
629+
table = sdata_blobs["table"]
630+
table.obs["area"] = pd.Categorical(["lo"] * table.n_obs)
631+
with pytest.raises(ValueError, match="not numeric"):
632+
measure_obs(sdata_blobs, "blobs_labels")
633+
assert "spatial" not in table.obsm # atomic: nothing written
634+
635+
def test_sparse_high_label_ids(self, sdata_blobs: SpatialData) -> None:
636+
# #5: sparse/high label ids (max id >> n_labels) are measured correctly (dense relabelling).
637+
arr = np.asarray(sdata_blobs["blobs_labels"].data)
638+
hi = (arr.astype(np.int64) * 1000) # ids become 1000, 2000, ... ; max id is huge, few labels
639+
measure_obs(sd_hi := _labels_sdata(hi), "lab", table_name="t")
640+
measure_obs(sd_lo := _labels_sdata(arr.astype(np.int64)), "lab", table_name="t")
641+
# relabelling values does not move pixels -> identical centroid set
642+
np.testing.assert_allclose(
643+
np.sort(sd_hi["t"].obsm["spatial"], axis=0), np.sort(sd_lo["t"].obsm["spatial"], axis=0)
644+
)
645+
646+
def test_polygon_with_radius_column_uses_geometric_area(self, sdata_blobs: SpatialData) -> None:
647+
# #7: dispatch on geometry type, not the "radius" column name -> polygons use geometry.area.
648+
gdf = sdata_blobs["blobs_polygons"].copy()
649+
gdf["radius"] = 5.0 # incidental column; must NOT trigger the circle (pi*r**2) branch
650+
sdata_blobs["pr"] = ShapesModel.parse(gdf)
651+
_add_shapes_table(sdata_blobs, "pr", name="pt")
652+
measure_obs(sdata_blobs, "pr", table_name="pt")
653+
order = sdata_blobs["pt"].obs["instance_id"].to_numpy()
654+
expected = gdf.geometry.area.to_numpy()[[list(gdf.index).index(i) for i in order]]
655+
np.testing.assert_allclose(sdata_blobs["pt"].obs["area"].to_numpy(), expected, rtol=1e-9)
656+
596657
def test_flags_select_outputs(self, sdata_blobs: SpatialData) -> None:
597658
measure_obs(sdata_blobs, "blobs_labels", area=False, diameter=False)
598659
table = sdata_blobs["table"]

0 commit comments

Comments
 (0)