feat: seurat v3 with dask csr (#3340)

Co-authored-by: Philipp A. <flying-sheep@web.de>

Author: ilan-gold

docs/release-notes/3340.feature.md

@@ -0,0 +1 @@
+{func}`scanpy.pp.highly_variable_genes` flavors `seurat_v3` and `seurat_v3_paper` are now `dask`-compatible {smaller}`I Gold`

src/scanpy/_utils/__init__.py

@@ -768,6 +768,15 @@ def dematrix[SA: _SupportedArray](x: SA | np.matrix) -> SA:
     return x
 
 
+def raise_if_dask_feature_axis_chunked(x: Any):
+    if isinstance(x, DaskArray) and x.chunksize[1] != x.shape[1]:
+        msg = (
+            "Only dask arrays with chunking along the first axis are supported. "
+            f"Got chunksize {x.chunksize} with shape {x.shape}. "
+        )
+        raise ValueError(msg)
+
+
 def select_groups(
     adata: AnnData,
     groups_order_subset: Iterable[str] | Literal["all"] = "all",

src/scanpy/experimental/pp/_normalization.py

@@ -8,12 +8,7 @@
 
 from ... import logging as logg
 from ..._compat import CSBase, warn
-from ..._utils import (
-    _doc_params,
-    _empty,
-    check_nonnegative_integers,
-    view_to_actual,
-)
+from ..._utils import _doc_params, _empty, check_nonnegative_integers, view_to_actual
 from ...experimental._docs import (
     doc_adata,
     doc_check_values,

src/scanpy/preprocessing/_highly_variable_genes.py

@@ -2,6 +2,7 @@
 
 import warnings
 from dataclasses import dataclass
+from functools import singledispatch
 from inspect import signature
 from typing import TYPE_CHECKING, TypedDict, cast
 
@@ -12,9 +13,13 @@
 from fast_array_utils import stats
 
 from .. import logging as logg
-from .._compat import CSBase, DaskArray, old_positionals, warn
+from .._compat import CSBase, CSRBase, DaskArray, old_positionals, warn
 from .._settings import Verbosity, settings
-from .._utils import check_nonnegative_integers, sanitize_anndata
+from .._utils import (
+    check_nonnegative_integers,
+    raise_if_dask_feature_axis_chunked,
+    sanitize_anndata,
+)
 from ..get import _get_obs_rep
 from ._distributed import materialize_as_ndarray
 from ._simple import filter_genes
@@ -28,6 +33,91 @@
     from .._types import HVGFlavor
 
 
+@singledispatch
+def clip_square_sum(
+    data_batch: np.ndarray, clip_val: np.ndarray
+) -> tuple[np.ndarray, np.ndarray]:
+    """Clip data_batch by clip_val.
+
+    Parameters
+    ----------
+    data_batch
+        The data to be clipped
+    clip_val
+        Clip by these values (must be broadcastable to the input data)
+
+    Returns
+    -------
+        The clipeed data
+    """
+    batch_counts = data_batch.astype(np.float64).copy()
+    clip_val_broad = np.broadcast_to(clip_val, batch_counts.shape)
+    np.putmask(
+        batch_counts,
+        batch_counts > clip_val_broad,
+        clip_val_broad,
+    )
+
+    squared_batch_counts_sum = np.square(batch_counts).sum(axis=0)
+    batch_counts_sum = batch_counts.sum(axis=0)
+    return squared_batch_counts_sum, batch_counts_sum
+
+
+@clip_square_sum.register(DaskArray)
+def _(data_batch: DaskArray, clip_val: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    n_blocks = data_batch.blocks.size
+
+    def sum_and_sum_squares_clipped_from_block(block):
+        return np.vstack(clip_square_sum(block, clip_val))[None, ...]
+
+    squared_batch_counts_sum, batch_counts_sum = (
+        data_batch.map_blocks(
+            sum_and_sum_squares_clipped_from_block,
+            new_axis=(1,),
+            chunks=((1,) * n_blocks, (2,), (data_batch.shape[1],)),
+            meta=np.array([]),
+            dtype=np.float64,
+        )
+        .sum(axis=0)
+        .compute()
+    )
+    return squared_batch_counts_sum, batch_counts_sum
+
+
+@clip_square_sum.register(CSBase)
+def _(data_batch: CSBase, clip_val: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    batch_counts = data_batch if isinstance(data_batch, CSRBase) else data_batch.tocsr()
+
+    return _sum_and_sum_squares_clipped(
+        batch_counts.indices,
+        batch_counts.data,
+        n_cols=batch_counts.shape[1],
+        clip_val=clip_val,
+        nnz=batch_counts.nnz,
+    )
+
+
+# parallel=False needed for accuracy
+@numba.njit(cache=True, parallel=False)  # noqa: TID251
+def _sum_and_sum_squares_clipped(
+    indices: NDArray[np.integer],
+    data: NDArray[np.floating],
+    *,
+    n_cols: int,
+    clip_val: NDArray[np.float64],
+    nnz: int,
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+    squared_batch_counts_sum = np.zeros(n_cols, dtype=np.float64)
+    batch_counts_sum = np.zeros(n_cols, dtype=np.float64)
+    for i in numba.prange(nnz):
+        idx = indices[i]
+        element = min(np.float64(data[i]), clip_val[idx])
+        squared_batch_counts_sum[idx] += element**2
+        batch_counts_sum[idx] += element
+
+    return squared_batch_counts_sum, batch_counts_sum
+
+
 def _highly_variable_genes_seurat_v3(  # noqa: PLR0912, PLR0915
     adata: AnnData,
     *,
@@ -70,23 +160,28 @@ def _highly_variable_genes_seurat_v3(  # noqa: PLR0912, PLR0915
         raise ImportError(msg) from e
     df = pd.DataFrame(index=adata.var_names)
     data = _get_obs_rep(adata, layer=layer)
+    raise_if_dask_feature_axis_chunked(data)
 
     if check_values and not check_nonnegative_integers(data):
         msg = f"`{flavor=!r}` expects raw count data, but non-integers were found."
         warn(msg, UserWarning)
 
     df["means"], df["variances"] = stats.mean_var(data, axis=0, correction=1)
 
-    if batch_key is None:
-        batch_info = pd.Categorical(np.zeros(adata.shape[0], dtype=int))
-    else:
-        batch_info = adata.obs[batch_key].to_numpy()
+    batch_info = (
+        pd.Categorical(np.zeros(adata.shape[0], dtype=int))
+        if batch_key is None
+        else adata.obs[batch_key].to_numpy()
+    )
 
     norm_gene_vars = []
     for b in np.unique(batch_info):
         data_batch = data[batch_info == b]
 
         mean, var = stats.mean_var(data_batch, axis=0, correction=1)
+        # These get computed anyway for loess
+        if isinstance(mean, DaskArray):
+            mean, var = mean.compute(), var.compute()
         not_const = var > 0
         estimat_var = np.zeros(data.shape[1], dtype=np.float64)
 
@@ -99,28 +194,10 @@ def _highly_variable_genes_seurat_v3(  # noqa: PLR0912, PLR0915
 
         # clip large values as in Seurat
         n_obs = data_batch.shape[0]
-        vmax = np.sqrt(n_obs)
-        clip_val = reg_std * vmax + mean
-        if isinstance(data_batch, CSBase):
-            batch_counts = data_batch.tocsr()
-            squared_batch_counts_sum, batch_counts_sum = _sum_and_sum_squares_clipped(
-                batch_counts.indices,
-                batch_counts.data,
-                n_cols=batch_counts.shape[1],
-                clip_val=clip_val,
-                nnz=batch_counts.nnz,
-            )
-        else:
-            batch_counts = data_batch.astype(np.float64).copy()
-            clip_val_broad = np.broadcast_to(clip_val, batch_counts.shape)
-            np.putmask(
-                batch_counts,
-                batch_counts > clip_val_broad,
-                clip_val_broad,
-            )
-
-            squared_batch_counts_sum = np.square(batch_counts).sum(axis=0)
-            batch_counts_sum = batch_counts.sum(axis=0)
+        clip_val = reg_std * np.sqrt(n_obs) + mean
+        squared_batch_counts_sum, batch_counts_sum = clip_square_sum(
+            data_batch, clip_val
+        )
 
         norm_gene_var = (1 / ((n_obs - 1) * np.square(reg_std))) * (
             (n_obs * np.square(mean))
@@ -142,10 +219,12 @@ def _highly_variable_genes_seurat_v3(  # noqa: PLR0912, PLR0915
     ma_ranked = np.ma.masked_invalid(ranked_norm_gene_vars)
     median_ranked = np.ma.median(ma_ranked, axis=0).filled(np.nan)
 
-    df["gene_name"] = df.index
-    df["highly_variable_nbatches"] = num_batches_high_var
-    df["highly_variable_rank"] = median_ranked
-    df["variances_norm"] = np.mean(norm_gene_vars, axis=0)
+    df = df.assign(
+        gene_name=df.index,
+        highly_variable_nbatches=num_batches_high_var,
+        highly_variable_rank=median_ranked,
+        variances_norm=np.mean(norm_gene_vars, axis=0),
+    )
     if flavor == "seurat_v3":
         sort_cols = ["highly_variable_rank", "highly_variable_nbatches"]
         sort_ascending = [True, False]
@@ -173,10 +252,13 @@ def _highly_variable_genes_seurat_v3(  # noqa: PLR0912, PLR0915
             "    'variances', float vector (adata.var)\n"
             "    'variances_norm', float vector (adata.var)"
         )
-        adata.var["highly_variable"] = df["highly_variable"].to_numpy()
-        adata.var["highly_variable_rank"] = df["highly_variable_rank"].to_numpy()
-        adata.var["means"] = df["means"].to_numpy()
-        adata.var["variances"] = df["variances"].to_numpy()
+        for to_numpy_key in [
+            "highly_variable",
+            "highly_variable_rank",
+            "means",
+            "variances",
+        ]:
+            adata.var[to_numpy_key] = df[to_numpy_key].to_numpy()
         adata.var["variances_norm"] = (
             df["variances_norm"].to_numpy().astype("float64", copy=False)
         )
@@ -193,27 +275,7 @@ def _highly_variable_genes_seurat_v3(  # noqa: PLR0912, PLR0915
             df = df.iloc[df["highly_variable"].to_numpy(), :]
 
         return df
-
-
-# parallel=False needed for accuracy
-@numba.njit(cache=True, parallel=False)  # noqa: TID251
-def _sum_and_sum_squares_clipped(
-    indices: NDArray[np.integer],
-    data: NDArray[np.floating],
-    *,
-    n_cols: int,
-    clip_val: NDArray[np.float64],
-    nnz: int,
-) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
-    squared_batch_counts_sum = np.zeros(n_cols, dtype=np.float64)
-    batch_counts_sum = np.zeros(n_cols, dtype=np.float64)
-    for i in numba.prange(nnz):
-        idx = indices[i]
-        element = min(np.float64(data[i]), clip_val[idx])
-        squared_batch_counts_sum[idx] += element**2
-        batch_counts_sum[idx] += element
-
-    return squared_batch_counts_sum, batch_counts_sum
+    return None
 
 
 @dataclass

src/scanpy/preprocessing/_pca/_dask.py

@@ -7,6 +7,7 @@
 import scipy.linalg
 from fast_array_utils import stats
 
+from scanpy._utils import raise_if_dask_feature_axis_chunked
 from scanpy._utils._doctests import doctest_needs
 
 from ..._compat import CSBase
@@ -52,13 +53,7 @@ def fit(self, x: DaskArray) -> PCAEighDaskFit:
                 f"Got {x._meta.format} as meta."
             )
             raise ValueError(msg)
-        if x.chunksize[1] != x.shape[1]:
-            msg = (
-                "Only dask arrays with chunking along the first axis are supported. "
-                f"Got chunksize {x.chunksize} with shape {x.shape}. "
-                "Rechunking should be simple and cost nothing from AnnData's on-disk format when the on-disk layout has this chunking."
-            )
-            raise ValueError(msg)
+        raise_if_dask_feature_axis_chunked(x)
         self.__class__ = PCAEighDaskFit
         self = cast("PCAEighDaskFit", self)  # noqa: PLW0642
 

src/testing/scanpy/_helpers/__init__.py

@@ -129,9 +129,9 @@ def as_dense_dask_array(*args, **kwargs) -> DaskArray:
     from anndata.tests.helpers import as_dense_dask_array
 
     a = as_dense_dask_array(*args, **kwargs)
+    # Newer versions of as_dense_dask_array chunk all axes by halve when the input is not a dask array.
     if (
         pkg_version("anndata") < Version("0.11")
-        and a.chunksize == a.shape
         and not isinstance(args[0], DaskArray)  # keep chunksize intact
     ):
         from anndata.tests.helpers import _half_chunk_size