(feat): pre-processing functions for dask with sparse chunks

scanpy/_compat.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from functools import partial
+from functools import partial, singledispatch
 from pathlib import Path
 
+import dask.array as da
+import numpy as np
 from legacy_api_wrap import legacy_api
 from packaging import version
+from scipy import sparse as sp
 
 try:
     from functools import cache
@@ -81,3 +84,29 @@ def pkg_version(package):
 
 
 old_positionals = partial(legacy_api, category=FutureWarning)
+
+
+@singledispatch
+def sum(X: np.ndarray | sp.spmatrix, axis=None):
+    return np.sum(X, axis=axis)
+
+
+@sum.register
+def _(X: da.Array, axis=None):
+    def sum_drop_keepdims(*args, **kwargs):
+        kwargs.pop("computing_meta", None)
+        if isinstance(X._meta, (sp.spmatrix, np.matrix)) or isinstance(
+            args[0], (sp.spmatrix, np.matrix)
+        ):  # forcing the `_meta` to be a sparse array really isn't desirable?
+            kwargs.pop("keepdims", None)
+            if isinstance(kwargs["axis"], tuple):
+                kwargs["axis"] = kwargs["axis"][0]
+        return da.chunk.sum(*args, **kwargs)
+
+    dtype = getattr(np.zeros(1, dtype=X.dtype).sum(), "dtype", object)
+
+    # operates on `np.matrix` for some reason with sparse chunks in dask so need explicit casting
+    def aggregate_sum(*args, **kwargs):
+        return da.chunk.sum(np.array(args[0]), **kwargs)
+
+    return da.reduction(X, sum_drop_keepdims, aggregate_sum, axis=axis, dtype=dtype)

scanpy/preprocessing/_distributed.py

@@ -35,15 +35,12 @@ def materialize_as_ndarray(
 
 
 def materialize_as_ndarray(
-    a: ArrayLike | tuple[ArrayLike | ZappyArray | DaskArray, ...],
+    a: DaskArray | ArrayLike | tuple[ArrayLike | ZappyArray | DaskArray, ...],
 ) -> tuple[np.ndarray] | np.ndarray:
     """Compute distributed arrays and convert them to numpy ndarrays."""
+    if isinstance(a, DaskArray):
+        return a.compute()
     if not isinstance(a, tuple):
         return np.asarray(a)
 
-    if not any(isinstance(arr, DaskArray) for arr in a):
-        return tuple(np.asarray(arr) for arr in a)
-
-    import dask.array as da
-
-    return da.compute(*a, sync=True)
+    return tuple(materialize_as_ndarray(arr) for arr in a)

scanpy/preprocessing/_normalization.py

@@ -8,7 +8,7 @@
 from sklearn.utils import sparsefuncs
 
 from .. import logging as logg
-from .._compat import DaskArray, old_positionals
+from .._compat import DaskArray, old_positionals, sum
 from .._utils import view_to_actual
 from ..get import _get_obs_rep, _set_obs_rep
 
@@ -35,7 +35,7 @@ def _normalize_data(X, counts, after=None, copy: bool = False):
     elif isinstance(counts, np.ndarray):
         np.divide(X, counts[:, None], out=X)
     else:
-        X = np.divide(X, counts[:, None])  # dask does not support kwarg "out"
+        X = X / counts[:, None]
     return X
 
 
@@ -187,22 +187,23 @@ def normalize_total(
 
     gene_subset = None
     msg = "normalizing counts per cell"
+
+    counts_per_cell = sum(X, axis=1)
     if exclude_highly_expressed:
-        counts_per_cell = X.sum(1)  # original counts per cell
+        counts_per_cell = sum(X, axis=1)  # original counts per cell
         counts_per_cell = np.ravel(counts_per_cell)
 
         # at least one cell as more than max_fraction of counts per cell
 
-        gene_subset = (X > counts_per_cell[:, None] * max_fraction).sum(0)
+        gene_subset = sum((X > counts_per_cell[:, None] * max_fraction), axis=0)
         gene_subset = np.asarray(np.ravel(gene_subset) == 0)
 
         msg += (
             ". The following highly-expressed genes are not considered during "
             f"normalization factor computation:\n{adata.var_names[~gene_subset].tolist()}"
         )
-        counts_per_cell = X[:, gene_subset].sum(1)
-    else:
-        counts_per_cell = X.sum(1)
+        counts_per_cell = sum(X[:, gene_subset], axis=1)
+
     start = logg.info(msg)
     counts_per_cell = np.ravel(counts_per_cell)
 
@@ -253,3 +254,4 @@ def normalize_total(
         return adata
     elif not inplace:
         return dat
+    return None

scanpy/preprocessing/_simple.py

@@ -17,7 +17,7 @@
 from sklearn.utils import check_array, sparsefuncs
 
 from .. import logging as logg
-from .._compat import old_positionals
+from .._compat import old_positionals, sum
 from .._settings import settings as sett
 from .._utils import (
     AnyRandom,
@@ -162,7 +162,7 @@ def filter_cells(
     X = data  # proceed with processing the data matrix
     min_number = min_counts if min_genes is None else min_genes
     max_number = max_counts if max_genes is None else max_genes
-    number_per_cell = np.sum(
+    number_per_cell = sum(
         X if min_genes is None and max_genes is None else X > 0, axis=1
     )
     if issparse(X):
@@ -172,7 +172,7 @@ def filter_cells(
     if max_number is not None:
         cell_subset = number_per_cell <= max_number
 
-    s = materialize_as_ndarray(np.sum(~cell_subset))
+    s = materialize_as_ndarray(sum(~cell_subset))
     if s > 0:
         msg = f"filtered out {s} cells that have "
         if min_genes is not None or min_counts is not None:
@@ -278,7 +278,7 @@ def filter_genes(
     X = data  # proceed with processing the data matrix
     min_number = min_counts if min_cells is None else min_cells
     max_number = max_counts if max_cells is None else max_cells
-    number_per_gene = np.sum(
+    number_per_gene = sum(
         X if min_cells is None and max_cells is None else X > 0, axis=0
     )
     if issparse(X):
@@ -288,7 +288,7 @@ def filter_genes(
     if max_number is not None:
         gene_subset = number_per_gene <= max_number
 
-    s = np.sum(~gene_subset)
+    s = sum(~gene_subset)
     if s > 0:
         msg = f"filtered out {s} genes that are detected "
         if min_cells is not None or min_counts is not None:

scanpy/testing/_helpers/data.py

@@ -18,11 +18,15 @@
 
 from typing import TYPE_CHECKING
 
+import dask.array as da
+from dask import delayed
+from scipy import sparse
+
 import scanpy as sc
 
 if TYPE_CHECKING:
     from anndata import AnnData
-
+    from anndata._core.sparse_dataset import SparseDataset
 # Functions returning the same objects (easy to misuse)
 
 
@@ -76,3 +80,49 @@
 
 def pbmc3k_normalized() -> AnnData:
     return _pbmc3k_normalized().copy()
+
+
+class CSRCallable:
+    """Dummy class to bypass dask checks"""
+
+    def __new__(cls, shape, dtype):
+        return csr_callable(shape, dtype)
+
+
+def csr_callable(shape: tuple[int, int], dtype) -> sparse.csr_matrix:
+    if len(shape) == 0:
+        shape = (0, 0)
+    if len(shape) == 1:
+        shape = (shape[0], 0)
+    elif len(shape) == 2:
+        pass
+    else:
+        raise ValueError(shape)
+
+    return sparse.csr_matrix(shape, dtype=dtype)
+
+
+def make_dask_chunk(x: SparseDataset, start: int, end: int) -> da.Array:
+    def take_slice(x, idx):
+        return x[idx]
+
+    return da.from_delayed(
+        delayed(take_slice)(x, slice(start, end)),
+        dtype=x.dtype,
+        shape=(end - start, x.shape[1]),
+        meta=CSRCallable,
+    )
+
+
+def sparse_dataset_as_dask(x: SparseDataset, stride: int):
+    n_chunks, rem = divmod(x.shape[0], stride)
+
+    chunks = []
+    cur_pos = 0
+    for i in range(n_chunks):
+        chunks.append(make_dask_chunk(x, cur_pos, cur_pos + stride))
+        cur_pos += stride
+    if rem:
+        chunks.append(make_dask_chunk(x, cur_pos, x.shape[0]))
+
+    return da.concatenate(chunks, axis=0)

scanpy/testing/_pytest/params.py

@@ -66,15 +66,7 @@ def param_with(
     at for (_, spsty), ats in MAP_ARRAY_TYPES.items() if spsty == "dense" for at in ats
 )
 
-ARRAY_TYPES_SUPPORTED = tuple(
-    (
-        param_with(at, marks=[pytest.mark.xfail(reason="sparse-in-dask not supported")])
-        if attrs == ("dask", "sparse")
-        else at
-    )
-    for attrs, ats in MAP_ARRAY_TYPES.items()
-    for at in ats
-)
+ARRAY_TYPES_SUPPORTED = tuple(at for ats in MAP_ARRAY_TYPES.values() for at in ats)
 """
 Sparse matrices in dask arrays aren’t officially supported upstream,
 so add xfail to them.

scanpy/tests/_data/10x-10k-subset.zarr/.zattrs

@@ -0,0 +1,4 @@
+{
+    "encoding-type": "anndata",
+    "encoding-version": "0.1.0"
+}

scanpy/tests/_data/10x-10k-subset.zarr/X/.zarray

@@ -1,7 +1,7 @@
 {
     "chunks": [
-        2000,
-        1000
+        1250,
+        125
     ],
     "compressor": {
         "blocksize": 0,

scanpy/tests/_data/10x-10k-subset.zarr/X/.zattrs

@@ -0,0 +1,4 @@
+{
+    "encoding-type": "array",
+    "encoding-version": "0.2.0"
+}

scanpy/tests/_data/10x-10k-subset.zarr/layers/.zattrs

@@ -0,0 +1,4 @@
+{
+    "encoding-type": "dict",
+    "encoding-version": "0.1.0"
+}

scanpy/tests/_data/10x-10k-subset.zarr/layers/.zgroup

@@ -0,0 +1,3 @@
+{
+    "zarr_format": 2
+}

scanpy/tests/_data/10x-10k-subset.zarr/layers/CSR_X/.zattrs

@@ -0,0 +1,8 @@
+{
+    "encoding-type": "csr_matrix",
+    "encoding-version": "0.1.0",
+    "shape": [
+        10000,
+        1000
+    ]
+}

scanpy/tests/_data/10x-10k-subset.zarr/layers/CSR_X/.zgroup

@@ -0,0 +1,3 @@
+{
+    "zarr_format": 2
+}

scanpy/tests/_data/10x-10k-subset.zarr/layers/CSR_X/data/.zarray

@@ -0,0 +1,20 @@
+{
+    "chunks": [
+        43245
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<i8",
+    "fill_value": 0,
+    "filters": null,
+    "order": "C",
+    "shape": [
+        691914
+    ],
+    "zarr_format": 2
+}

scanpy/tests/_data/10x-10k-subset.zarr/layers/CSR_X/indices/.zarray

@@ -0,0 +1,20 @@
+{
+    "chunks": [
+        86490
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<i4",
+    "fill_value": 0,
+    "filters": null,
+    "order": "C",
+    "shape": [
+        691914
+    ],
+    "zarr_format": 2
+}

scanpy/tests/_data/10x-10k-subset.zarr/layers/CSR_X/indptr/.zarray

@@ -0,0 +1,20 @@
+{
+    "chunks": [
+        10001
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<i4",
+    "fill_value": 0,
+    "filters": null,
+    "order": "C",
+    "shape": [
+        10001
+    ],
+    "zarr_format": 2
+}

scanpy/tests/_data/10x-10k-subset.zarr/obs/.zattrs

@@ -0,0 +1,6 @@
+{
+    "_index": "index",
+    "column-order": [],
+    "encoding-type": "dataframe",
+    "encoding-version": "0.2.0"
+}

scanpy/tests/_data/10x-10k-subset.zarr/obs/.zgroup

@@ -0,0 +1,3 @@
+{
+    "zarr_format": 2
+}

scanpy/tests/_data/10x-10k-subset.zarr/obs/.zarray → scanpy/tests/_data/10x-10k-subset.zarr/obs/index/.zarray

@@ -9,14 +9,13 @@
         "id": "blosc",
         "shuffle": 1
     },
-    "dtype": [
-        [
-            "index",
-            "|S18"
-        ]
+    "dtype": "|O",
+    "fill_value": 0,
+    "filters": [
+        {
+            "id": "vlen-utf8"
+        }
     ],
-    "fill_value": "AAAAAAAAAAAAAAAAAAAAAAAA",
-    "filters": null,
     "order": "C",
     "shape": [
         10000