|
26 | 26 | import spatialdata as sd |
27 | 27 | from anndata import AnnData |
28 | 28 | from cycler import Cycler, cycler |
| 29 | +from dask.array.core import slices_from_chunks |
29 | 30 | from datashader.core import Canvas |
30 | 31 | from geopandas import GeoDataFrame |
31 | 32 | from matplotlib import colors, patheffects, rcParams |
|
59 | 60 | from skimage.util import map_array |
60 | 61 | from spatialdata import ( |
61 | 62 | SpatialData, |
| 63 | + deepcopy as sd_deepcopy, |
62 | 64 | get_element_annotators, |
63 | 65 | get_extent, |
64 | 66 | get_values, |
|
67 | 69 | ) |
68 | 70 | from spatialdata._core.query.relational_query import _locate_value |
69 | 71 | 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 | +) |
71 | 80 | from spatialdata.transformations.operations import get_transformation |
72 | 81 | from spatialdata.transformations.transformations import Scale, Translation |
73 | 82 | from spatialdata.transformations.transformations import Sequence as TransformSequence |
@@ -4417,3 +4426,356 @@ def _convert_alpha_to_datashader_range(alpha: float) -> float: |
4417 | 4426 | """Convert alpha from the range [0, 1] to the range [0, 255] used in datashader.""" |
4418 | 4427 | # prevent a value of 255, bc that led to fully colored test plots instead of just colored points/shapes |
4419 | 4428 | 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