Respect value_key when aggregating image channels

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 = []
 

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: