(feat): sc.get.aggregate via dask

docs/release-notes/3700.feature.md

@@ -0,0 +1 @@
+Make {func}`scanpy.get.aggregate` `dask` compatible over all aggregations except median. {smaller}`I Gold`

src/scanpy/get/_aggregated.py

@@ -1,15 +1,16 @@
 from __future__ import annotations
 
-from functools import singledispatch
-from typing import TYPE_CHECKING, Literal
+from functools import partial, singledispatch
+from typing import TYPE_CHECKING, Literal, TypedDict
 
 import numpy as np
 import pandas as pd
 from anndata import AnnData, utils
+from fast_array_utils.stats._power import power as fau_power  # TODO: upstream
 from scipy import sparse
 from sklearn.utils.sparsefuncs import csc_median_axis_0
 
-from scanpy._compat import CSBase
+from scanpy._compat import CSBase, CSRBase, DaskArray
 
 from .._utils import _resolve_axis, get_literal_vals
 from .get import _check_mask
@@ -19,7 +20,7 @@
 
     from numpy.typing import NDArray
 
-    Array = np.ndarray | CSBase
+    Array = np.ndarray | CSBase | DaskArray
 
 # Used with get_literal_vals
 AggType = Literal["count_nonzero", "mean", "sum", "var", "median"]
@@ -330,13 +331,98 @@ def _aggregate(
     *,
     mask: NDArray[np.bool_] | None = None,
     dof: int = 1,
-):
+) -> dict[AggType, np.ndarray | DaskArray]:
     msg = f"Data type {type(data)} not supported for aggregation"
     raise NotImplementedError(msg)
 
 
+class MeanVarDict(TypedDict):
+    mean: DaskArray
+    var: DaskArray
+
+
+def aggregate_dask_mean_var(
+    data: DaskArray,
+    by: pd.Categorical,
+    *,
+    mask: NDArray[np.bool_] | None = None,
+    dof: int = 1,
+) -> MeanVarDict:
+    mean = aggregate_dask(data, by, "mean", mask=mask, dof=dof)["mean"]
+    sq_mean = aggregate_dask(fau_power(data, 2), by, "mean", mask=mask, dof=dof)["mean"]
+    # TODO: If we don't compute here, the results are not deterministic under the process cluster for sparse.
+    if isinstance(data._meta, CSRBase):
+        sq_mean = sq_mean.compute()
+    var = sq_mean - fau_power(mean, 2)
+    if dof != 0:
+        group_counts = np.bincount(by.codes)
+        var *= (group_counts / (group_counts - dof))[:, np.newaxis]
+    return {"var": var, "mean": mean}
+
+
+@_aggregate.register(DaskArray)
+def aggregate_dask(
+    data: DaskArray,
+    by: pd.Categorical,
+    func: AggType | Iterable[AggType],
+    *,
+    mask: NDArray[np.bool_] | None = None,
+    dof: int = 1,
+) -> dict[AggType, DaskArray]:
+    if not isinstance(data._meta, CSRBase | np.ndarray):
+        msg = f"Got {type(data._meta)} meta in DaskArray but only csr_matrix/csr_array and ndarray are supported."
+        raise ValueError(msg)
+    if data.chunksize[1] != data.shape[1]:
+        msg = "Feature axis must be unchunked"
+        raise ValueError(msg)
+
+    def aggregate_chunk_sum_or_count_nonzero(
+        chunk: Array, *, func: Literal["count_nonzero", "sum"], block_info=None
+    ):
+        subset = slice(*block_info[0]["array-location"][0])
+        by_subsetted = by[subset]
+        mask_subsetted = mask[subset] if mask is not None else mask
+        res = _aggregate(chunk, by_subsetted, func, mask=mask_subsetted, dof=dof)[func]
+        return res[None, :]
+
+    funcs = set([func] if isinstance(func, str) else func)
+    if "median" in funcs:
+        msg = "Dask median calculation not supported.  If you want a median-of-medians calculation, please open an issue."
+        raise NotImplementedError(msg)
+    has_mean, has_var = (v in funcs for v in ["mean", "var"])
+    funcs_no_var_or_mean = funcs - {"var", "mean"}
+    aggregated = {
+        f: data.map_blocks(
+            partial(aggregate_chunk_sum_or_count_nonzero, func=func),
+            new_axis=(1,),
+            chunks=((1,) * data.blocks.size, (len(by.categories),), (data.shape[1],)),
+            meta=np.array(
+                [],
+                dtype=np.float64
+                if func not in ["count_nonzero", "sum"]
+                else data.dtype,
+            ),
+        ).sum(axis=0)
+        for f in funcs_no_var_or_mean
+    }
+    if has_var:
+        aggredated_mean_var = aggregate_dask_mean_var(data, by, mask=mask, dof=dof)
+        aggregated["var"] = aggredated_mean_var["var"]
+        if has_mean:
+            aggregated["mean"] = aggredated_mean_var["mean"]
+    # division must come after, not before, the summation for numerical precision
+    # i.e., we can't just call map blocks over the mean function.
+    elif has_mean:
+        group_counts = np.bincount(by.codes)
+        aggregated["mean"] = (
+            aggregate_dask(data, by, "sum", mask=mask, dof=dof)["sum"]
+            / group_counts[:, None]
+        )
+    return aggregated
+
+
 @_aggregate.register(pd.DataFrame)
-def aggregate_df(data, by, func, *, mask=None, dof=1):
+def aggregate_df(data, by, func, *, mask=None, dof=1) -> dict[AggType, np.ndarray]:
     return _aggregate(data.values, by, func, mask=mask, dof=dof)
 
 

src/scanpy/preprocessing/_qc.py

@@ -416,6 +416,9 @@ def _(mtx: DaskArray, ns: Collection[int]) -> DaskArray:
     if not isinstance(mtx._meta, CSRBase | np.ndarray):
         msg = f"DaskArray must have csr matrix or ndarray meta, got {mtx._meta}."
         raise ValueError(msg)
+    if mtx.chunksize[1] != mtx.shape[1]:
+        msg = f"{mtx} must not be chunked along the feature axis"
+        raise ValueError(msg)
     return mtx.map_blocks(
         lambda x: top_segment_proportions(x, ns), meta=np.array([])
     ).compute()

src/testing/scanpy/_pytest/params.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from functools import wraps
 from importlib.metadata import version
 from typing import TYPE_CHECKING
 
@@ -10,18 +11,18 @@
 from packaging.version import Version
 from scipy import sparse
 
-from .._helpers import (
-    as_dense_dask_array,
-    as_sparse_dask_array,
-)
+from .._helpers import as_dense_dask_array, as_sparse_dask_array
 from .._pytest.marks import needs
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterable
     from typing import Any, Literal
 
+    import numpy as np
     from _pytest.mark.structures import ParameterSet
 
+    from ....scanpy._compat import DaskArray
+
 
 skipif_no_sparray = pytest.mark.skipif(
     Version(version("anndata")) < Version("0.11"),
@@ -41,6 +42,18 @@ def param_with(
     )
 
 
+def _chunked_1d(
+    f: Callable[[np.ndarray], DaskArray],
+) -> Callable[[np.ndarray], DaskArray]:
+    @wraps(f)
+    def wrapper(a: np.ndarray) -> DaskArray:
+        da = f(a)
+        return da.rechunk((da.chunksize[0], -1))
+
+    wrapper.__name__ = f"{wrapper.__name__}-1d_chunked"
+    return wrapper
+
+
 MAP_ARRAY_TYPES: dict[
     tuple[Literal["mem", "dask"], Literal["dense", "sparse"]],
     tuple[ParameterSet, ...],
@@ -51,20 +64,21 @@ def param_with(
         pytest.param(sparse.csc_matrix, id="scipy_csc_mat"),  # noqa: TID251
         pytest.param(sparse.csr_array, id="scipy_csr_arr", marks=[skipif_no_sparray]),  # noqa: TID251
     ),
-    ("dask", "dense"): (
+    ("dask", "dense"): tuple(
         pytest.param(
-            as_dense_dask_array,
+            wrapper(as_dense_dask_array),
             marks=[needs.dask, pytest.mark.anndata_dask_support],
-            id="dask_array_dense",
-        ),
+            id=f"dask_array_dense{suffix}",
+        )
+        for wrapper, suffix in [(lambda x: x, ""), (_chunked_1d, "-1d_chunked")]
     ),
-    ("dask", "sparse"): (
+    ("dask", "sparse"): tuple(
         pytest.param(
-            as_sparse_dask_array,
+            wrapper(as_sparse_dask_array),
             marks=[needs.dask, pytest.mark.anndata_dask_support],
-            id="dask_array_sparse",
-        ),
-        # probably not necessary to also do csc
+            id=f"dask_array_sparse{suffix}",
+        )
+        for wrapper, suffix in [(lambda x: x, ""), (_chunked_1d, "-1d_chunked")]
     ),
 }
 

tests/test_aggregated.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import anndata as ad
 import numpy as np
 import pandas as pd
@@ -8,11 +10,24 @@
 from scipy import sparse
 
 import scanpy as sc
+from scanpy._compat import DaskArray
 from scanpy._utils import _resolve_axis, get_literal_vals
 from scanpy.get._aggregated import AggType
 from testing.scanpy._helpers import assert_equal
 from testing.scanpy._helpers.data import pbmc3k_processed
-from testing.scanpy._pytest.params import ARRAY_TYPES_MEM
+from testing.scanpy._pytest.marks import needs
+from testing.scanpy._pytest.params import ARRAY_TYPES as ARRAY_TYPES_ALL
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from scanpy._compat import CSRBase
+
+ARRAY_TYPES = [
+    at
+    for at in ARRAY_TYPES_ALL
+    if at.id not in {"dask_array_dense", "dask_array_sparse"}
+]
 
 
 @pytest.fixture(params=get_literal_vals(AggType))
@@ -93,16 +108,18 @@ def test_mask(axis):
     assert np.all(by_name["0"].layers["sum"] == 0)
 
 
-@pytest.mark.parametrize("array_type", ARRAY_TYPES_MEM)
+@pytest.mark.parametrize("array_type", ARRAY_TYPES)
 def test_aggregate_vs_pandas(metric, array_type):
     adata = pbmc3k_processed().raw.to_adata()
     adata = adata[
         adata.obs["louvain"].isin(adata.obs["louvain"].cat.categories[:5]), :1_000
     ].copy()
     adata.X = array_type(adata.X)
+    xfail_dask_median(adata, metric)
     adata.obs["percent_mito_binned"] = pd.cut(adata.obs["percent_mito"], bins=5)
     result = sc.get.aggregate(adata, ["louvain", "percent_mito_binned"], metric)
-
+    if isinstance(adata.X, DaskArray):
+        adata.X = adata.X.compute()
     if metric == "count_nonzero":
         expected = (
             (adata.to_df() != 0)
@@ -124,7 +141,8 @@ def test_aggregate_vs_pandas(metric, array_type):
     )
     expected.index.name = None
     expected.columns.name = None
-
+    if isinstance(result.layers[metric], DaskArray):
+        result.layers[metric] = result.layers[metric].compute()
     result_df = result.to_df(layer=metric)
     result_df.index.name = None
     result_df.columns.name = None
@@ -139,16 +157,17 @@ def test_aggregate_vs_pandas(metric, array_type):
     pd.testing.assert_frame_equal(result_df, expected, check_dtype=False, atol=1e-5)
 
 
-@pytest.mark.parametrize("array_type", ARRAY_TYPES_MEM)
+@pytest.mark.parametrize("array_type", ARRAY_TYPES)
 def test_aggregate_axis(array_type, metric):
     adata = pbmc3k_processed().raw.to_adata()
     adata = adata[
         adata.obs["louvain"].isin(adata.obs["louvain"].cat.categories[:5]), :1_000
     ].copy()
     adata.X = array_type(adata.X)
+    xfail_dask_median(adata, metric)
     expected = sc.get.aggregate(adata, ["louvain"], metric)
-    actual = sc.get.aggregate(adata.T, ["louvain"], metric, axis=1).T
-
+    actual = sc.get.aggregate(adata.T, ["louvain"], metric, axis=1)
+    actual = actual.T
     assert_equal(expected, actual)
 
 
@@ -194,6 +213,44 @@ def test_aggregate_incorrect_dim():
         sc.get.aggregate(adata, ["louvain"], "sum", axis="foo")
 
 
+def to_bad_chunking(x: CSRBase):
+    import dask.array as da
+
+    return da.from_array(
+        x,
+        chunks=(x.shape[0] // 2, x.shape[1] // 2),
+        meta=sparse.csr_matrix(np.array([])),  # noqa: TID251
+    )
+
+
+def to_csc(x: CSRBase):
+    import dask.array as da
+
+    return da.from_array(
+        x.tocsc(),
+        chunks=(x.shape[0] // 2, x.shape[1]),
+        meta=sparse.csc_matrix(np.array([])),  # noqa: TID251
+    )
+
+
+@needs.dask
+@pytest.mark.anndata_dask_support
+@pytest.mark.parametrize(
+    ("func", "error_msg"),
+    [
+        pytest.param(to_csc, "only csr_matrix", id="csc"),
+        pytest.param(
+            to_bad_chunking, "Feature axis must be unchunked", id="bad_chunking"
+        ),
+    ],
+)
+def test_aggregate_bad_dask_array(func: Callable[[CSRBase], DaskArray], error_msg: str):
+    adata = pbmc3k_processed().raw.to_adata()
+    adata.X = func(adata.X)
+    with pytest.raises(ValueError, match=error_msg):
+        sc.get.aggregate(adata, ["louvain"], "sum")
+
+
 @pytest.mark.parametrize("axis_name", ["obs", "var"])
 def test_aggregate_axis_specification(axis_name):
     axis, axis_name = _resolve_axis(axis_name)
@@ -387,15 +444,24 @@ def test_combine_categories(label_cols, cols, expected):
     pd.testing.assert_frame_equal(reconstructed_df, result_label_df)
 
 
-@pytest.mark.parametrize("array_type", ARRAY_TYPES_MEM)
+def xfail_dask_median(adata, metric):
+    if isinstance(adata.X, DaskArray) and metric == "median":
+        pytest.xfail("Median calculation not implemented for Dask")
+
+
+@pytest.mark.parametrize("array_type", ARRAY_TYPES)
 def test_aggregate_arraytype(array_type, metric):
     adata = pbmc3k_processed().raw.to_adata()
     adata = adata[
         adata.obs["louvain"].isin(adata.obs["louvain"].cat.categories[:5]), :1_000
     ].copy()
     adata.X = array_type(adata.X)
+    xfail_dask_median(adata, metric)
     aggregate = sc.get.aggregate(adata, ["louvain"], metric)
-    assert isinstance(aggregate.layers[metric], np.ndarray)
+    assert isinstance(
+        aggregate.layers[metric],
+        DaskArray if isinstance(adata.X, DaskArray) else np.ndarray,
+    )
 
 
 def test_aggregate_obsm_varm():