From 02340aacbb102a140d064d7a9e8e57a84058d373 Mon Sep 17 00:00:00 2001 From: Qinzi27 Date: Fri, 18 Sep 2026 04:46:52 +1000 Subject: [PATCH] Respect value_key when aggregating image channels --- src/spatialdata/_core/operations/aggregate.py | 32 +++++-- tests/core/operations/test_aggregations.py | 86 +++++++++++++++++++ 2 files changed, 112 insertions(+), 6 deletions(-) diff --git a/src/spatialdata/_core/operations/aggregate.py b/src/spatialdata/_core/operations/aggregate.py index b4772df1b..412ac668b 100644 --- a/src/spatialdata/_core/operations/aggregate.py +++ b/src/spatialdata/_core/operations/aggregate.py @@ -59,18 +59,21 @@ def aggregate( The regions to aggregate by: if `by_sdata` is None, must be a SpatialElement, otherwise must be a string specifying the name of the SpatialElement in `by_sdata` value_key - Name (or list of names) of the columns containing the values to aggregate; can refer both to numerical or - categorical values. If the values are categorical, `value_key` can't be a list. + Name (or list of names) of the columns or image channels containing the values to aggregate; can refer both + to numerical or categorical values. If the values are categorical, `value_key` can't be a list. The key can be: - the name of a column(s) in the dataframe (Dask `DataFrame` for points or `GeoDataFrame` for shapes); - the name of obs column(s) in the associated `AnnData` table (for points, shapes and labels); - the name of a var(s), referring to the column(s) of the X matrix in the table (for points, shapes and - labels). + labels); + - the name of a channel(s) in the `c` coordinate of an image. - If nothing is passed here, it defaults to the equivalent of a column of ones. - Defaults to `FEATURE_KEY` for points (if present). + For images, `None` selects all channels; otherwise, channels are selected in the requested order. Empty lists, + duplicate requests, unknown channels, and non-unique image channel names raise `ValueError` when selecting + channels. For points and shapes, it defaults to the equivalent of a column of ones, or to `FEATURE_KEY` for + points (if present). agg_func Aggregation function to apply over point values, e.g. `"mean"`, `"sum"`, `"count"`. Passed to :func:`pandas.DataFrame.groupby.agg` or to :func:`xrspatial.zonal_stats` @@ -200,7 +203,7 @@ def aggregate( raise NotImplementedError("fractions = True is not yet supported for raster aggregation") assert isinstance(values_, DataArray | DataTree) assert isinstance(by_, DataArray | DataTree) - adata = _aggregate_image_by_labels(values=values_, by=by_, agg_func=agg_func, **kwargs) + adata = _aggregate_image_by_labels(values=values_, by=by_, agg_func=agg_func, value_key=value_key, **kwargs) if adata is None: raise NotImplementedError(f"Cannot aggregate {values_type} by {by_type}") @@ -260,6 +263,7 @@ def _aggregate_image_by_labels( values: DataArray | DataTree, by: DataArray | DataTree, agg_func: str | list[str] = "mean", + value_key: str | list[str] | None = None, **kwargs: Any, ) -> ad.AnnData: """ @@ -274,6 +278,8 @@ def _aggregate_image_by_labels( agg_func Aggregation function to apply over point values, e.g. "mean", "sum", "count" from :func:`xrspatial.zonal_stats`. + value_key + Image channel name(s) to aggregate, in the requested order. If `None`, aggregate all channels. kwargs Additional keyword arguments to pass to :func:`xrspatial.zonal_stats`. @@ -299,6 +305,20 @@ def _aggregate_image_by_labels( assert isinstance(values_variable, DataArray) values = values_variable + if value_key is not None: + channels = [value_key] if isinstance(value_key, str) else value_key + if not channels: + raise ValueError("`value_key` must not be empty when selecting image channels.") + if len(channels) != len(set(channels)): + raise ValueError("`value_key` must not contain duplicate image channels.") + available_channels = values.get_index("c") + if not available_channels.is_unique: + raise ValueError("Image channel names must be unique when selecting with `value_key`.") + missing = [channel for channel in channels if channel not in available_channels] + if missing: + raise ValueError(f"Image channels {missing} specified by `value_key` were not found.") + values = values.sel(c=channels) + agg_func = [agg_func] if isinstance(agg_func, str) else agg_func outs = [] diff --git a/tests/core/operations/test_aggregations.py b/tests/core/operations/test_aggregations.py index d471b79cc..9f92f37d6 100644 --- a/tests/core/operations/test_aggregations.py +++ b/tests/core/operations/test_aggregations.py @@ -8,6 +8,7 @@ from anndata.tests.helpers import assert_equal from geopandas import GeoDataFrame from numpy.random import default_rng +from xarray.testing import assert_identical from spatialdata import aggregate, to_polygons from spatialdata._core._deepcopy import deepcopy as _deepcopy @@ -359,6 +360,91 @@ def test_aggregate_image_by_labels(labels_blobs, image_schema, labels_schema) -> assert len(out) == 3 +@pytest.fixture(params=[None, [2]], ids=["single_scale", "multiscale"]) +def image_labels_for_channel_selection(request): + # Background is deliberately bright; regions have different pixel counts. + image = Image2DModel.parse( + np.array( + [ + [[999, 2, 4, 999], [1, 3, 5, 7]], + [[999, 10, 14, 999], [2, 4, 6, 8]], + [[999, 6, 8, 999], [3, 5, 7, 9]], + ], + dtype=float, + ), + c_coords=["DAPI", "CD3", "CD20"], + scale_factors=request.param, + ) + labels = Labels2DModel.parse(np.array([[0, 1, 1, 0], [2, 2, 2, 2]], dtype=np.int32)) + return SpatialData(images={"image": image}, labels={"labels": labels}) + + +@pytest.mark.parametrize("value_key", [None, "CD3", ["CD3"], ["CD20", "CD3"]]) +@pytest.mark.parametrize("agg_func", ["mean", "sum", ["mean", "sum", "count"]]) +def test_aggregate_image_by_labels_value_key(image_labels_for_channel_selection, value_key, agg_func): + sdata = image_labels_for_channel_selection + image, labels = sdata.images["image"], sdata.labels["labels"] + original_image, original_labels = image.copy(deep=True), labels.copy(deep=True) + out = aggregate(values=image, by=labels, value_key=value_key, agg_func=agg_func).tables["table"] + + channels = ( + ["DAPI", "CD3", "CD20"] if value_key is None else [value_key] if isinstance(value_key, str) else value_key + ) + stats = [agg_func] if isinstance(agg_func, str) else agg_func + # Independent, hand-calculated expectations; rows are label IDs 1 and 2. + expected = { + "DAPI": {"mean": [3, 4], "sum": [6, 16], "count": [2, 4]}, + "CD3": {"mean": [12, 5], "sum": [24, 20], "count": [2, 4]}, + "CD20": {"mean": [7, 6], "sum": [14, 24], "count": [2, 4]}, + } + names = [f"channel_{channel}_{stat}" for channel in channels for stat in stats] + assert out.var_names.tolist() == names + np.testing.assert_allclose( + out.X.toarray(), np.column_stack([expected[c][stat] for c in channels for stat in stats]) + ) + assert out.obs_names.tolist() == ["1", "2"] + assert out.obs["instance_id"].tolist() == [1, 2] + assert out.obs["region"].tolist() == ["by", "by"] + assert out.uns[TableModel.ATTRS_KEY] == {"region": "by", "region_key": "region", "instance_key": "instance_id"} + + all_channels = aggregate(values=image, by=labels, agg_func=agg_func).tables["table"] + assert_equal(out, all_channels[:, names].copy()) + assert_identical(image, original_image) + assert_identical(labels, original_labels) + + +@pytest.mark.parametrize( + ("value_key", "match"), + [ + ([], "must not be empty"), + (["CD3", "CD3"], "must not contain duplicate"), + ("missing", "missing.*not found"), + (["CD3", "missing"], "missing.*not found"), + ], +) +def test_aggregate_image_by_labels_invalid_value_key(image_labels_for_channel_selection, value_key, match): + sdata = image_labels_for_channel_selection + with pytest.raises(ValueError, match=match): + aggregate(values=sdata.images["image"], by=sdata.labels["labels"], value_key=value_key) + + +def test_aggregate_image_by_labels_value_key_spatialdata(image_labels_for_channel_selection): + sdata = image_labels_for_channel_selection + out = sdata.aggregate(values="image", by="labels", value_key="CD3", zone_ids=[2]).tables["table"] + assert out.var_names.tolist() == ["channel_CD3_sum"] + np.testing.assert_allclose(out.X.toarray(), [[20]]) + assert out.obs["instance_id"].tolist() == [2] + assert out.obs["region"].tolist() == ["labels"] + assert out.uns[TableModel.ATTRS_KEY]["region"] == "labels" + + +def test_aggregate_image_by_labels_value_key_nonunique_channels(): + image = Image2DModel.parse(np.ones((2, 2, 2)), c_coords=["CD3", "CD3"]) + labels = Labels2DModel.parse(np.array([[0, 1], [2, 2]], dtype=np.int32)) + with pytest.raises(ValueError, match="Image channel names must be unique"): + aggregate(values=image, by=labels, value_key="CD3") + + @pytest.mark.parametrize("values", ["blobs_image", "blobs_points", "blobs_circles", "blobs_polygons"]) @pytest.mark.parametrize("by", ["blobs_labels", "blobs_circles", "blobs_polygons"]) def test_aggregate_requiring_alignment(sdata_blobs: SpatialData, values, by) -> None: