chore: add csc-in-dask tests (#3870)

Author: ilan-gold

src/testing/scanpy/_pytest/params.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from functools import wraps
+from functools import partial, wraps
 from importlib.metadata import version
 from typing import TYPE_CHECKING
 
@@ -29,6 +29,24 @@
     reason="scipy cs{rc}_array not supported in anndata<0.11",
 )
 
+anndata_test_utils_supports_typ_kwarg = Version(version("anndata")) >= Version("0.12.6")
+
+
+def gen_csr_csc_params_wrapper(
+    func: Callable,
+    format: Literal["csr", "csc"],
+    matrix_or_array: Literal["matrix", "array"],
+):
+    def wrapper(arr):
+        if anndata_test_utils_supports_typ_kwarg:
+            return _chunked_1d(
+                partial(func, typ=getattr(sparse, f"{format}_{matrix_or_array}"))
+            )(arr)
+        return _chunked_1d(func)(arr)
+
+    wrapper.__name__ = f"{func.__name__}-1d_chunked-{format}_{matrix_or_array}"
+    return wrapper
+
 
 def param_with(
     at: ParameterSet,
@@ -48,7 +66,11 @@ def _chunked_1d(
     @wraps(f)
     def wrapper(a: np.ndarray) -> DaskArray:
         da = f(a)
-        return da.rechunk((da.chunksize[0], -1))
+        return da.rechunk(
+            (da.chunksize[0], -1)
+            if not hasattr(da._meta, "format") or da._meta.format == "csr"
+            else (-1, da.chunksize[1])
+        )
 
     wrapper.__name__ = f"{wrapper.__name__}-1d_chunked"
     return wrapper
@@ -75,10 +97,31 @@ def wrapper(a: np.ndarray) -> DaskArray:
     ("dask", "sparse"): tuple(
         pytest.param(
             wrapper(as_sparse_dask_matrix),
-            marks=[needs.dask],
+            marks=[needs.dask, skip_csc_mark]
+            if skip_csc_mark is not None
+            else [needs.dask],
             id=f"dask_array_sparse{suffix}",
         )
-        for wrapper, suffix in [(lambda x: x, ""), (_chunked_1d, "-1d_chunked")]
+        for wrapper, suffix, skip_csc_mark in [
+            (lambda x: x, "", None),
+            *(
+                (
+                    partial(
+                        gen_csr_csc_params_wrapper,
+                        format=format,
+                        matrix_or_array=matrix_or_array,
+                    ),
+                    f"-1d_chunked-{format}_{matrix_or_array}",
+                    pytest.mark.skipif(
+                        not anndata_test_utils_supports_typ_kwarg and format == "csc",
+                        reason="anndata < 0.12.6 lacked the required kwargs to enable csc matrix test utils.",
+                    ),
+                )
+                for format in ["csr", "csc"]
+                # TODO: use `array` as well once anndata 0.13 drops
+                for matrix_or_array in ["matrix"]
+            ),
+        ]
     ),
 }
 

tests/test_aggregated.py

@@ -22,10 +22,16 @@
 
     from scanpy._compat import CSRBase
 
-ARRAY_TYPES = [
+VALID_ARRAY_TYPES = [
     at
     for at in ARRAY_TYPES_ALL
-    if at.id not in {"dask_array_dense", "dask_array_sparse"}
+    if at.id
+    not in {
+        "dask_array_dense",
+        "dask_array_sparse",
+        "dask_array_sparse-1d_chunked-csc_array",
+        "dask_array_sparse-1d_chunked-csc_matrix",
+    }
 ]
 
 
@@ -118,7 +124,7 @@ def test_mask(axis):
     assert np.all(by_name["0"].layers["sum"] == 0)
 
 
-@pytest.mark.parametrize("array_type", ARRAY_TYPES)
+@pytest.mark.parametrize("array_type", VALID_ARRAY_TYPES)
 def test_aggregate_vs_pandas(
     metric: AggType, array_type, request: pytest.FixtureRequest
 ):
@@ -160,7 +166,7 @@ def test_aggregate_vs_pandas(
     pd.testing.assert_frame_equal(result_df, expected, check_dtype=False, atol=1e-5)
 
 
-@pytest.mark.parametrize("array_type", ARRAY_TYPES)
+@pytest.mark.parametrize("array_type", VALID_ARRAY_TYPES)
 def test_aggregate_axis(array_type, metric, request: pytest.FixtureRequest):
     adata = pbmc3k_processed().raw.to_adata()
     adata = adata[
@@ -456,7 +462,7 @@ 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)
+@pytest.mark.parametrize("array_type", VALID_ARRAY_TYPES)
 def test_aggregate_arraytype(
     array_type, metric: AggType, request: pytest.FixtureRequest
 ):

tests/test_highly_variable_genes.py

@@ -649,7 +649,12 @@ def test_seurat_v3_bad_chunking(adata, array_type, flavor):
     ],
 )
 @pytest.mark.parametrize(
-    "array_type", [p for p in ARRAY_TYPES if "dask" not in p.id or "1d_chunked" in p.id]
+    "array_type",
+    [
+        p
+        for p in ARRAY_TYPES
+        if "dask" not in p.id or ("1d_chunked" in p.id and "csr" in p.id)
+    ],
 )
 @pytest.mark.parametrize("batch_key", [None, "batch"])
 def test_subset_inplace_consistency(flavor, array_type, batch_key):
@@ -728,7 +733,9 @@ def test_subset_inplace_consistency(flavor, array_type, batch_key):
     ],
 )
 @pytest.mark.parametrize("batch_key", [None, "batch"], ids=["single", "batched"])
-@pytest.mark.parametrize("to_dask", [p for p in ARRAY_TYPES if "1d_chunked" in p.id])
+@pytest.mark.parametrize(
+    "to_dask", [p for p in ARRAY_TYPES if "1d_chunked" in p.id and "csr" in p.id]
+)
 def test_dask_consistency(adata: AnnData, flavor, batch_key, to_dask):
     # current blob produces singularities in loess....maybe a bad sign of the data?
     if "seurat_v3" in flavor:

tests/test_pca.py

@@ -58,18 +58,24 @@
     [-1.50180389, 5.56886849, 1.64034442, 2.24476032, -0.05109001],
 ])
 
-
-ARRAY_TYPES = [
+# These are array types which are expected to work with the current PCA implementation.
+VALID_ARRAY_TYPES = [
     param_with(
         at,
         marks=[needs.dask_ml] if at.id == "dask_array_dense-1d_chunked" else [],
     )
     for at in ARRAY_TYPES_ALL
-    if at.id not in {"dask_array_dense", "dask_array_sparse"}
+    if at.id
+    not in {
+        "dask_array_dense",
+        "dask_array_sparse",
+        "dask_array_sparse-1d_chunked-csc_array",
+        "dask_array_sparse-1d_chunked-csc_matrix",
+    }
 ]
 
 
-@pytest.fixture(params=ARRAY_TYPES)
+@pytest.fixture(params=VALID_ARRAY_TYPES)
 def array_type(request: pytest.FixtureRequest) -> ArrayType:
     return request.param
 
@@ -93,10 +99,14 @@ def gen_pca_params(
         xfail_reason = "dask without 1d chunking scheme not supported"
         yield None, None, xfail_reason
         return
-    if id == "dask_array_sparse-1d_chunked" and not zero_center:
+    if "dask_array_sparse-1d_chunked" in id and not zero_center:
         xfail_reason = "Sparse-in-dask with zero_center=False not implemented yet"
         yield None, None, xfail_reason
         return
+    if "dask_array_sparse-1d_chunked-csc" in id:
+        xfail_reason = "Sparse-in-dask with csc blocks not implemented yet"
+        yield None, None, xfail_reason
+        return
     if svd_solver_type is None:
         yield None, None, None
         return
@@ -137,7 +147,12 @@ def possible_solvers(
             svd_solvers = {"auto", "full", "tsqr", "randomized", "covariance_eigh"}
         case (dc, False) if id == "dask_array_dense-1d_chunked":
             svd_solvers = {"tsqr", "randomized"}
-        case (dc, True) if id == "dask_array_sparse-1d_chunked":
+        case (dc, True) if (
+            # See https://github.com/scverse/scanpy/blob/216b21d91312b899e939db9636d9ab20e7c29d77/src/testing/scanpy/_pytest/params.py#L88-L103
+            # for why we need two checks (i.e., before and after allowing CSC matrices)
+            "dask_array_sparse-1d_chunked-csr" in id
+            or id == "dask_array_sparse-1d_chunked"
+        ):
             svd_solvers = {"covariance_eigh"}
         case (type() as dc, True) if issubclass(dc, CSBase):
             svd_solvers = {"arpack"} | SKLEARN_ADDITIONAL
@@ -148,7 +163,7 @@ def possible_solvers(
         case (helpers.asarray, False):
             svd_solvers = {"arpack", "randomized"}
         case _:
-            pytest.fail(f"Unknown {array_type=} ({zero_center=})")
+            pytest.fail(f"Unknown {array_type=} ({zero_center=}) ({id=})")
 
     if svd_solver_type == "invalid":
         svd_solvers = all_svd_solvers - svd_solvers
@@ -178,7 +193,7 @@ def possible_solvers(
                 f"{svd_solver or svd_solver_type}-{'xfail' if xfail_reason else warn_pat_expected}"
             ),
         )
-        for array_type in ARRAY_TYPES
+        for array_type in VALID_ARRAY_TYPES
         for zero_center in [True, False]
         for svd_solver_type in [None, "valid", "invalid"]
         for svd_solver, warn_pat_expected, xfail_reason in gen_pca_params(
@@ -542,10 +557,13 @@ def test_pca_rep(rep: Literal["layer", "obsm"]) -> None:
 @pytest.mark.parametrize(
     "other_array_type",
     [
-        lambda x: x.toarray(),
-        *(at.values[0] for at in ARRAY_TYPES if "1d_chunked" in at.id),
+        pytest.param(lambda x: x.toarray(), id="dense"),
+        *(
+            pytest.param(at.values[0], id=at.id)
+            for at in VALID_ARRAY_TYPES
+            if "1d_chunked" in at.id
+        ),
     ],
-    ids=["dense-mem", "sparse-dask", "dense-dask"],
 )
 def test_covariance_eigh_impls(other_array_type):
     warnings.filterwarnings("error")
@@ -590,8 +608,8 @@ def test_sparse_dask_input_errors(msg_re: str, op: Callable[[DaskArray], DaskArr
     adata_sparse.X = op(
         next(
             at.values[0]
-            for at in ARRAY_TYPES
-            if at.id == "dask_array_sparse-1d_chunked"
+            for at in VALID_ARRAY_TYPES
+            if "dask_array_sparse-1d_chunked" in at.id
         )(adata_sparse.X)
     )
 
@@ -612,7 +630,9 @@ def test_sparse_dask_input_errors(msg_re: str, op: Callable[[DaskArray], DaskArr
 def test_cov_sparse_dask(dtype, dtype_arg, rtol):
     x_arr = A_list.astype(dtype)
     x = next(
-        at.values[0] for at in ARRAY_TYPES if at.id == "dask_array_sparse-1d_chunked"
+        at.values[0]
+        for at in VALID_ARRAY_TYPES
+        if "dask_array_sparse-1d_chunked" in at.id
     )(x_arr)
     cov, gram, mean = _cov_sparse_dask(x, return_gram=True, dtype=dtype_arg)
     np.testing.assert_allclose(mean, np.mean(x_arr, axis=0))

tests/test_preprocessing.py

@@ -280,7 +280,12 @@ def test_sample_copy_backed_error(tmp_path):
 
 @pytest.mark.parametrize("array_type", ARRAY_TYPES)
 @pytest.mark.parametrize("max_value", [None, 1.0], ids=["no_clip", "clip"])
-def test_scale_matrix_types(array_type, zero_center, max_value):
+def test_scale_matrix_types(
+    *,
+    array_type: Callable,
+    zero_center: bool,
+    max_value: float | None,
+):
     adata = pbmc68k_reduced()
     adata.X = adata.raw.X
     adata_casted = adata.copy()

tests/test_qc_metrics.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from contextlib import nullcontext
+
 import numpy as np
 import pandas as pd
 import pytest
@@ -9,7 +11,7 @@
 from scipy import sparse
 
 import scanpy as sc
-from scanpy._compat import DaskArray
+from scanpy._compat import CSCBase, DaskArray
 from scanpy.preprocessing._qc import (
     describe_obs,
     describe_var,
@@ -83,8 +85,16 @@ def test_top_segments(request: pytest.FixtureRequest, array_type):
         reason = "DaskArray with feature axis chunking not yet supported"
         request.applymarker(pytest.mark.xfail(reason=reason))
     a = array_type(np.ones((300, 100)))
-    with maybe_dask_process_context():
+    is_csc_dask = isinstance(a, DaskArray) and isinstance(a._meta, CSCBase)
+    with (
+        maybe_dask_process_context(),
+        pytest.raises(ValueError, match=r"DaskArray must have csr")
+        if is_csc_dask
+        else nullcontext(),
+    ):
         seg = top_segment_proportions(a, [50, 100])
+    if is_csc_dask:
+        return
     assert (seg[:, 0] == 0.5).all()
     assert (seg[:, 1] == 1.0).all()
 
@@ -107,10 +117,22 @@ def test_top_proportions(request: pytest.FixtureRequest, array_type):
 # While many of these are trivial,
 # they’re also just making sure the metrics are there
 def test_qc_metrics(adata_prepared: AnnData):
-    with maybe_dask_process_context():
+    is_csc_dask = isinstance(adata_prepared.X, DaskArray) and isinstance(
+        adata_prepared.X._meta, CSCBase
+    )
+    with (
+        maybe_dask_process_context(),
+        (
+            pytest.raises(ValueError, match=r"DaskArray must have csr")
+            if is_csc_dask
+            else nullcontext()
+        ),
+    ):
         sc.pp.calculate_qc_metrics(
             adata_prepared, qc_vars=["mito", "negative"], inplace=True
         )
+    if is_csc_dask:
+        return
     x = (
         adata_prepared.X.compute()
         if isinstance(adata_prepared.X, DaskArray)
@@ -159,14 +181,26 @@ def test_qc_metrics(adata_prepared: AnnData):
 
 
 def test_qc_metrics_idempotent(adata_prepared: AnnData):
-    with maybe_dask_process_context():
+    is_csc_dask = isinstance(adata_prepared.X, DaskArray) and isinstance(
+        adata_prepared.X._meta, CSCBase
+    )
+    with (
+        maybe_dask_process_context(),
+        (
+            pytest.raises(ValueError, match=r"DaskArray must have csr")
+            if is_csc_dask
+            else nullcontext()
+        ),
+    ):
         sc.pp.calculate_qc_metrics(
             adata_prepared, qc_vars=["mito", "negative"], inplace=True
         )
         old_obs, old_var = adata_prepared.obs.copy(), adata_prepared.var.copy()
         sc.pp.calculate_qc_metrics(
             adata_prepared, qc_vars=["mito", "negative"], inplace=True
         )
+    if is_csc_dask:
+        return
     assert set(adata_prepared.obs.columns) == set(old_obs.columns)
     assert set(adata_prepared.var.columns) == set(old_var.columns)
     for col in adata_prepared.obs:
@@ -176,7 +210,15 @@ def test_qc_metrics_idempotent(adata_prepared: AnnData):
 
 
 def test_qc_metrics_no_log1p(adata_prepared: AnnData):
-    with maybe_dask_process_context():
+    with (
+        maybe_dask_process_context(),
+        (
+            pytest.raises(ValueError, match=r"DaskArray must have csr")
+            if isinstance(adata_prepared.X, DaskArray)
+            and isinstance(adata_prepared.X._meta, CSCBase)
+            else nullcontext()
+        ),
+    ):
         sc.pp.calculate_qc_metrics(
             adata_prepared, qc_vars=["mito", "negative"], log1p=False, inplace=True
         )