Merge pull request #23195 from seberg/public-rng-spawn

API: Add `rng.spawn()`, `bit_gen.spawn()`, and `bit_gen.seed_seq`

9 files changed, 213 insertions, 8 deletions

diff --git a/doc/release/upcoming_changes/23195.improvement.rst b/doc/release/upcoming_changes/23195.improvement.rst
new file mode 100644
index 000000000..38b33e849
--- /dev/null
+++ b/doc/release/upcoming_changes/23195.improvement.rst
@@ -0,0 +1,20 @@
+Ability to directly spawn random number generators
+--------------------------------------------------
+`numpy.random.Generator.spawn` now allows to directly spawn new
+independent child generators via the `numpy.random.SeedSequence.spawn`
+mechanism.
+`numpy.random.BitGenerator.spawn` does the same for the underlying
+bit generator.
+
+Additionally, `numpy.random.BitGenerator.seed_seq` now gives direct
+access to the seed sequence used for initializing the bit generator.
+This allows for example::
+
+    seed = 0x2e09b90939db40c400f8f22dae617151
+    rng = np.random.default_rng(seed)
+    child_rng1, child_rng2 = rng.spawn(2)
+
+    # safely use rng, child_rng1, and child_rng2
+
+Previously, this was hard to do without passing the ``SeedSequence``
+explicitly.  Please see `numpy.random.SeedSequence` for more information.
diff --git a/doc/source/reference/random/bit_generators/index.rst b/doc/source/reference/random/bit_generators/index.rst
index d93f38d0b..14c19a6bd 100644
--- a/doc/source/reference/random/bit_generators/index.rst
+++ b/doc/source/reference/random/bit_generators/index.rst
@@ -50,6 +50,8 @@ The included BitGenerators are:
     Philox <philox>
     SFC64 <sfc64>
 
+.. _seeding_and_entropy:
+
 Seeding and Entropy
 ===================
 
@@ -127,6 +129,16 @@ of 12 instances:
 
 .. end_block
 
+If you already have an initial random generator instance, you can shorten
+the above by using the `~BitGenerator.spawn` method:
+
+.. code-block:: python
+
+    from numpy.random import PCG64, SeedSequence
+    # High quality initial entropy
+    entropy = 0x87351080e25cb0fad77a44a3be03b491
+    base_bitgen = PCG64(entropy)
+    generators = base_bitgen.spawn(12)
 
 An alternative way is to use the fact that a `~SeedSequence` can be initialized
 by a tuple of elements. Here we use a base entropy value and an integer
diff --git a/doc/source/reference/random/generator.rst b/doc/source/reference/random/generator.rst
index dc71cb1f9..e08395b17 100644
--- a/doc/source/reference/random/generator.rst
+++ b/doc/source/reference/random/generator.rst
@@ -18,12 +18,13 @@ can be changed by passing an instantized BitGenerator to ``Generator``.
     :members: __init__
     :exclude-members: __init__
 
-Accessing the BitGenerator
---------------------------
+Accessing the BitGenerator and Spawning
+---------------------------------------
 .. autosummary::
    :toctree: generated/
 
    ~numpy.random.Generator.bit_generator
+   ~numpy.random.Generator.spawn
 
 Simple random data
 ------------------
diff --git a/doc/source/reference/random/parallel.rst b/doc/source/reference/random/parallel.rst
index b625d34b7..b4934a0ca 100644
--- a/doc/source/reference/random/parallel.rst
+++ b/doc/source/reference/random/parallel.rst
@@ -12,6 +12,11 @@ or distributed).
 `~SeedSequence` spawning
 ------------------------
 
+NumPy allows you to spawn new (with very high probability) independent
+`~BitGenerator` and `~Generator` instances via their ``spawn()`` method.
+This spawning is implemented by the `~SeedSequence` used for initializing
+the bit generators random stream.
+
 `~SeedSequence` `implements an algorithm`_ to process a user-provided seed,
 typically as an integer of some size, and to convert it into an initial state for
 a `~BitGenerator`. It uses hashing techniques to ensure that low-quality seeds
@@ -53,15 +58,25 @@ wrap this together into an API that is easy to use and difficult to misuse.
 
 .. end_block
 
-Child `~SeedSequence` objects can also spawn to make grandchildren, and so on.
-Each `~SeedSequence` has its position in the tree of spawned `~SeedSequence`
-objects mixed in with the user-provided seed to generate independent (with very
-high probability) streams.
+For convenience the direct use of `~SeedSequence` is not necessary.
+The above ``streams`` can be spawned directly from a parent generator
+via `~Generator.spawn`:
+
+.. code-block:: python
+
+  parent_rng = default_rng(12345)
+  streams = parent_rng.spawn(10)
+
+.. end_block
+
+Child objects can also spawn to make grandchildren, and so on.
+Each child has a `~SeedSequence` with its position in the tree of spawned
+child objects mixed in with the user-provided seed to generate independent
+(with very high probability) streams.
 
 .. code-block:: python
 
-  grandchildren = child_seeds[0].spawn(4)
-  grand_streams = [default_rng(s) for s in grandchildren]
+  grandchildren = streams[0].spawn(4)
 
 .. end_block
 
diff --git a/numpy/random/_generator.pyi b/numpy/random/_generator.pyi
index f0d814fef..23c04e472 100644
--- a/numpy/random/_generator.pyi
+++ b/numpy/random/_generator.pyi
@@ -72,6 +72,7 @@ class Generator:
     def __reduce__(self) -> tuple[Callable[[str], Generator], tuple[str], dict[str, Any]]: ...
     @property
     def bit_generator(self) -> BitGenerator: ...
+    def spawn(self, n_children: int) -> list[Generator]: ...
     def bytes(self, length: int) -> bytes: ...
     @overload
     def standard_normal(  # type: ignore[misc]
diff --git a/numpy/random/_generator.pyx b/numpy/random/_generator.pyx
index 83a4b2ad5..faf19eaf2 100644
--- a/numpy/random/_generator.pyx
+++ b/numpy/random/_generator.pyx
@@ -238,6 +238,58 @@ cdef class Generator:
         """
         return self._bit_generator
 
+    def spawn(self, int n_children):
+        """
+        Create new independent child generators.
+
+        See :ref:`seedsequence-spawn` for additional notes on spawning
+        children.
+
+        .. versionadded:: 1.25.0
+
+        Returns
+        -------
+        child_generators : list of Generators
+
+        Raises
+        ------
+        TypeError
+            When the underlying SeedSequence does not implement spawning.
+
+        See Also
+        --------
+        random.BitGenerator.spawn, random.SeedSequence.spawn :
+            Equivalent method on the bit generator and seed sequence.
+        bit_generator :
+            The bit generator instance used by the generator.
+
+        Examples
+        --------
+        Starting from a seeded default generator:
+
+        >>> # High quality entropy created with: f"0x{secrets.randbits(128):x}"
+        >>> entropy = 0x3034c61a9ae04ff8cb62ab8ec2c4b501
+        >>> rng = np.random.default_rng(entropy)
+
+        Create two new generators for example for parallel executation:
+
+        >>> child_rng1, child_rng2 = rng.spawn(2)
+
+        Drawn numbers from each are independent but derived from the initial
+        seeding entropy:
+
+        >>> rng.uniform(), child_rng1.uniform(), child_rng2.uniform()
+        (0.19029263503854454, 0.9475673279178444, 0.4702687338396767)
+
+        It is safe to spawn additional children from the original ``rng`` or
+        the children:
+
+        >>> more_child_rngs = rng.spawn(20)
+        >>> nested_spawn = child_rng1.spawn(20)
+
+        """
+        return [type(self)(g) for g in self._bit_generator.spawn(n_children)]
+
     def random(self, size=None, dtype=np.float64, out=None):
         """
         random(size=None, dtype=np.float64, out=None)
@@ -4825,6 +4877,8 @@ def default_rng(seed=None):
     -----
     If ``seed`` is not a `BitGenerator` or a `Generator`, a new `BitGenerator`
     is instantiated. This function does not manage a default global instance.
+
+    See :ref:`seeding_and_entropy` for more information about seeding.
     
     Examples
     --------
diff --git a/numpy/random/bit_generator.pyi b/numpy/random/bit_generator.pyi
index e6e3b10cd..8b9779cad 100644
--- a/numpy/random/bit_generator.pyi
+++ b/numpy/random/bit_generator.pyi
@@ -96,6 +96,9 @@ class BitGenerator(abc.ABC):
     def state(self) -> Mapping[str, Any]: ...
     @state.setter
     def state(self, value: Mapping[str, Any]) -> None: ...
+    @property
+    def seed_seq(self) -> ISeedSequence: ...
+    def spawn(self, n_children: int) -> list[BitGenerator]: ...
     @overload
     def random_raw(self, size: None = ..., output: Literal[True] = ...) -> int: ...  # type: ignore[misc]
     @overload
diff --git a/numpy/random/bit_generator.pyx b/numpy/random/bit_generator.pyx
index 47804c487..06f8c9753 100644
--- a/numpy/random/bit_generator.pyx
+++ b/numpy/random/bit_generator.pyx
@@ -212,6 +212,9 @@ class ISpawnableSeedSequence(ISeedSequence):
         Spawn a number of child `SeedSequence` s by extending the
         `spawn_key`.
 
+        See :ref:`seedsequence-spawn` for additional notes on spawning
+        children.
+
         Parameters
         ----------
         n_children : int
@@ -451,6 +454,9 @@ cdef class SeedSequence():
         Spawn a number of child `SeedSequence` s by extending the
         `spawn_key`.
 
+        See :ref:`seedsequence-spawn` for additional notes on spawning
+        children.
+
         Parameters
         ----------
         n_children : int
@@ -458,6 +464,12 @@ cdef class SeedSequence():
         Returns
         -------
         seqs : list of `SeedSequence` s
+
+        See Also
+        --------
+        random.Generator.spawn, random.BitGenerator.spawn :
+            Equivalent method on the generator and bit generator.
+
         """
         cdef uint32_t i
 
@@ -551,6 +563,53 @@ cdef class BitGenerator():
     def state(self, value):
         raise NotImplementedError('Not implemented in base BitGenerator')
 
+    @property
+    def seed_seq(self):
+        """
+        Get the seed sequence used to initialize the bit generator.
+
+        .. versionadded:: 1.25.0
+
+        Returns
+        -------
+        seed_seq : ISeedSequence
+            The SeedSequence object used to initialize the BitGenerator.
+            This is normally a `np.random.SeedSequence` instance.
+
+        """
+        return self._seed_seq
+
+    def spawn(self, int n_children):
+        """
+        Create new independent child bit generators.
+
+        See :ref:`seedsequence-spawn` for additional notes on spawning
+        children.  Some bit generators also implement ``jumped``
+        as a different approach for creating independent streams.
+
+        .. versionadded:: 1.25.0
+
+        Returns
+        -------
+        child_bit_generators : list of BitGenerators
+
+        Raises
+        ------
+        TypeError
+            When the underlying SeedSequence does not implement spawning.
+
+        See Also
+        --------
+        random.Generator.spawn, random.SeedSequence.spawn :
+            Equivalent method on the generator and seed sequence.
+
+        """
+        if not isinstance(self._seed_seq, ISpawnableSeedSequence):
+            raise TypeError(
+                "The underlying SeedSequence does not implement spawning.")
+
+        return [type(self)(seed=s) for s in self._seed_seq.spawn(n_children)]
+
     def random_raw(self, size=None, output=True):
         """
         random_raw(self, size=None)
diff --git a/numpy/random/tests/test_direct.py b/numpy/random/tests/test_direct.py
index 58d966adf..fa2ae866b 100644
--- a/numpy/random/tests/test_direct.py
+++ b/numpy/random/tests/test_direct.py
@@ -148,6 +148,46 @@ def test_seedsequence():
     assert len(dummy.spawn(10)) == 10
 
 
+def test_generator_spawning():
+    """ Test spawning new generators and bit_generators directly.
+    """
+    rng = np.random.default_rng()
+    seq = rng.bit_generator.seed_seq
+    new_ss = seq.spawn(5)
+    expected_keys = [seq.spawn_key + (i,) for i in range(5)]
+    assert [c.spawn_key for c in new_ss] == expected_keys
+
+    new_bgs = rng.bit_generator.spawn(5)
+    expected_keys = [seq.spawn_key + (i,) for i in range(5, 10)]
+    assert [bg.seed_seq.spawn_key for bg in new_bgs] == expected_keys
+
+    new_rngs = rng.spawn(5)
+    expected_keys = [seq.spawn_key + (i,) for i in range(10, 15)]
+    found_keys = [rng.bit_generator.seed_seq.spawn_key for rng in new_rngs]
+    assert found_keys == expected_keys
+
+    # Sanity check that streams are actually different:
+    assert new_rngs[0].uniform() != new_rngs[1].uniform()
+
+
+def test_non_spawnable():
+    from numpy.random.bit_generator import ISeedSequence
+
+    class FakeSeedSequence:
+        def generate_state(self, n_words, dtype=np.uint32):
+            return np.zeros(n_words, dtype=dtype)
+
+    ISeedSequence.register(FakeSeedSequence)
+
+    rng = np.random.default_rng(FakeSeedSequence())
+
+    with pytest.raises(TypeError, match="The underlying SeedSequence"):
+        rng.spawn(5)
+
+    with pytest.raises(TypeError, match="The underlying SeedSequence"):
+        rng.bit_generator.spawn(5)
+
+
 class Base:
     dtype = np.uint64
     data2 = data1 = {}