diff --git a/.gitignore b/.gitignore index 8158195d..acebbf38 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ __pycache__/ /benchmarks/ test-data/ .vscode/ +PanGPA.log # Distribution / packaging /dist/ diff --git a/docs/release-notes/0.17.0.md b/docs/release-notes/0.17.0.md index 1706a90d..4006f7a2 100644 --- a/docs/release-notes/0.17.0.md +++ b/docs/release-notes/0.17.0.md @@ -8,6 +8,7 @@ * Split {func}`~rapids_singlecell.gr.calculate_niche` into {func}`~rapids_singlecell.gr.calculate_niche_neighborhood`, {func}`~rapids_singlecell.gr.calculate_niche_utag` and {func}`~rapids_singlecell.gr.calculate_niche_cellcharter`, with ``mask``, ``library_key`` and cross-flavor ``min_niche_size``, following {mod}`squidpy` {pr}`758` {smaller}`S Dicks` * Speed up {func}`~rapids_singlecell.pp.harmony_integrate` and make it reproducible by seeding k-means from a deterministic `float64` fit on a bounded, batch-stratified random subsample instead of a non-reproducible `float32` fit over all cells. `dtype` now defaults to `numpy.float32` {pr}`756` {smaller}`S Dicks` * Derive unset {func}`~rapids_singlecell.pp.harmony_integrate` stopping rules from `flavor`: `harmony2` follows Harmony2 defaults, `harmony1` still follows harmony-pytorch {pr}`756` {smaller}`S Dicks` +* Add support for {class}`numpy.random.Generator` to all functions previously accepting a ``random_state`` parameter, which is renamed to ``rng`` following `scanpy`. Passing ``random_state`` still works and keeps producing the same results; passing both raises {smaller}`S Dicks` ```{rubric} Misc ``` @@ -17,3 +18,7 @@ ```{rubric} Deprecations ``` * Deprecate {func}`~rapids_singlecell.gr.calculate_niche`; the new functions replace ``copy`` with ``inplace`` and drop ``gmm_init`` and the leiden ``random_state`` {pr}`758` {smaller}`S Dicks` + +```{rubric} Removals +``` +* Remove the ``random_seed`` parameter of {func}`~rapids_singlecell.pp.scrublet_simulate_doublets`, following `scanpy`. Use ``rng`` (or ``random_state``) instead {smaller}`S Dicks` diff --git a/src/rapids_singlecell/_utils/_random.py b/src/rapids_singlecell/_utils/_random.py new file mode 100644 index 00000000..85c62dd3 --- /dev/null +++ b/src/rapids_singlecell/_utils/_random.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +from collections.abc import Sequence +from functools import wraps +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from collections.abc import Callable + +__all__ = [ + "RNGLike", + "SeedLike", + "_LegacyRandom", + "_LegacyRng", + "_accepts_legacy_random_state", + "_legacy_random_state", + "_seed_from_rng", +] + +type SeedLike = int | np.integer | Sequence[int] | np.random.SeedSequence +type RNGLike = np.random.Generator | np.random.BitGenerator +type _LegacyRandom = int | np.random.RandomState | None + +_SEED_BOUND = 2**32 +"""cuML, cuGraph and CuPy all take a 32-bit unsigned seed.""" + + +class _LegacyRng: + """Marks a seed that arrived through the superseded `random_state` argument. + + Unlike scanpy's class of the same name, this is only a marker. Integer-only + GPU consumers use it to recover the legacy seed, while host-side consumers + can recover the original :class:`~numpy.random.RandomState` object and + continue its exact stream. + """ + + __slots__ = ("arg",) + + def __init__(self, arg: _LegacyRandom) -> None: + self.arg = arg + + def __repr__(self) -> str: + return f"_LegacyRng({self.arg!r})" + + +def _accepts_legacy_random_state[**P, R]( + default: _LegacyRandom, / +) -> Callable[[Callable[P, R]], Callable[P, R]]: + """Let a function taking `rng` still be called with `random_state`. + + A `random_state` argument is wrapped in a :class:`_LegacyRng` and passed as + `rng`. Passing both is an error. If neither is given, `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: + raise TypeError("Specify at most one of `rng` and `random_state`.") + case True, False: + kwargs["rng"] = _LegacyRng(kwargs.pop("random_state")) + case False, False: + kwargs["rng"] = _LegacyRng(default) + return func(*args, **kwargs) + + return wrapper + + return decorator + + +def _seed_from_rng(rng: SeedLike | RNGLike | _LegacyRng | None, /) -> int | None: + """The integer seed to hand to cuML, cuGraph, CuPy or a kernel. + + A `random_state` integer is forwarded untouched so that existing calls keep + producing the same results; anything else is drawn from the generator. + """ + if isinstance(rng, _LegacyRng): + if rng.arg is None: + return None + if isinstance(rng.arg, (int, np.integer)): + return int(rng.arg) + return int(rng.arg.randint(0, _SEED_BOUND)) + return int(np.random.default_rng(rng).integers(0, _SEED_BOUND)) + + +def _legacy_random_state( + rng: SeedLike | RNGLike | _LegacyRng | None, / +) -> _LegacyRandom: + """The value to hand to host-side APIs that still take a `random_state`.""" + if isinstance(rng, _LegacyRng): + return rng.arg + [bit_generator] = np.random.default_rng(rng).bit_generator.spawn(1) + return np.random.RandomState(bit_generator) diff --git a/src/rapids_singlecell/preprocessing/_harmony_integrate.py b/src/rapids_singlecell/preprocessing/_harmony_integrate.py index 19b9a6a5..4056dc4b 100644 --- a/src/rapids_singlecell/preprocessing/_harmony_integrate.py +++ b/src/rapids_singlecell/preprocessing/_harmony_integrate.py @@ -6,12 +6,20 @@ import cupy as cp import numpy as np +from rapids_singlecell._utils._random import ( + RNGLike, + SeedLike, + _accepts_legacy_random_state, + _seed_from_rng, +) + if TYPE_CHECKING: from anndata import AnnData from ._harmony import COLSUM_ALGO +@_accepts_legacy_random_state(0) def harmony_integrate( adata: AnnData, key: str | list[str], @@ -34,7 +42,7 @@ def harmony_integrate( correction_method: Literal["fast", "batched"] | None = None, colsum_algo: COLSUM_ALGO | None = None, block_proportion: float = 0.05, - random_state: int = 0, + rng: SeedLike | RNGLike | None = None, verbose: bool = False, ) -> None: """Integrate different experiments using the Harmony algorithm :cite:p:`Korsunsky2019,Patikas2026`. @@ -158,8 +166,9 @@ def harmony_integrate( Proportion of cells updated per clustering sub-iteration. Smaller values produce more stochastic updates. Larger values are faster but may converge to different solutions. - random_state - Random seed for reproducibility. + rng + Random seed or :class:`~numpy.random.Generator` for reproducibility. + The superseded `random_state` argument is still accepted. verbose Whether to print benchmarking and convergence information. @@ -169,6 +178,8 @@ def harmony_integrate( containing principal components adjusted by Harmony such that different experiments are integrated. """ + random_state = _seed_from_rng(rng) + from ._harmony import harmonize # Resolve flavor into internal flags diff --git a/src/rapids_singlecell/preprocessing/_neighbors/__init__.py b/src/rapids_singlecell/preprocessing/_neighbors/__init__.py index 00d74bbf..f1c57b9d 100644 --- a/src/rapids_singlecell/preprocessing/_neighbors/__init__.py +++ b/src/rapids_singlecell/preprocessing/_neighbors/__init__.py @@ -6,6 +6,12 @@ import cupy as cp import numpy as np +from rapids_singlecell._utils._random import ( + RNGLike, + SeedLike, + _accepts_legacy_random_state, + _seed_from_rng, +) from rapids_singlecell.preprocessing._neighbors._helper import ( _check_metrics, _check_neighbors_X, @@ -32,13 +38,14 @@ ] +@_accepts_legacy_random_state(0) def neighbors( adata: AnnData, n_neighbors: int = 15, n_pcs: int | None = None, *, use_rep: str | None = None, - random_state: AnyRandom = 0, + rng: SeedLike | RNGLike | None = None, algorithm: _Algorithms = "brute", metric: _Metrics = "euclidean", metric_kwds: Mapping[str, Any] = MappingProxyType({}), @@ -70,8 +77,9 @@ def neighbors( If None, the representation is chosen automatically: For .n_vars < 50, .X is used, otherwise `'X_pca'` is used. If `'X_pca'` is not present, it's computed with default parameters or `n_pcs` if present. - random_state - A numpy random seed. + rng + Random seed or :class:`~numpy.random.Generator` for reproducibility. + The superseded `random_state` argument is still accepted. algorithm The query algorithm to use. Valid options are: `'brute'` @@ -174,6 +182,8 @@ def neighbors( neighbors. """ + random_state = _seed_from_rng(rng) + adata = adata.copy() if copy else adata if adata.is_view: @@ -246,6 +256,7 @@ def neighbors( return adata if copy else None +@_accepts_legacy_random_state(0) def bbknn( adata: AnnData, neighbors_within_batch: int = 3, @@ -253,7 +264,7 @@ def bbknn( *, batch_key: str | None = None, use_rep: str | None = None, - random_state: AnyRandom = 0, + rng: SeedLike | RNGLike | None = None, algorithm: _Algorithms_bbknn = "brute", metric: _Metrics = "euclidean", metric_kwds: Mapping[str, Any] = MappingProxyType({}), @@ -284,8 +295,9 @@ def bbknn( If `None`, the representation is chosen automatically: For `.n_vars < 50`, `.X` is used, otherwise `'X_pca'` is used. If `'X_pca'` is not present, it's computed with default parameters or `n_pcs` if present. - random_state - A numpy random seed. + rng + Random seed or :class:`~numpy.random.Generator` for reproducibility. + The superseded `random_state` argument is still accepted. algorithm The query algorithm to use. Valid options are: @@ -359,6 +371,8 @@ def bbknn( connectivities and distances. """ + random_state = _seed_from_rng(rng) + if batch_key is None: raise ValueError("Please provide a batch key to perform batch-balanced KNN.") diff --git a/src/rapids_singlecell/preprocessing/_pca.py b/src/rapids_singlecell/preprocessing/_pca.py index 3758b835..f863d977 100644 --- a/src/rapids_singlecell/preprocessing/_pca.py +++ b/src/rapids_singlecell/preprocessing/_pca.py @@ -12,6 +12,12 @@ from scipy.sparse import issparse from rapids_singlecell._compat import DaskArray +from rapids_singlecell._utils._random import ( + RNGLike, + SeedLike, + _accepts_legacy_random_state, + _seed_from_rng, +) from rapids_singlecell.get import X_to_GPU, _check_mask, _get_obs_rep from ._utils import _check_gpu_X @@ -62,6 +68,7 @@ def _resolve_mask_var( return mask_var_param, _check_mask(adata, mask_var, "var") +@_accepts_legacy_random_state(0) def pca( data: AnnData | ArrayTypesDask, n_comps: int | None = None, @@ -69,7 +76,7 @@ def pca( layer: str = None, zero_center: bool = True, svd_solver: str | None = None, - random_state: int | None = 0, + rng: SeedLike | RNGLike | None = None, mask_var: NDArray[np.bool] | str | None = _empty, use_highly_variable: bool | None = None, dtype: str = "float32", @@ -145,8 +152,9 @@ def pca( `'jacobi'` cuML: Jacobi iterative solver. Faster but less accurate. For dense arrays only. - random_state - Random state for initialization. + rng + Random seed or :class:`~numpy.random.Generator` for initialization. + The superseded `random_state` argument is still accepted. mask_var Mask to use for the PCA computation. @@ -216,6 +224,8 @@ def pca( Explained variance, equivalent to the eigenvalues of the \ covariance matrix. """ + random_state = _seed_from_rng(rng) + if not isinstance(data, AnnData): if layer is not None: raise ValueError("`layer` can only be used with an AnnData object.") diff --git a/src/rapids_singlecell/preprocessing/_scrublet/__init__.py b/src/rapids_singlecell/preprocessing/_scrublet/__init__.py index 867cf990..34fdcd3e 100644 --- a/src/rapids_singlecell/preprocessing/_scrublet/__init__.py +++ b/src/rapids_singlecell/preprocessing/_scrublet/__init__.py @@ -10,6 +10,12 @@ from scanpy import logging as logg from rapids_singlecell import preprocessing as pp +from rapids_singlecell._utils._random import ( + RNGLike, + SeedLike, + _accepts_legacy_random_state, + _LegacyRng, +) from rapids_singlecell.get import _get_obs_rep from . import pipeline @@ -20,6 +26,18 @@ from rapids_singlecell.preprocessing._neighbors import _Metrics +type _ScrubletRandom = int | np.random.RandomState | np.random.Generator | None + + +def _normalize_scrublet_rng( + rng: SeedLike | RNGLike | _LegacyRng | None, / +) -> _ScrubletRandom: + if isinstance(rng, _LegacyRng): + return rng.arg + return np.random.default_rng(rng) + + +@_accepts_legacy_random_state(0) def scrublet( adata: AnnData, adata_sim: AnnData | None = None, @@ -40,7 +58,7 @@ def scrublet( threshold: float | None = None, verbose: bool = True, copy: bool = False, - random_state: AnyRandom = 0, + rng: SeedLike | RNGLike | None = None, ) -> AnnData | None: """\ Predict doublets using Scrublet :cite:p:`Wolock2019`. @@ -124,8 +142,10 @@ def scrublet( copy If :data:`True`, return a copy of the input ``adata`` with Scrublet results added. Otherwise, Scrublet results are added in place. - random_state - Initial state for doublet simulation and nearest neighbors. + rng + Random seed or :class:`~numpy.random.Generator` for doublet simulation + and nearest neighbors. + The superseded `random_state` argument is still accepted. Returns ------- @@ -155,6 +175,8 @@ def scrublet( scores for observed transcriptomes and simulated doublets. """ + rng = _normalize_scrublet_rng(rng) + if copy: adata = adata.copy() @@ -162,7 +184,12 @@ def scrublet( 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 = None, + *, + rng: _ScrubletRandom, + ): # With no adata_sim we assume the regular use case, starting with raw # counts and simulating doublets @@ -187,12 +214,17 @@ def _run_scrublet(ad_obs: AnnData, ad_sim: AnnData | None = None): # Simulate the doublets based on the raw expressions from the normalised # and filtered object. + simulation_rng = ( + {"rng": rng} + if isinstance(rng, np.random.Generator) + else {"random_state": rng} + ) ad_sim = scrublet_simulate_doublets( ad_obs, layer="raw", sim_doublet_ratio=sim_doublet_ratio, synthetic_doublet_umi_subsampling=synthetic_doublet_umi_subsampling, - random_seed=random_state, + **simulation_rng, ) if log_transform: @@ -217,7 +249,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, - random_state=random_state, + random_state=rng, verbose=verbose, ) @@ -233,12 +265,18 @@ def _run_scrublet(ad_obs: AnnData, ad_sim: AnnData | None = None): # 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)) + if isinstance(rng, np.random.Generator) + else [rng] * 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]) @@ -259,7 +297,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 @@ -286,7 +324,7 @@ def _scrublet_call_doublets( knn_dist_metric: _Metrics = "euclidean", get_doublet_neighbor_parents: bool = False, threshold: float | None = None, - random_state: AnyRandom = 0, + random_state: AnyRandom | np.random.Generator = 0, verbose: bool = True, ) -> AnnData: """\ @@ -368,6 +406,12 @@ def _scrublet_call_doublets( Dictionary of Scrublet parameters """ + meta_random_state = ( + {} + if isinstance(random_state, np.random.Generator) + else {"random_state": random_state} + ) + # Estimate n_neighbors if not provided, and create scrublet object. if n_neighbors is None: @@ -443,7 +487,7 @@ def _scrublet_call_doublets( .get("sim_doublet_ratio", None) ), "n_neighbors": n_neighbors, - "random_state": random_state, + **meta_random_state, }, } @@ -468,13 +512,14 @@ def _scrublet_call_doublets( return adata_obs +@_accepts_legacy_random_state(0) def scrublet_simulate_doublets( adata: AnnData, *, layer: str | None = None, sim_doublet_ratio: float = 2.0, synthetic_doublet_umi_subsampling: float = 1.0, - random_seed: AnyRandom = 0, + rng: SeedLike | RNGLike | None = None, ) -> AnnData: """ Simulate doublets by adding the counts of random observed transcriptome pairs. @@ -517,8 +562,10 @@ def scrublet_simulate_doublets( scores for observed transcriptomes and simulated doublets. """ + rng = _normalize_scrublet_rng(rng) + X = _get_obs_rep(adata, layer=layer) - scrub = Scrublet(X, random_state=random_seed) + scrub = Scrublet(X, random_state=rng) scrub.simulate_doublets( sim_doublet_ratio=sim_doublet_ratio, diff --git a/src/rapids_singlecell/preprocessing/_scrublet/core.py b/src/rapids_singlecell/preprocessing/_scrublet/core.py index 213ff116..e17751dc 100644 --- a/src/rapids_singlecell/preprocessing/_scrublet/core.py +++ b/src/rapids_singlecell/preprocessing/_scrublet/core.py @@ -16,7 +16,7 @@ from .sparse_utils import subsample_counts if TYPE_CHECKING: - from numpy.random import RandomState + from numpy.random import Generator, RandomState from numpy.typing import NDArray from rapids_singlecell._utils import AnyRandom @@ -70,12 +70,12 @@ class Scrublet: n_neighbors: InitVar[int | None] = None expected_doublet_rate: float = 0.1 stdev_doublet_rate: float = 0.02 - random_state: InitVar[AnyRandom] = 0 + random_state: InitVar[AnyRandom | Generator] = 0 # private fields _n_neighbors: int = field(init=False, repr=False) - _random_state: RandomState = field(init=False, repr=False) + _random_state: RandomState | Generator = field(init=False, repr=False) _counts_obs: sparse.csc_matrix = field(init=False, repr=False) _total_counts_obs: NDArray[np.integer] = field(init=False, repr=False) @@ -171,7 +171,7 @@ def __post_init__( counts_obs: sparse.csr_matrix | sparse.csc_matrix | NDArray[np.integer], total_counts_obs: NDArray[np.integer] | None, n_neighbors: int | None, - random_state: AnyRandom, + random_state: AnyRandom | Generator, ) -> None: self._counts_obs = sparse.csc_matrix(counts_obs) self._total_counts_obs = ( @@ -184,7 +184,11 @@ def __post_init__( if n_neighbors is None else n_neighbors ) - self._random_state = get_random_state(random_state) + self._random_state = ( + random_state + if isinstance(random_state, np.random.Generator) + else get_random_state(random_state) + ) def simulate_doublets( self, @@ -220,11 +224,15 @@ def simulate_doublets( n_obs = self._counts_obs.shape[0] n_sim = int(n_obs * sim_doublet_ratio) - pair_ix = sample_comb( - (n_obs, n_obs), - n_sim, - **_random_state_kwargs(sample_comb, self._random_state), - ) + if isinstance(self._random_state, np.random.Generator): + flat_ix = self._random_state.choice(n_obs**2, size=n_sim, replace=False) + pair_ix = np.vstack(np.unravel_index(flat_ix, (n_obs, n_obs))).T + else: + pair_ix = sample_comb( + (n_obs, n_obs), + n_sim, + **_random_state_kwargs(sample_comb, self._random_state), + ) E1 = cast("sparse.csc_matrix", self._counts_obs[pair_ix[:, 0], :]) E2 = cast("sparse.csc_matrix", self._counts_obs[pair_ix[:, 1], :]) diff --git a/src/rapids_singlecell/preprocessing/_scrublet/sparse_utils.py b/src/rapids_singlecell/preprocessing/_scrublet/sparse_utils.py index eec45a15..34f2651c 100644 --- a/src/rapids_singlecell/preprocessing/_scrublet/sparse_utils.py +++ b/src/rapids_singlecell/preprocessing/_scrublet/sparse_utils.py @@ -9,10 +9,9 @@ from rapids_singlecell.preprocessing._utils import _get_mean_var, get_random_state if TYPE_CHECKING: + from numpy.random import Generator, RandomState from numpy.typing import NDArray - from rapids_singlecell._utils import AnyRandom - def sparse_multiply( E: sparse.csr_matrix | sparse.csc_matrix | NDArray[np.float64], @@ -47,21 +46,31 @@ def subsample_counts( *, rate: float, original_totals, - random_seed: AnyRandom = 0, + random_seed: int | RandomState | Generator | None = 0, ) -> tuple[sparse.csr_matrix | sparse.csc_matrix, NDArray[np.int64]]: if rate < 1: - random_seed = get_random_state(random_seed) + is_generator = isinstance(random_seed, np.random.Generator) + if not is_generator: + random_seed = get_random_state(random_seed) dtype = E.dtype E.data = cp.array( random_seed.binomial(np.round(E.data.get()).astype(int), rate), dtype=dtype ) current_totals = E.sum(1).ravel() unsampled_orig_totals = original_totals - current_totals - unsampled_downsamp_totals = cp.random.binomial( - cp.round(unsampled_orig_totals).astype(int), - rate, - dtype=dtype, - ) + if is_generator: + unsampled_downsamp_totals = cp.asarray( + random_seed.binomial( + np.round(unsampled_orig_totals.get()).astype(int), rate + ), + dtype=dtype, + ) + else: + unsampled_downsamp_totals = cp.random.binomial( + cp.round(unsampled_orig_totals).astype(int), + rate, + dtype=dtype, + ) final_downsamp_totals = current_totals + unsampled_downsamp_totals else: final_downsamp_totals = original_totals diff --git a/src/rapids_singlecell/tools/_clustering.py b/src/rapids_singlecell/tools/_clustering.py index 051e001c..af6b0b11 100644 --- a/src/rapids_singlecell/tools/_clustering.py +++ b/src/rapids_singlecell/tools/_clustering.py @@ -12,6 +12,13 @@ from scanpy.tools._utils import _choose_graph from scanpy.tools._utils_clustering import rename_groups, restrict_adjacency +from rapids_singlecell._utils._random import ( + RNGLike, + SeedLike, + _accepts_legacy_random_state, + _seed_from_rng, +) + from ._utils import _choose_representation if TYPE_CHECKING: @@ -127,11 +134,12 @@ def mapper(pair): return g +@_accepts_legacy_random_state(0) def leiden( adata: AnnData, resolution: float | list[float] = 1.0, *, - random_state: int | None = 0, + rng: SeedLike | RNGLike | None = None, theta: float = 1.0, restrict_to: tuple[str, Sequence[str]] | None = None, key_added: str = "leiden", @@ -160,8 +168,10 @@ def leiden( (called gamma in the modularity formula). Higher values lead to more clusters. If a list of values is provided, the Leiden algorithm will be run for each value in the list. - random_state - Change the initialization of the optimization. Defaults to 0. + rng + Random seed or :class:`~numpy.random.Generator` changing the + initialization of the optimization. Defaults to 0. + The superseded `random_state` argument is still accepted. theta Called theta in the Leiden algorithm, this is used to scale modularity @@ -211,6 +221,8 @@ def leiden( """ # Adjacency graph + random_state = _seed_from_rng(rng) + adata = adata.copy() if copy else adata dtype = _check_dtype(dtype) @@ -462,6 +474,7 @@ def louvain( return adata if copy else None +@_accepts_legacy_random_state(42) def kmeans( adata: AnnData, n_clusters: int = 8, @@ -469,7 +482,7 @@ def kmeans( *, use_rep: str = "X_pca", n_init: int = 1, - random_state: float = 42, + rng: SeedLike | RNGLike | None = None, key_added: str = "kmeans", copy: bool = False, **kwargs, @@ -492,9 +505,10 @@ def kmeans( computed with default parameters or `n_pcs` if present. n_init Number of initializations to run the KMeans algorithm - random_state - if you want results to be the same when you restart Python, select a - state. Default is 42. + rng + Random seed or :class:`~numpy.random.Generator`; fix it if you want + results to be the same when you restart Python. Default is 42. + The superseded `random_state` argument is still accepted. key_added `adata.obs` key under which to add the cluster labels. copy @@ -503,6 +517,8 @@ def kmeans( Additional keyword arguments for KMeans. """ + random_state = _seed_from_rng(rng) + from cuml.cluster import KMeans adata = adata.copy() if copy else adata diff --git a/src/rapids_singlecell/tools/_draw_graph.py b/src/rapids_singlecell/tools/_draw_graph.py index faacbbaf..b31332c4 100644 --- a/src/rapids_singlecell/tools/_draw_graph.py +++ b/src/rapids_singlecell/tools/_draw_graph.py @@ -8,6 +8,12 @@ from scanpy.tools._utils import get_init_pos_from_paga from rapids_singlecell._compat import _random_state_kwargs +from rapids_singlecell._utils._random import ( + RNGLike, + SeedLike, + _accepts_legacy_random_state, + _seed_from_rng, +) from ._clustering import _create_graph from ._utils import _validate_init_pos @@ -16,12 +22,13 @@ from anndata import AnnData +@_accepts_legacy_random_state(0) def draw_graph( adata: AnnData, *, init_pos: str | bool | None = None, max_iter: int = 500, - random_state: int | None = 0, + rng: SeedLike | RNGLike | None = None, ) -> None: """ Force-directed graph drawing :cite:p:`Fruchterman1991,Jacomy2014`. @@ -45,10 +52,12 @@ def draw_graph( No error occurs when the algorithm terminates in this manner. Good short-term quality can be achieved with 50-100 iterations. Above 1000 iterations is discouraged. - random_state - Random state to use when initializing layout and generating - samples. Defaults to 0. If `None` is passed, a hash of process id, - time, and hostname is used by `cugraph`. + rng + Random seed or :class:`~numpy.random.Generator` used when + initializing layout and generating samples. Defaults to 0. If + `None` is passed, a hash of process id, time, and hostname is + used by `cugraph`. + The superseded `random_state` argument is still accepted. Returns ------- @@ -57,6 +66,8 @@ def draw_graph( X_draw_graph_layout_fa : `adata.obsm` Coordinates of graph layout. """ + random_state = _seed_from_rng(rng) + from cugraph.layout import force_atlas2 # Adjacency graph diff --git a/src/rapids_singlecell/tools/_score_genes.py b/src/rapids_singlecell/tools/_score_genes.py index d905a7c6..1167aa16 100644 --- a/src/rapids_singlecell/tools/_score_genes.py +++ b/src/rapids_singlecell/tools/_score_genes.py @@ -9,6 +9,13 @@ import pandas as pd from rapids_singlecell._compat import DaskArray +from rapids_singlecell._utils._random import ( + RNGLike, + SeedLike, + _accepts_legacy_random_state, + _LegacyRng, + _seed_from_rng, +) from rapids_singlecell.get import X_to_GPU, _get_obs_rep from rapids_singlecell.preprocessing._utils import _check_gpu_X, _check_use_raw @@ -20,6 +27,7 @@ from anndata import AnnData +@_accepts_legacy_random_state(0) def score_genes( adata: AnnData, gene_list: Sequence[str] | pd.Index, @@ -29,7 +37,7 @@ def score_genes( gene_pool: Sequence[str] | pd.Index | None = None, n_bins: int = 25, score_name: str = "score", - random_state: int | None = 0, + rng: SeedLike | RNGLike | None = None, copy: bool = False, use_raw: bool | None = None, layer: str | None = None, @@ -59,8 +67,9 @@ def score_genes( Number of expression level bins for sampling. score_name Name of the field to be added in `.obs`. - random_state - The random seed for sampling. + rng + Random seed or :class:`~numpy.random.Generator` for sampling. + The superseded `random_state` argument is still accepted. copy Copy `adata` or modify it inplace. use_raw @@ -80,8 +89,12 @@ def score_genes( X = _get_obs_rep(adata, layer=layer, use_raw=use_raw) X = X_to_GPU(X) _check_gpu_X(X, allow_dask=True) - if random_state is not None: - np.random.seed(random_state) + sample_rng = None + if isinstance(rng, _LegacyRng): + if rng.arg is not None: + np.random.seed(_seed_from_rng(rng)) + else: + sample_rng = np.random.default_rng(rng) var_names = adata.raw.var_names if use_raw else adata.var_names gene_list, gene_pool = _check_score_genes_args(var_names, gene_list, gene_pool) @@ -95,6 +108,7 @@ def score_genes( ctrl_as_ref=ctrl_as_ref, ctrl_size=ctrl_size, n_bins=n_bins, + sample_rng=sample_rng, ): control_genes = control_genes.union(r_genes) @@ -157,6 +171,7 @@ def _score_genes_bins( ctrl_as_ref: bool, ctrl_size: int, n_bins: int, + sample_rng: np.random.Generator | None = None, ) -> Generator[pd.Index[str], None, None]: # average expression of genes idx = cp.array(var_names.isin(gene_pool), dtype=cp.bool_) @@ -173,7 +188,10 @@ def _score_genes_bins( keep_ctrl_in_obs_cut = False if ctrl_as_ref else obs_cut.index.isin(gene_list) # now pick `ctrl_size` genes from every cut - for cut in np.unique(obs_cut.loc[gene_list]): + cuts = np.unique(obs_cut.loc[gene_list]) + # spawn a sub-rng per cut, like scanpy, so this stays parallelizable + sub_rngs = [None] * len(cuts) if sample_rng is None else sample_rng.spawn(len(cuts)) + for cut, sub_rng in zip(cuts, sub_rngs, strict=True): r_genes: pd.Index[str] = obs_cut[(obs_cut == cut) & ~keep_ctrl_in_obs_cut].index if len(r_genes) == 0: msg = ( @@ -182,7 +200,7 @@ def _score_genes_bins( ) warnings.warn(msg) if ctrl_size < len(r_genes): - r_genes = r_genes.to_series().sample(ctrl_size).index + r_genes = r_genes.to_series().sample(ctrl_size, random_state=sub_rng).index if ctrl_as_ref: # otherwise `r_genes` is already filtered r_genes = r_genes.difference(gene_list) yield r_genes diff --git a/src/rapids_singlecell/tools/_umap.py b/src/rapids_singlecell/tools/_umap.py index 0d6ad896..311e12d2 100644 --- a/src/rapids_singlecell/tools/_umap.py +++ b/src/rapids_singlecell/tools/_umap.py @@ -15,6 +15,13 @@ from rapids_singlecell._compat import _random_state_kwargs from rapids_singlecell._utils import _get_logger_level +from rapids_singlecell._utils._random import ( + RNGLike, + SeedLike, + _accepts_legacy_random_state, + _legacy_random_state, + _LegacyRng, +) from ._utils import _choose_representation, _validate_init_pos @@ -24,6 +31,7 @@ _InitPos = Literal["auto", "spectral", "random", "paga"] +@_accepts_legacy_random_state(0) def umap( adata: AnnData, *, @@ -34,7 +42,7 @@ def umap( alpha: float = 1.0, negative_sample_rate: int = 5, init_pos: _InitPos | np.ndarray | cp.ndarray | str | None = "auto", - random_state: int = 0, + rng: SeedLike | RNGLike | None = None, a: float | None = None, b: float | None = None, key_added: str | None = None, @@ -90,8 +98,10 @@ def umap( .. note:: If your embedding looks odd it's recommended setting `init_pos` to 'random'. - random_state - `int`, `random_state` is the seed used by the random number generator + rng + Random seed or :class:`~numpy.random.Generator` used by the random + number generator. + The superseded `random_state` argument is still accepted. a More specific parameters controlling the embedding. If `None` these values are set automatically as determined by `min_dist` and @@ -126,6 +136,8 @@ def umap( UMAP parameters `a`, `b`, and `random_state` (if specified). """ + random_state = _legacy_random_state(rng) + adata = adata.copy() if copy else adata if neighbors_key is None: @@ -144,7 +156,11 @@ def umap( stored_params = { "a": a, "b": b, - **({"random_state": random_state} if random_state != 0 else {}), + **( + {"random_state": rng.arg} + if isinstance(rng, _LegacyRng) and rng.arg != 0 + else {} + ), } neigh_params = neighbors["params"]