use sub-generators

Author: flying-sheep

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" },
 ]

src/scanpy/_utils/random.py

@@ -108,11 +108,19 @@ def wrap_global(
             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):
@@ -129,11 +137,11 @@ def wrapper(self: _FakeRandomGen, *args, **kwargs):
 _FakeRandomGen._delegate()
 
 
-def _if_legacy_apply_global(rng: np.random.Generator) -> np.random.Generator:
-    """Re-apply legacy `random_state` semantics when `rng` is a `_FakeRandomGen`.
+def _if_legacy_apply_global(rng: np.random.Generator, /) -> np.random.Generator:
+    """Wrap the global legacy RNG if `rng` is a `_FakeRandomGen`.
 
-    This resets the global legacy RNG from the original `_arg` and returns a
-    generator which continues drawing from the same internal state.
+    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
@@ -142,20 +150,20 @@ def _if_legacy_apply_global(rng: np.random.Generator) -> np.random.Generator:
 
 
 def _legacy_random_state(
-    rng: SeedLike | RNGLike | None, *, always_state: bool = False
+    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
-    rng = np.random.default_rng(rng)
-    return np.random.RandomState(rng.bit_generator.spawn(1)[0])
+    [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,
+    random_state_default: _LegacyRandom, /
 ) -> Callable[[Callable[P, R]], Callable[P, R]]:
     """Make a function accept `random_state: _LegacyRandom` and pass it as `rng`.
 

src/scanpy/neighbors/__init__.py

@@ -10,11 +10,12 @@
 
 import numpy as np
 import scipy
+from packaging.version import Version
 from scipy import sparse
 
 from .. import _utils
 from .. import logging as logg
-from .._compat import CSBase, CSRBase, SpBase, warn
+from .._compat import CSBase, CSRBase, SpBase, pkg_version, warn
 from .._settings import settings
 from .._utils import NeighborsView, _doc_params, get_literal_vals
 from .._utils.random import (
@@ -46,9 +47,8 @@
     # TODO: make `type` when https://github.com/sphinx-doc/sphinx/pull/13508 is released
     RPForestDict: TypeAlias = Mapping[str, Mapping[str, np.ndarray]]  # noqa: UP040
 
-N_DCS: int = 15  # default number of diffusion components
-# Backwards compat, constants should be defined in only one place.
-N_PCS: int = settings.N_PCS
+
+SCIPY_1_17 = pkg_version("scipy") >= Version("1.17")
 
 
 class KwdsForTransformer(TypedDict):
@@ -208,6 +208,10 @@ def neighbors(  # noqa: PLR0913
     :doc:`/how-to/knn-transformers`
 
     """
+    meta_random_state = (
+        dict(random_state=rng._arg) if isinstance(rng, _FakeRandomGen) else {}
+    )
+
     if distances is None:
         if metric is None:
             metric = "euclidean"
@@ -235,9 +239,8 @@ def neighbors(  # noqa: PLR0913
             if p.name in {"use_rep", "knn", "n_pcs", "metric_kwds"}
             if params[p.name] != p.default
         }
-        if not isinstance(rng, _FakeRandomGen) or rng._arg != 0:
+        if meta_random_state.get("random_state") != 0:  # rng or random_state was passed
             ignored.add("rng/random_state")
-            rng = _FakeRandomGen(0)
         if ignored:
             warn(
                 f"Parameter(s) ignored if `distances` is given: {ignored}",
@@ -270,8 +273,8 @@ def neighbors(  # noqa: PLR0913
         key_added,
         n_neighbors=neighbors_.n_neighbors,
         method=method,
-        random_state=_legacy_random_state(rng),
         metric=metric,
+        **meta_random_state,
         **({} if not metric_kwds else dict(metric_kwds=metric_kwds)),
         **({} if use_rep is None else dict(use_rep=use_rep)),
         **({} if n_pcs is None else dict(n_pcs=n_pcs)),
@@ -849,15 +852,13 @@ def compute_transitions(self, *, density_normalize: bool = True) -> None:
         self._transitions_sym = self.Z @ conn_norm @ self.Z
         logg.info("    finished", time=start)
 
-    @_accepts_legacy_random_state(0)
     def compute_eigen(
         self,
         *,
         n_comps: int = 15,
-        sym: bool | None = None,
         sort: Literal["decrease", "increase"] = "decrease",
         rng: np.random.Generator,
-    ):
+    ) -> None:
         """Compute eigen decomposition of transition matrix.
 
         Parameters
@@ -886,6 +887,9 @@ def compute_eigen(
             plotting.
 
         """
+        [rng_init, rng_eigsh] = np.random.default_rng(rng).spawn(2)
+        del rng
+
         np.set_printoptions(precision=10)
         if self._transitions_sym is None:
             msg = "Run `.compute_transitions` first."
@@ -903,9 +907,14 @@ def compute_eigen(
             matrix = matrix.astype(np.float64)
 
             # Setting the random initial vector
-            v0 = rng.standard_normal(matrix.shape[0])
+            v0 = rng_init.standard_normal(matrix.shape[0])
             evals, evecs = sparse.linalg.eigsh(
-                matrix, k=n_comps, which=which, ncv=ncv, v0=v0
+                matrix,
+                k=n_comps,
+                which=which,
+                ncv=ncv,
+                v0=v0,
+                **(dict(rng=rng_eigsh) if SCIPY_1_17 else {}),
             )
             evals, evecs = evals.astype(np.float32), evecs.astype(np.float32)
         if sort == "decrease":

src/scanpy/preprocessing/_pca/_compat.py

@@ -22,6 +22,9 @@
     from ..._utils.random import RNGLike, SeedLike
 
 
+SCIPY_1_15 = pkg_version("scikit-learn") >= Version("1.5.0rc1")
+
+
 @_accepts_legacy_random_state(None)
 def _pca_compat_sparse(
     x: CSBase,
@@ -33,7 +36,11 @@ def _pca_compat_sparse(
 ) -> tuple[NDArray[np.floating], PCA]:
     """Sparse PCA for scikit-learn <1.4."""
     rng = np.random.default_rng(rng)
-    random_init = rng.uniform(size=np.min(x.shape))
+    # this exists only to be stored in our PCA container object
+    random_state_meta = _legacy_random_state(rng)
+    [rng_init, rng_svds] = rng.spawn(2)
+    del rng
+
     x = check_array(x, accept_sparse=["csr", "csc"])
 
     if mu is None:
@@ -55,11 +62,15 @@ def rmat_op(v: NDArray[np.floating]):
         rmatmat=rmat_op,
     )
 
-    u, s, v = svds(linop, solver=solver, k=n_pcs, v0=random_init)
-    # u_based_decision was changed in https://github.com/scikit-learn/scikit-learn/pull/27491
-    u, v = svd_flip(
-        u, v, u_based_decision=pkg_version("scikit-learn") < Version("1.5.0rc1")
+    random_init = rng_init.uniform(size=np.min(x.shape))
+    kw = (
+        dict(rng=rng_svds)
+        if SCIPY_1_15
+        else dict(random_state=_legacy_random_state(rng_svds))
     )
+    u, s, v = svds(linop, solver=solver, k=n_pcs, v0=random_init, **kw)
+    # u_based_decision was changed in https://github.com/scikit-learn/scikit-learn/pull/27491
+    u, v = svd_flip(u, v, u_based_decision=not SCIPY_1_15)
     idx = np.argsort(-s)
     v = v[idx, :]
 
@@ -71,9 +82,7 @@ def rmat_op(v: NDArray[np.floating]):
 
     from sklearn.decomposition import PCA
 
-    pca = PCA(
-        n_components=n_pcs, svd_solver=solver, random_state=_legacy_random_state(rng)
-    )
+    pca = PCA(n_components=n_pcs, svd_solver=solver, random_state=random_state_meta)
     pca.explained_variance_ = ev
     pca.explained_variance_ratio_ = ev_ratio
     pca.components_ = v

src/scanpy/preprocessing/_scrublet/__init__.py

@@ -10,7 +10,7 @@
 
 from ... import logging as logg
 from ... import preprocessing as pp
-from ..._utils.random import _accepts_legacy_random_state, _legacy_random_state
+from ..._utils.random import _accepts_legacy_random_state, _FakeRandomGen
 from ...get import _get_obs_rep
 from . import pipeline
 from .core import Scrublet
@@ -177,10 +177,12 @@ def scrublet(  # noqa: PLR0913
 
     adata_obs = adata.copy()
 
-    def _run_scrublet(ad_obs: AnnData, ad_sim: AnnData | None = None):
+    def _run_scrublet(
+        ad_obs: AnnData, ad_sim: AnnData | None, *, rng: np.random.Generator
+    ):
+        rng_sim, rng_call = rng.spawn(2)
         # With no adata_sim we assume the regular use case, starting with raw
         # counts and simulating doublets
-
         if ad_sim is None:
             pp.filter_genes(ad_obs, min_cells=3)
             pp.filter_cells(ad_obs, min_genes=3)
@@ -207,7 +209,7 @@ def _run_scrublet(ad_obs: AnnData, ad_sim: AnnData | None = None):
                 layer="raw",
                 sim_doublet_ratio=sim_doublet_ratio,
                 synthetic_doublet_umi_subsampling=synthetic_doublet_umi_subsampling,
-                rng=rng,
+                rng=rng_sim,
             )
             del ad_obs.layers["raw"]
             if log_transform:
@@ -232,7 +234,7 @@ def _run_scrublet(ad_obs: AnnData, ad_sim: AnnData | None = None):
             knn_dist_metric=knn_dist_metric,
             get_doublet_neighbor_parents=get_doublet_neighbor_parents,
             threshold=threshold,
-            rng=rng,
+            rng=rng_call,
             verbose=verbose,
         )
 
@@ -249,12 +251,14 @@ def _run_scrublet(ad_obs: AnnData, ad_sim: AnnData | None = None):
         # Run Scrublet independently on batches and return just the
         # scrublet-relevant parts of the objects to add to the input object
         batches = np.unique(adata.obs[batch_key])
+        sub_rngs = rng.spawn(len(batches))
         scrubbed = [
             _run_scrublet(
                 adata_obs[adata_obs.obs[batch_key] == batch].copy(),
                 adata_sim,
+                rng=sub_rng,
             )
-            for batch in batches
+            for batch, sub_rng in zip(batches, sub_rngs, strict=True)
         ]
         scrubbed_obs = pd.concat([scrub["obs"] for scrub in scrubbed]).astype(
             adata.obs.dtypes
@@ -274,7 +278,7 @@ def _run_scrublet(ad_obs: AnnData, ad_sim: AnnData | None = None):
         adata.uns["scrublet"]["batched_by"] = batch_key
 
     else:
-        scrubbed = _run_scrublet(adata_obs, adata_sim)
+        scrubbed = _run_scrublet(adata_obs, adata_sim, rng=rng)
 
         # Copy outcomes to input object from our processed version
         adata.obs["doublet_score"] = scrubbed["obs"]["doublet_score"]
@@ -385,6 +389,12 @@ def _scrublet_call_doublets(  # noqa: PLR0913
         Dictionary of Scrublet parameters
 
     """
+    meta_random_state = (
+        dict(random_state=rng._arg) if isinstance(rng, _FakeRandomGen) else {}
+    )
+    rng_scrub, rng_pca = rng.spawn(2)
+    del rng
+
     # Estimate n_neighbors if not provided, and create scrublet object.
 
     if n_neighbors is None:
@@ -398,7 +408,7 @@ def _scrublet_call_doublets(  # noqa: PLR0913
         n_neighbors=n_neighbors,
         expected_doublet_rate=expected_doublet_rate,
         stdev_doublet_rate=stdev_doublet_rate,
-        rng=rng,
+        rng=rng_scrub,
     )
 
     # Ensure normalised matrix sparseness as Scrublet does
@@ -424,13 +434,11 @@ def _scrublet_call_doublets(  # noqa: PLR0913
 
     if mean_center:
         logg.info("Embedding transcriptomes using PCA...")
-        pipeline.pca(
-            scrub, n_prin_comps=n_prin_comps, svd_solver="arpack", rng=scrub._rng
-        )
+        pipeline.pca(scrub, n_prin_comps=n_prin_comps, svd_solver="arpack", rng=rng_pca)
     else:
         logg.info("Embedding transcriptomes using Truncated SVD...")
         pipeline.truncated_svd(
-            scrub, n_prin_comps=n_prin_comps, algorithm="arpack", rng=scrub._rng
+            scrub, n_prin_comps=n_prin_comps, algorithm="arpack", rng=rng_pca
         )
 
     # Score the doublets
@@ -463,7 +471,7 @@ def _scrublet_call_doublets(  # noqa: PLR0913
                 .get("sim_doublet_ratio", None)
             ),
             "n_neighbors": n_neighbors,
-            "random_state": _legacy_random_state(rng),
+            **meta_random_state,
         },
     }