Add support for np.random.Generator

cirq-core/cirq/__init__.py

@@ -536,6 +536,7 @@
     MeasurementKey,
     MeasurementType,
     PeriodicValue,
+    PRNG_OR_SEED_LIKE,
     RANDOM_STATE_OR_SEED_LIKE,
     state_vector_to_probabilities,
     SympyCondition,

cirq-core/cirq/protocols/json_test_data/spec.py

@@ -154,6 +154,7 @@
         'QUANTUM_STATE_LIKE',
         'QubitOrderOrList',
         'RANDOM_STATE_OR_SEED_LIKE',
+        'PRNG_OR_SEED_LIKE',
         'STATE_VECTOR_LIKE',
         'Sweepable',
         'TParamKey',

cirq-core/cirq/value/__init__.py

@@ -65,3 +65,6 @@
 from cirq.value.type_alias import TParamKey, TParamVal, TParamValComplex
 
 from cirq.value.value_equality_attr import value_equality
+
+
+from cirq.value.prng import parse_prng, CustomPRNG, PRNG_OR_SEED_LIKE

cirq-core/cirq/value/prng.py

@@ -0,0 +1,80 @@
+# Copyright 2024 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abc
+from typing import TypeVar, Union, overload
+
+import numpy as np
+
+from cirq._doc import document
+
+
+class CustomPRNG(abc.ABC): ...
+
+
+_CUSTOM_PRNG_T = TypeVar("_CUSTOM_PRNG_T", bound=CustomPRNG)
+_PRNG_T = Union[np.random.Generator, np.random.RandomState, _CUSTOM_PRNG_T]
+_SEED_T = Union[int, None]
+PRNG_OR_SEED_LIKE = Union[None, int, np.random.RandomState, np.random.Generator, _CUSTOM_PRNG_T]
+
+document(
+    PRNG_OR_SEED_LIKE,
+    """A pseudorandom number generator or object that can be converted to one.
+
+    If an integer or None, turns into a `np.random.Generator` seeded with that
+    value.
+
+    If none of the above, it is used unmodified. In this case, it is assumed
+    that the object implements whatever methods are required for the use case
+    at hand. For example, it might be an existing instance of `np.random.Generator`
+    or `np.random.RandomState` or a custom pseudorandom number generator implementation
+    and in that case, it has to inherit `cirq.value.CustomPRNG`.
+    """,
+)
+
+
+@overload
+def parse_prng(prng_or_seed: _SEED_T) -> np.random.Generator: ...
+
+
+@overload
+def parse_prng(prng_or_seed: np.random.Generator) -> np.random.Generator: ...
+
+
+@overload
+def parse_prng(prng_or_seed: np.random.RandomState) -> np.random.RandomState: ...
+
+
+@overload
+def parse_prng(prng_or_seed: _CUSTOM_PRNG_T) -> _CUSTOM_PRNG_T: ...
+
+
+def parse_prng(
+    prng_or_seed: PRNG_OR_SEED_LIKE,
+) -> Union[np.random.Generator, np.random.RandomState, _CUSTOM_PRNG_T]:
+    """Interpret an object as a pseudorandom number generator.
+
+    If `prng_or_seed` is None or an integer, returns `np.random.default_rng(prng_or_seed)`.
+    Otherwise, returns `prng_or_seed` unmodified.
+
+    Args:
+        prng_or_seed: The object to be used as or converted to a pseudorandom
+            number generator.
+
+    Returns:
+        The pseudorandom number generator object.
+    """
+    if prng_or_seed is None or isinstance(prng_or_seed, int):
+        return np.random.default_rng(prng_or_seed)
+    return prng_or_seed

cirq-core/cirq/value/prng_test.py

@@ -0,0 +1,48 @@
+# Copyright 2024 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List, Union
+import numpy as np
+import cirq
+
+
+class TestPrng(cirq.value.CustomPRNG):
+
+    def random(self, size):
+        return tuple(range(size))
+
+
+def _sample(prng):
+    return tuple(prng.random(10))
+
+
+def test_parse_rng() -> None:
+    eq = cirq.testing.EqualsTester()
+
+    # An `np.random.Generator` or a seed.
+    group_inputs: List[Union[int, np.random.Generator]] = [42, np.random.default_rng(42)]
+    group: List[np.random.Generator] = [cirq.value.parse_prng(s) for s in group_inputs]
+    eq.add_equality_group(*[_sample(g) for g in group])
+
+    # A None seed.
+    prng: np.random.Generator = cirq.value.parse_prng(None)
+    eq.add_equality_group(_sample(prng))
+
+    # Custom PRNG.
+    custom_prng: TestPrng = cirq.value.parse_prng(TestPrng())
+    eq.add_equality_group(_sample(custom_prng))
+
+    # RandomState PRNG.
+    random_state: np.random.RandomState = cirq.value.parse_prng(np.random.RandomState(42))
+    eq.add_equality_group(_sample(random_state))

cirq-core/cirq/value/random_state.py

@@ -33,6 +33,8 @@
     at hand. For example, it might be an existing instance of
     `np.random.RandomState` or a custom pseudorandom number generator
     implementation.
+
+    Note: prefer to use cirq.PRNG_OR_SEED_LIKE.
     """,
 )