perf: parallel downsample

.vscode/launch.json

@@ -19,7 +19,10 @@
             "purpose": ["debug-test"],
             "console": "internalConsole",
             "justMyCode": false,
-            "env": { "PYTEST_ADDOPTS": "--color=yes" },
+            "env": {
+                "PYTEST_ADDOPTS": "--color=yes",
+                "NUMBA_DEBUG_TYPEINFER": "1",
+            },
             "presentation": { "hidden": true },
         },
     ],

benchmarks/benchmarks/preprocessing_counts.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+from inspect import signature
 from itertools import product
 from typing import TYPE_CHECKING
 
@@ -15,10 +16,19 @@
 from ._utils import get_count_dataset
 
 if TYPE_CHECKING:
+    from typing import Any
+
     from ._utils import Dataset, KeyCount
 
 
-# ASV suite
+def cache_adata(dataset: Dataset, layer: KeyCount) -> None:
+    """Without this caching, asv was running several processes which meant the data was repeatedly downloaded."""
+    adata, batch_key = get_count_dataset(dataset, layer=layer)
+    assert "lop1p" not in adata.uns
+    adata.uns["batch_key"] = batch_key
+    adata.write_h5ad(f"{dataset}_{layer}.h5ad")
+
+
 class PreprocessingCountsSuite:  # noqa: D101
     params: tuple[list[Dataset], list[KeyCount]] = (
         ["pbmc68k_reduced", "pbmc3k"],
@@ -27,12 +37,8 @@ class PreprocessingCountsSuite:  # noqa: D101
     param_names = ("dataset", "layer")
 
     def setup_cache(self) -> None:
-        """Without this caching, asv was running several processes which meant the data was repeatedly downloaded."""
         for dataset, layer in product(*self.params):
-            adata, batch_key = get_count_dataset(dataset, layer=layer)
-            assert "lop1p" not in adata.uns
-            adata.uns["batch_key"] = batch_key
-            adata.write_h5ad(f"{dataset}_{layer}.h5ad")
+            cache_adata(dataset, layer)
 
     def setup(self, dataset, layer) -> None:
         self.adata = ad.read_h5ad(f"{dataset}_{layer}.h5ad")
@@ -65,6 +71,40 @@ def peakmem_scrublet(self, *_) -> None:
     #     sc.pp.highly_variable_genes(self.adata, flavor="seurat_v3_paper")
 
 
+class PreprocessingCountsRngSuite:  # noqa: D101
+    params: tuple[list[Dataset], list[str], list[str]] = (
+        ["pbmc68k_reduced", "pbmc3k"],
+        ["rng", "random_state"],
+    )
+    param_names = ("dataset", "layer")
+
+    def setup_cache(self) -> None:
+        for dataset in self.params[0]:
+            cache_adata(dataset, "counts")
+
+    def setup(self, dataset, rng_arg) -> None:
+        if (
+            rng_arg == "rng"
+            and "rng" not in signature(sc.pp.downsample_counts).parameters
+        ):
+            raise NotImplementedError
+        self.adata = ad.read_h5ad(f"{dataset}_counts.h5ad")
+        self.rng_kw: Any = {rng_arg: 0}
+        self.total = self.adata.X.sum() / 10
+
+    def time_downsample_per_cell(self, *_) -> None:
+        sc.pp.downsample_counts(self.adata, counts_per_cell=3, **self.rng_kw)
+
+    def peakmem_downsample_per_cell(self, *_) -> None:
+        sc.pp.downsample_counts(self.adata, counts_per_cell=3, **self.rng_kw)
+
+    def time_downsample_total(self, *_) -> None:
+        sc.pp.downsample_counts(self.adata, total_counts=self.total, **self.rng_kw)
+
+    def peakmem_downsample_total(self, *_) -> None:
+        sc.pp.downsample_counts(self.adata, total_counts=self.total, **self.rng_kw)
+
+
 class FastSuite:
     """Suite for fast preprocessing operations."""
 
@@ -75,11 +115,8 @@ class FastSuite:
     param_names = ("dataset", "layer")
 
     def setup_cache(self) -> None:
-        """Without this caching, asv was running several processes which meant the data was repeatedly downloaded."""
         for dataset, layer in product(*self.params):
-            adata, _ = get_count_dataset(dataset, layer=layer)
-            assert "lop1p" not in adata.uns
-            adata.write_h5ad(f"{dataset}_{layer}.h5ad")
+            cache_adata(dataset, layer)
 
     def setup(self, dataset, layer) -> None:
         self.adata = ad.read_h5ad(f"{dataset}_{layer}.h5ad")

docs/release-notes/3983.feat.md

@@ -0,0 +1 @@
+Add support for {class}`numpy.random.Generator` to all functions previously accepting a `random_state` parameter {smaller}`P Angerer`

hatch.toml

@@ -36,6 +36,7 @@ overrides.matrix.deps.python = [
     { if = [ "low-vers" ], value = "3.12" },
 ]
 overrides.matrix.deps.extra-dependencies = [
+    { if = [ "stable" ], value = "scipy>=1.17" },
     { if = [ "pre" ], value = "anndata @ git+https://github.com/scverse/anndata.git" },
     { if = [ "pre" ], value = "pandas>=3rc0" },
 ]

pyproject.toml

@@ -272,12 +272,15 @@ filterwarnings = [
     # we want to see and eventually fix these
     "default::numba.core.errors.NumbaPerformanceWarning",
     "default:.*TSNE.*random.*to.*pca:FutureWarning",                       # we should set init=obsm["X_pca"] or so
+    "default::numba.core.errors.NumbaTypeSafetyWarning",
     # matplotlib <3.11 uses old pyparsing APIs
     "ignore::pyparsing.warnings.PyparsingDeprecationWarning",
     # igraph vs leidenalg warning
     "ignore:The `igraph` implementation of leiden clustering:UserWarning",
     # everybody uses this zarr 3 feature, including us, XArray, lots of data out there …
     "ignore:Consolidated metadata is currently not part:UserWarning",
+    # joblib fallback to serial mode in restricted multiprocessing environments
+    "ignore:.*joblib will operate in serial mode:UserWarning",
 ]
 
 [tool.coverage]

src/scanpy/_docs.py

@@ -0,0 +1,17 @@
+"""Shared docstrings for general parameters."""
+
+from __future__ import annotations
+
+__all__ = ["doc_rng"]
+
+doc_rng = """\
+rng
+    Random number generation to control stochasticity.
+
+    If a type:`SeedLike` value, it’s used to seed a new random number generator;
+    If a :class:`numpy.random.Generator`, `rng`’s state will be directly advanced;
+    If :data:`None`, a non-reproducible random number generator is used.
+    See :func:`numpy.random.default_rng` for more details.
+
+    The default value matches legacy scanpy behavior and will change to `None` in scanpy 2.0.\
+"""

src/scanpy/_utils/numba.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from functools import cache
+
+import numba as nb
+import numpy as np
+from numba.extending import overload_method
+from numba.np.numpy_support import is_nonelike
+
+__all__ = ["add_np_generator_choice"]
+
+
+@cache  # make sure this is only called once
+def add_np_generator_choice():
+    """Add `np.random.Generator.choice` to numba if necessary."""
+
+    @nb.njit  # noqa: TID251
+    def test(rng: np.random.Generator) -> int:
+        return rng.choice(1)
+
+    try:
+        test(np.random.default_rng())
+    except nb.TypingError:
+        _register_np_generator_choice()
+
+
+def _register_np_generator_choice():
+
+    @overload_method(nb.types.NumPyRandomGeneratorType, "choice")
+    def NumPyRandomGeneratorType_choice(inst, a, size=None, replace=True):  # noqa: FBT002, N802
+        if is_nonelike(size):
+
+            def impl(inst, a, size=None, replace=True):  # noqa: FBT002
+                return inst.integers(0, a)
+
+        else:
+
+            def impl(inst, a, size=None, replace=True):  # noqa: FBT002
+                if replace:
+                    return inst.integers(0, a, size=size)
+                else:
+                    pool = np.arange(a)
+                    for i in range(size):
+                        j = inst.integers(i, a)
+                        pool[i], pool[j] = pool[j], pool[i]
+                    return pool[:size].copy()
+
+        return impl