gh-961: use rng_dispatcher throughout code (#962)

paddyroddy · connoraird · web-flow · commit 8f4bb6b711eb · 2026-01-15T17:39:38.000Z
Co-authored-by: Connor Aird &lt;c.aird@ucl.ac.uk&gt;
diff --git a/glass/_array_api_utils.py b/glass/_array_api_utils.py
@@ -29,6 +29,8 @@
 
     from glass._types import AnyArray, FloatArray, UnifiedGenerator
 
+SEED = 42
+
 
 class CompatibleBackendNotFoundError(Exception):
     """
@@ -83,12 +85,14 @@ def default_xp() -> ModuleType:
     return import_numpy()
 
 
-def rng_dispatcher(*, xp: ModuleType) -> UnifiedGenerator:
+def rng_dispatcher(*, seed: int | AArray = SEED, xp: ModuleType) -> UnifiedGenerator:
     """
     Dispatch a random number generator based on the provided array's backend.
 
     Parameters
     ----------
+    seed
+        Seed for the random number generator.
     xp
         The array library backend to use for array operations.
 
@@ -101,8 +105,6 @@ def rng_dispatcher(*, xp: ModuleType) -> UnifiedGenerator:
     NotImplementedError
         If the array backend is not supported.
     """
-    seed = 42
-
     if xp.__name__ == "jax.numpy":
         import glass.jax  # noqa: PLC0415
 
@@ -128,10 +130,7 @@ class Generator:
 
     __slots__ = ("axp", "nxp", "rng")
 
-    def __init__(
-        self,
-        seed: int | bool | AArray | None = None,  # noqa: FBT001
-    ) -> None:
+    def __init__(self, *, seed: int | AArray = SEED) -> None:
         """
         Initialize the Generator.
 
@@ -146,7 +145,7 @@ def __init__(
 
         self.axp = array_api_strict
         self.nxp = numpy
-        self.rng = self.nxp.random.default_rng(seed=seed)
+        self.rng = rng_dispatcher(seed=seed, xp=self.nxp)
 
     def random(
         self,
@@ -171,7 +170,7 @@ def random(
             Array of random floats.
         """
         dtype = dtype if dtype is not None else self.nxp.float64
-        return self.axp.asarray(self.rng.random(size, dtype, out))  # type: ignore[arg-type]
+        return self.axp.asarray(self.rng.random(size, dtype, out))  # type: ignore[arg-type,call-arg]
 
     def normal(
         self,
@@ -216,7 +215,7 @@ def poisson(
         -------
             Array of samples from the Poisson distribution.
         """
-        return self.axp.asarray(self.rng.poisson(lam, size))
+        return self.axp.asarray(self.rng.poisson(lam, size))  # type: ignore[arg-type]
 
     def standard_normal(
         self,
@@ -241,7 +240,7 @@ def standard_normal(
             Array of samples from the standard normal distribution.
         """
         dtype = dtype if dtype is not None else self.nxp.float64
-        return self.axp.asarray(self.rng.standard_normal(size, dtype, out))  # type: ignore[arg-type]
+        return self.axp.asarray(self.rng.standard_normal(size, dtype, out))  # type: ignore[arg-type,call-arg]
 
     def uniform(
         self,
@@ -265,7 +264,7 @@ def uniform(
         -------
             Array of samples from the uniform distribution.
         """
-        return self.axp.asarray(self.rng.uniform(low, high, size))
+        return self.axp.asarray(self.rng.uniform(low, high, size))  # type: ignore[arg-type]
 
 
 class XPAdditions:
diff --git a/glass/fields.py b/glass/fields.py
@@ -25,7 +25,14 @@
     from collections.abc import Callable, Generator, Iterable, Iterator, Sequence
     from typing import Literal
 
-    from glass._types import AnyArray, ComplexArray, FloatArray, IntArray, T
+    from glass._types import (
+        AnyArray,
+        ComplexArray,
+        FloatArray,
+        IntArray,
+        T,
+        UnifiedGenerator,
+    )
 
 
 try:
@@ -318,7 +325,7 @@ def _generate_grf(
     nside: int,
     *,
     ncorr: int | None = None,
-    rng: np.random.Generator | None = None,
+    rng: UnifiedGenerator | None = None,
 ) -> Generator[FloatArray]:
     """
     Iteratively sample Gaussian random fields (internal use).
@@ -356,7 +363,7 @@ def _generate_grf(
         If all gls are empty.
     """
     if rng is None:
-        rng = np.random.default_rng(42)
+        rng = _utils.rng_dispatcher(xp=np)
 
     # number of gls and number of fields
     ngls = len(gls)  # type: ignore[arg-type]
@@ -386,7 +393,7 @@ def _generate_grf(
     for j, a, s in conditional_dist:
         # standard normal random variates for alm
         # sample real and imaginary parts, then view as complex number
-        rng.standard_normal(n * (n + 1), np.float64, z.view(np.float64))
+        rng.standard_normal(n * (n + 1), np.float64, z.view(np.float64))  # type: ignore[arg-type,call-arg]
 
         # scale by standard deviation of the conditional distribution
         # variance is distributed over real and imaginary part
diff --git a/tests/benchmarks/test_fields.py b/tests/benchmarks/test_fields.py
@@ -175,7 +175,7 @@ def function_to_benchmark() -> list[Any]:
         generator = glass.fields._generate_grf(
             gls,
             nside,
-            rng=urngb if use_rng else None,  # type: ignore[arg-type]
+            rng=urngb if use_rng else None,
             ncorr=ncorr,
         )
         return generator_consumer.consume(generator)  # type: ignore[no-any-return]
diff --git a/tests/core/test_array_api_utils.py b/tests/core/test_array_api_utils.py
@@ -13,7 +13,6 @@
 # check if available for testing
 HAVE_ARRAY_API_STRICT = importlib.util.find_spec("array_api_strict") is not None
 HAVE_JAX = importlib.util.find_spec("jax") is not None
-SEED = 42
 
 
 def test_rng_dispatcher_numpy() -> None:
@@ -39,61 +38,59 @@ def test_rng_dispatcher_array_api_strict() -> None:
 
 @pytest.mark.skipif(not HAVE_ARRAY_API_STRICT, reason="test requires array_api_strict")
 def test_init() -> None:
-    rng = _utils.Generator(SEED)
+    rng = _utils.Generator()
     assert isinstance(rng, _utils.Generator)
 
 
 @pytest.mark.skipif(not HAVE_ARRAY_API_STRICT, reason="test requires array_api_strict")
 def test_random() -> None:
     import array_api_strict
-    from array_api_strict._array_object import Array
 
-    rng = _utils.Generator(SEED)
+    rng = _utils.rng_dispatcher(xp=array_api_strict)
     rvs = rng.random(size=10_000)
     assert rvs.shape == (10_000,)
     assert array_api_strict.min(rvs) >= 0.0
     assert array_api_strict.max(rvs) < 1.0
-    assert isinstance(rvs, Array)
+    assert isinstance(rvs, array_api_strict._array_object.Array)
 
 
 @pytest.mark.skipif(not HAVE_ARRAY_API_STRICT, reason="test requires array_api_strict")
 def test_normal() -> None:
-    from array_api_strict._array_object import Array
+    import array_api_strict
 
-    rng = _utils.Generator(SEED)
+    rng = _utils.rng_dispatcher(xp=array_api_strict)
     rvs = rng.normal(1, 2, size=10_000)
     assert rvs.shape == (10_000,)
-    assert isinstance(rvs, Array)
+    assert isinstance(rvs, array_api_strict._array_object.Array)
 
 
 @pytest.mark.skipif(not HAVE_ARRAY_API_STRICT, reason="test requires array_api_strict")
 def test_standard_normal() -> None:
-    from array_api_strict._array_object import Array
+    import array_api_strict
 
-    rng = _utils.Generator(SEED)
+    rng = _utils.rng_dispatcher(xp=array_api_strict)
     rvs = rng.standard_normal(size=10_000)
     assert rvs.shape == (10_000,)
-    assert isinstance(rvs, Array)
+    assert isinstance(rvs, array_api_strict._array_object.Array)
 
 
 @pytest.mark.skipif(not HAVE_ARRAY_API_STRICT, reason="test requires array_api_strict")
 def test_poisson() -> None:
-    from array_api_strict._array_object import Array
+    import array_api_strict
 
-    rng = _utils.Generator(SEED)
+    rng = _utils.rng_dispatcher(xp=array_api_strict)
     rvs = rng.poisson(lam=1, size=10_000)
     assert rvs.shape == (10_000,)
-    assert isinstance(rvs, Array)
+    assert isinstance(rvs, array_api_strict._array_object.Array)
 
 
 @pytest.mark.skipif(not HAVE_ARRAY_API_STRICT, reason="test requires array_api_strict")
 def test_uniform() -> None:
     import array_api_strict
-    from array_api_strict._array_object import Array
 
-    rng = _utils.Generator(SEED)
+    rng = _utils.rng_dispatcher(xp=array_api_strict)
     rvs = rng.uniform(size=10_000)
     assert rvs.shape == (10_000,)
     assert array_api_strict.min(rvs) >= 0.0
     assert array_api_strict.max(rvs) < 1.0
-    assert isinstance(rvs, Array)
+    assert isinstance(rvs, array_api_strict._array_object.Array)
diff --git a/tests/core/test_fields.py b/tests/core/test_fields.py
@@ -8,6 +8,7 @@
 import pytest
 
 import glass
+import glass._array_api_utils as _utils
 import glass.fields
 
 if TYPE_CHECKING:
@@ -386,20 +387,19 @@ def test_generate_grf(compare: type[Compare]) -> None:
     gls = [np.asarray([1.0, 0.5, 0.1])]
     nside = 4
     ncorr = 1
-    seed = 42
 
     gaussian_fields = list(glass.fields._generate_grf(gls, nside))
 
     assert gaussian_fields[0].shape == (hp.nside2npix(nside),)
 
     # requires resetting the RNG for reproducibility
-    rng = np.random.default_rng(seed=seed)
+    rng = _utils.rng_dispatcher(xp=np)
     gaussian_fields = list(glass.fields._generate_grf(gls, nside, rng=rng))
 
     assert gaussian_fields[0].shape == (hp.nside2npix(nside),)
 
     # requires resetting the RNG for reproducibility
-    rng = np.random.default_rng(seed=seed)
+    rng = _utils.rng_dispatcher(xp=np)
     new_gaussian_fields = list(
         glass.fields._generate_grf(gls, nside, ncorr=ncorr, rng=rng),
     )
diff --git a/tests/core/test_jax.py b/tests/core/test_jax.py
@@ -1,17 +1,17 @@
+import jax.numpy as jnp
 import pytest
+from jax.typing import ArrayLike
 
-jax = pytest.importorskip("jax", reason="tests require jax")
-
-import jax.numpy as jnp  # noqa: E402
-from jax.typing import ArrayLike  # noqa: E402
+import glass._array_api_utils as _utils
+from glass.jax import Generator
 
-from glass.jax import Generator  # noqa: E402
+jax = pytest.importorskip("jax", reason="tests require jax")
 
 SEED = 42
 
 
 def test_init() -> None:
-    rng = Generator(SEED)
+    rng = _utils.rng_dispatcher(xp=jnp)
     assert isinstance(rng, Generator)
     assert isinstance(rng.key, jax.Array)
     assert jax.dtypes.issubdtype(rng.key.dtype, jax.dtypes.prng_key)
@@ -31,18 +31,18 @@ def test_from_key() -> None:
 
 
 def test_key() -> None:
-    rng = Generator(SEED)
-    rngkey, outkey = jax.random.split(rng.key, 2)
-    key = rng.split()
-    assert jnp.all(rng.key == rngkey)
+    rng = _utils.rng_dispatcher(xp=jnp)
+    rngkey, outkey = jax.random.split(rng.key, 2)  # type: ignore[union-attr]
+    key = rng.split()  # type: ignore[union-attr]
+    assert jnp.all(rng.key == rngkey)  # type: ignore[union-attr]
     assert jnp.all(key == outkey)
 
 
 def test_spawn() -> None:
-    rng = Generator(SEED)
-    key, *subkeys = jax.random.split(rng.key, 4)
-    subrngs = rng.spawn(3)
-    assert rng.key == key
+    rng = _utils.rng_dispatcher(xp=jnp)
+    key, *subkeys = jax.random.split(rng.key, 4)  # type: ignore[union-attr]
+    subrngs = rng.spawn(3)  # type: ignore[union-attr]
+    assert rng.key == key  # type: ignore[union-attr]
     assert isinstance(subrngs, list)
     assert len(subrngs) == 3
     for subrng, subkey in zip(subrngs, subkeys, strict=False):
@@ -51,48 +51,48 @@ def test_spawn() -> None:
 
 
 def test_random() -> None:
-    rng = Generator(SEED)
-    key = rng.key
+    rng = _utils.rng_dispatcher(xp=jnp)
+    key = rng.key  # type: ignore[union-attr]
     rvs = rng.random(size=10_000)
-    assert rng.key != key
+    assert rng.key != key  # type: ignore[union-attr]
     assert rvs.shape == (10_000,)
     assert jnp.min(rvs) >= 0.0
     assert jnp.max(rvs) < 1.0
     assert isinstance(rvs, ArrayLike)
 
 
 def test_normal() -> None:
-    rng = Generator(SEED)
-    key = rng.key
+    rng = _utils.rng_dispatcher(xp=jnp)
+    key = rng.key  # type: ignore[union-attr]
     rvs = rng.normal(1, 2, size=10_000)
-    assert rng.key != key
+    assert rng.key != key  # type: ignore[union-attr]
     assert rvs.shape == (10_000,)
     assert isinstance(rvs, ArrayLike)
 
 
 def test_standard_normal() -> None:
-    rng = Generator(SEED)
-    key = rng.key
+    rng = _utils.rng_dispatcher(xp=jnp)
+    key = rng.key  # type: ignore[union-attr]
     rvs = rng.standard_normal(size=10_000)
-    assert rng.key != key
+    assert rng.key != key  # type: ignore[union-attr]
     assert rvs.shape == (10_000,)
     assert isinstance(rvs, ArrayLike)
 
 
 def test_poisson() -> None:
-    rng = Generator(SEED)
-    key = rng.key
+    rng = _utils.rng_dispatcher(xp=jnp)
+    key = rng.key  # type: ignore[union-attr]
     rvs = rng.poisson(lam=1, size=10_000)
-    assert rng.key != key
+    assert rng.key != key  # type: ignore[union-attr]
     assert rvs.shape == (10_000,)
     assert isinstance(rvs, ArrayLike)
 
 
 def test_uniform() -> None:
-    rng = Generator(SEED)
-    key = rng.key
+    rng = _utils.rng_dispatcher(xp=jnp)
+    key = rng.key  # type: ignore[union-attr]
     rvs = rng.uniform(size=10_000)
-    assert rng.key != key
+    assert rng.key != key  # type: ignore[union-attr]
     assert rvs.shape == (10_000,)
     assert jnp.min(rvs) >= 0.0
     assert jnp.max(rvs) < 1.0
diff --git a/tests/fixtures/generators.py b/tests/fixtures/generators.py

Original file line number	Diff line number	Diff line change
`@@ -175,7 +175,7 @@ def function_to_benchmark() -> list[Any]:`
`175`	`175`	`generator = glass.fields._generate_grf(`
`176`	`176`	`gls,`
`177`	`177`	`nside,`
`178`		`- rng=urngb if use_rng else None, # type: ignore[arg-type]`
	`178`	`+ rng=urngb if use_rng else None,`
`179`	`179`	`ncorr=ncorr,`
`180`	`180`	`)`
`181`	`181`	`return generator_consumer.consume(generator) # type: ignore[no-any-return]`