feat: support np.random.Generator

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

@@ -278,6 +278,8 @@ filterwarnings = [
     "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/_utils/random.py

@@ -7,22 +7,27 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
-from sklearn.utils import check_random_state
+from sklearn.utils.random import check_random_state
 
 from . import ensure_igraph
 
 if TYPE_CHECKING:
-    from collections.abc import Generator
+    from collections.abc import Callable, Generator
+    from typing import Self
 
+    from numpy.random import BitGenerator
     from numpy.typing import NDArray
 
 
 __all__ = [
     "RNGLike",
     "SeedLike",
     "_LegacyRandom",
+    "_accepts_legacy_random_state",
+    "_if_legacy_apply_global",
+    "_legacy_random_state",
+    "_set_igraph_rng",
     "ith_k_tuple",
-    "legacy_numpy_gen",
     "random_k_tuples",
     "random_str",
 ]
@@ -38,34 +43,38 @@
 
 
 class _RNGIgraph:
-    """Random number generator for igraph so global seed is not changed.
+    """Random number generator for igraph so global random state is not changed.
 
     See :func:`igraph.set_random_number_generator` for the requirements.
     """
 
-    def __init__(self, random_state: int | np.random.RandomState = 0) -> None:
-        self._rng = check_random_state(random_state)
+    def __init__(self, rng: SeedLike | RNGLike | None) -> None:
+        self._rng = np.random.default_rng(rng)
 
     def getrandbits(self, k: int) -> int:
-        return self._rng.tomaxint() & ((1 << k) - 1)
+        if isinstance(self._rng, _FakeRandomGen):
+            i = self._rng._state.tomaxint()
+        else:
+            lims = np.iinfo(np.uint64)
+            i = int(self._rng.integers(0, lims.max, dtype=np.uint64, endpoint=True))
+        return i & ((1 << k) - 1)
 
-    def randint(self, a: int, b: int) -> int:
-        return self._rng.randint(a, b + 1)
+    def randint(self, a: int, b: int) -> np.int64:
+        """Can’t use `endpoint` here as _FakeRandomGen doesn’t support it."""
+        return self._rng.integers(a, b + 1)
 
     def __getattr__(self, attr: str):
         return getattr(self._rng, "normal" if attr == "gauss" else attr)
 
 
 @contextmanager
-def set_igraph_random_state(
-    random_state: int | np.random.RandomState,
-) -> Generator[None, None, None]:
+def _set_igraph_rng(rng: SeedLike | RNGLike | None) -> Generator[None]:
     ensure_igraph()
     import igraph
 
-    rng = _RNGIgraph(random_state)
+    ig_rng = _RNGIgraph(rng)
     try:
-        igraph.set_random_number_generator(rng)
+        igraph.set_random_number_generator(ig_rng)
         yield None
     finally:
         igraph.set_random_number_generator(random)
@@ -76,28 +85,55 @@ def set_igraph_random_state(
 ###################################
 
 
-def legacy_numpy_gen(
-    random_state: _LegacyRandom | None = None,
-) -> np.random.Generator:
-    """Return a random generator that behaves like the legacy one."""
-    if random_state is not None:
-        if isinstance(random_state, np.random.RandomState):
-            np.random.set_state(random_state.get_state(legacy=False))
-            return _FakeRandomGen(random_state)
-        np.random.seed(random_state)
-    return _FakeRandomGen(np.random.RandomState(np.random.get_bit_generator()))
+class _FakeRandomGen(np.random.Generator):
+    """A `Generator` that wraps a legacy `RandomState` instance.
 
+    To behave like a `RandomState`, it’s not enough to just use a MT19937 `bit_generator`
+    (as in `Generator(RandomState(seed).bit_generator)`),
+    so instead this hack uses the exact same random numbers as `RandomState(seed)`.
+    """
 
-class _FakeRandomGen(np.random.Generator):
+    _arg: _LegacyRandom
     _state: np.random.RandomState
 
-    def __init__(self, random_state: np.random.RandomState) -> None:
-        self._state = random_state
+    def __init__(
+        self, arg: _LegacyRandom, state: np.random.RandomState | None = None
+    ) -> None:
+        self._arg = arg
+        self._state = check_random_state(arg) if state is None else state
+
+    @property
+    def bit_generator(self) -> BitGenerator:
+        msg = "A _FakeRandomGen instance has no `bit_generator` attribute."
+        raise AttributeError(msg)
+
+    @classmethod
+    def wrap_global(
+        cls,
+        arg: _LegacyRandom = None,
+        state: np.random.RandomState | None = None,
+    ) -> Self:
+        """Create a generator that wraps the global `RandomState` backing the legacy `np.random` functions."""
+        if arg is not None:
+            if isinstance(arg, np.random.RandomState):
+                np.random.set_state(arg.get_state(legacy=False))
+                return _FakeRandomGen(arg, state)
+            np.random.seed(arg)
+        return _FakeRandomGen(arg, np.random.RandomState(np.random.get_bit_generator()))
+
+    def spawn(self, n_children: int) -> list[Self]:
+        """Return `self` `n_children` times.
+
+        In a real generator, the spawned children are independent,
+        but for backwards compatibility we return the same instance.
+        """
+        return [self] * n_children
 
     @classmethod
     def _delegate(cls) -> None:
+        names = dict(integers="randint")
         for name, meth in np.random.Generator.__dict__.items():
-            if name.startswith("_") or not callable(meth):
+            if name.startswith("_") or not callable(meth) or name in cls.__dict__:
                 continue
 
             def mk_wrapper(name: str, meth):
@@ -108,12 +144,66 @@ def wrapper(self: _FakeRandomGen, *args, **kwargs):
 
                 return wrapper
 
-            setattr(cls, name, mk_wrapper(name, meth))
+            setattr(cls, names.get(name, name), mk_wrapper(name, meth))
 
 
 _FakeRandomGen._delegate()
 
 
+def _if_legacy_apply_global(rng: np.random.Generator, /) -> np.random.Generator:
+    """Wrap the global legacy RNG if `rng` is a `_FakeRandomGen`.
+
+    This is used where our code used to  call `np.random.seed()`.
+    It’s a no-op if `rng` is not a `_FakeRandomGen`.
+    """
+    if not isinstance(rng, _FakeRandomGen):
+        return rng
+
+    return _FakeRandomGen.wrap_global(rng._arg, rng._state)
+
+
+def _legacy_random_state(
+    rng: SeedLike | RNGLike | None, /, *, always_state: bool = False
+) -> _LegacyRandom:
+    """Convert a np.random.Generator into a legacy `random_state` argument.
+
+    If `rng` is already a `_FakeRandomGen`, return its original `_arg` attribute.
+    """
+    if isinstance(rng, _FakeRandomGen):
+        return rng._state if always_state else rng._arg
+    [bitgen] = np.random.default_rng(rng).bit_generator.spawn(1)
+    return np.random.RandomState(bitgen)
+
+
+def _accepts_legacy_random_state[**P, R](
+    random_state_default: _LegacyRandom, /
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    """Make a function accept `random_state: _LegacyRandom` and pass it as `rng`.
+
+    If the decorated function is called with a `random_state` argument,
+    it’ll be wrapped in a :class:`_FakeRandomGen`.
+    Passing both ``rng`` and ``random_state`` at the same time is an error.
+    If neither is given, ``random_state_default`` is used.
+    """
+
+    def decorator(func: Callable[P, R]) -> Callable[P, R]:
+        @wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            match "random_state" in kwargs, "rng" in kwargs:
+                case True, True:
+                    msg = "Specify at most one of `rng` and `random_state`."
+                    raise TypeError(msg)
+                case True, False:
+                    kwargs["rng"] = _FakeRandomGen(kwargs.pop("random_state"))
+                case False, False:
+                    kwargs["rng"] = _FakeRandomGen(random_state_default)
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 ###################
 # Random k-tuples #
 ###################

src/scanpy/datasets/_datasets.py

@@ -12,13 +12,14 @@
 from .._compat import deprecated
 from .._settings import settings
 from .._utils._doctests import doctest_internet, doctest_needs
+from .._utils.random import _accepts_legacy_random_state, _legacy_random_state
 from ..readwrite import read, read_h5ad, read_visium
 from ._utils import check_datasetdir_exists
 
 if TYPE_CHECKING:
     from typing import Literal
 
-    from .._utils.random import _LegacyRandom
+    from .._utils.random import RNGLike, SeedLike
 
     type VisiumSampleID = Literal[
         "V1_Breast_Cancer_Block_A_Section_1",
@@ -54,13 +55,14 @@
 HERE = Path(__file__).parent
 
 
+@_accepts_legacy_random_state(0)
 def blobs(
     *,
     n_variables: int = 11,
     n_centers: int = 5,
     cluster_std: float = 1.0,
     n_observations: int = 640,
-    random_state: _LegacyRandom = 0,
+    rng: SeedLike | RNGLike | None = None,
 ) -> AnnData:
     """Gaussian Blobs.
 
@@ -75,7 +77,7 @@ def blobs(
     n_observations
         Number of observations. By default, this is the same observation number
         as in :func:`scanpy.datasets.krumsiek11`.
-    random_state
+    rng
         Determines random number generation for dataset creation.
 
     Returns
@@ -98,7 +100,7 @@ def blobs(
         n_features=n_variables,
         centers=n_centers,
         cluster_std=cluster_std,
-        random_state=random_state,
+        random_state=_legacy_random_state(rng),
     )
     return AnnData(x, obs=dict(blobs=y.astype(str)))
 

src/scanpy/experimental/_docs.py

@@ -60,8 +60,8 @@
 doc_pca_chunk = """\
 n_comps
     Number of principal components to compute in the PCA step.
-random_state
-    Random seed for setting the initial states for the optimization in the PCA step.
+rng
+    Random number generator for setting the initial states for the optimization in the PCA step.
 kwargs_pca
     Dictionary of further keyword arguments passed on to `scanpy.pp.pca()`.
 """

src/scanpy/experimental/pp/_normalization.py

@@ -9,6 +9,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.random import _accepts_legacy_random_state
 from ...experimental._docs import (
     doc_adata,
     doc_check_values,
@@ -27,6 +28,7 @@
     from typing import Any
 
     from ..._utils import Empty
+    from ..._utils.random import RNGLike, SeedLike
 
 
 def _pearson_residuals(
@@ -160,13 +162,14 @@ def normalize_pearson_residuals(
     check_values=doc_check_values,
     inplace=doc_inplace,
 )
+@_accepts_legacy_random_state(0)
 def normalize_pearson_residuals_pca(
     adata: AnnData,
     *,
     theta: float = 100,
     clip: float | None = None,
     n_comps: int | None = 50,
-    random_state: float = 0,
+    rng: SeedLike | RNGLike | None = None,
     kwargs_pca: Mapping[str, Any] = MappingProxyType({}),
     mask_var: np.ndarray | str | None | Empty = _empty,
     use_highly_variable: bool | None = None,
@@ -233,7 +236,7 @@ def normalize_pearson_residuals_pca(
     normalize_pearson_residuals(
         adata_pca, theta=theta, clip=clip, check_values=check_values
     )
-    pca(adata_pca, n_comps=n_comps, random_state=random_state, **kwargs_pca)
+    pca(adata_pca, n_comps=n_comps, rng=rng, **kwargs_pca)
     n_comps = adata_pca.obsm["X_pca"].shape[1]  # might be None
 
     if inplace:

src/scanpy/experimental/pp/_recipes.py

@@ -17,13 +17,17 @@
 )
 from scanpy.preprocessing import pca
 
+from ..._utils.random import _accepts_legacy_random_state
+
 if TYPE_CHECKING:
     from collections.abc import Mapping
     from typing import Any
 
     import pandas as pd
     from anndata import AnnData
 
+    from ..._utils.random import RNGLike, SeedLike
+
 
 @_doc_params(
     adata=doc_adata,
@@ -33,6 +37,7 @@
     check_values=doc_check_values,
     inplace=doc_inplace,
 )
+@_accepts_legacy_random_state(0)
 def recipe_pearson_residuals(  # noqa: PLR0913
     adata: AnnData,
     *,
@@ -42,7 +47,7 @@ def recipe_pearson_residuals(  # noqa: PLR0913
     batch_key: str | None = None,
     chunksize: int = 1000,
     n_comps: int | None = 50,
-    random_state: float | None = 0,
+    rng: SeedLike | RNGLike | None = None,
     kwargs_pca: Mapping[str, Any] = MappingProxyType({}),
     check_values: bool = True,
     inplace: bool = True,
@@ -133,7 +138,7 @@ def recipe_pearson_residuals(  # noqa: PLR0913
     experimental.pp.normalize_pearson_residuals(
         adata_pca, theta=theta, clip=clip, check_values=check_values
     )
-    pca(adata_pca, n_comps=n_comps, random_state=random_state, **kwargs_pca)
+    pca(adata_pca, n_comps=n_comps, rng=rng, **kwargs_pca)
 
     if inplace:
         normalization_param = adata_pca.uns["pearson_residuals_normalization"]