Refactor simulator RNG handling

Author: jp

cirq-core/cirq/sim/simulator_base.py

@@ -29,6 +29,7 @@
     Type,
     TypeVar,
     TYPE_CHECKING,
+    Union,
 )
 
 import numpy as np
@@ -93,21 +94,27 @@ def __init__(
         *,
         dtype: Type[np.complexfloating] = np.complex64,
         noise: 'cirq.NOISE_MODEL_LIKE' = None,
-        seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+        seed: Optional[Union[int, np.random.Generator, np.random.RandomState]] = None,
         split_untangled_states: bool = False,
     ):
         """Initializes the simulator.
 
         Args:
             dtype: The `numpy.dtype` used by the simulation.
             noise: A noise model to apply while simulating.
-            seed: The random seed to use for this simulator.
+            seed: The random seed or generator to use for this simulator.
             split_untangled_states: If True, optimizes simulation by running
                 unentangled qubit sets independently and merging those states
                 at the end.
         """
         self._dtype = dtype
-        self._prng = value.parse_random_state(seed)
+        if isinstance(seed, np.random.RandomState):
+            # Convert RandomState to Generator for backward compatibility
+            self._prng = np.random.default_rng(seed.get_state()[1][0])
+        elif isinstance(seed, np.random.Generator):
+            self._prng = seed
+        else:
+            self._prng = np.random.default_rng(seed)
         self._noise = devices.NoiseModel.from_noise_model_like(noise)
         self._split_untangled_states = split_untangled_states
 
@@ -228,6 +235,7 @@ def _run(
         circuit: 'cirq.AbstractCircuit',
         param_resolver: 'cirq.ParamResolver',
         repetitions: int,
+        rng: Optional[np.random.Generator] = None,
     ) -> Dict[str, np.ndarray]:
         """See definition in `cirq.SimulatesSamples`."""
         param_resolver = param_resolver or study.ParamResolver({})
@@ -254,7 +262,10 @@ def _run(
             assert step_result is not None
             measurement_ops = [cast(ops.GateOperation, op) for op in general_ops]
             return step_result.sample_measurement_ops(
-                measurement_ops, repetitions, seed=self._prng, _allow_repeated=True
+                measurement_ops,
+                repetitions,
+                seed=rng if rng is not None else self._prng,
+                _allow_repeated=True,
             )
 
         records: Dict['cirq.MeasurementKey', List[Sequence[Sequence[int]]]] = {}
@@ -395,9 +406,15 @@ def sample(
         self,
         qubits: List['cirq.Qid'],
         repetitions: int = 1,
-        seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+        seed: Optional[Union[int, np.random.Generator, np.random.RandomState]] = None,
     ) -> np.ndarray:
-        return self._sim_state.sample(qubits, repetitions, seed)
+        if isinstance(seed, np.random.RandomState):
+            rng = np.random.default_rng(seed.get_state()[1][0])
+        elif isinstance(seed, np.random.Generator):
+            rng = seed
+        else:
+            rng = np.random.default_rng(seed)
+        return self._sim_state.sample(qubits, repetitions, rng)
 
 
 class SimulationTrialResultBase(

cirq-core/cirq/sim/simulator_base_test.py

@@ -434,3 +434,29 @@ def test_inhomogeneous_measurement_count_padding():
     results = sim.run(c, repetitions=10)
     for i in range(10):
         assert np.sum(results.records['m'][i, :, :]) == 1
+
+
+def test_run_with_custom_rng():
+    sim = cirq.Simulator()
+    circuit = cirq.Circuit(cirq.H(q0), cirq.measure(q0))
+    rng1 = np.random.default_rng(seed=1234)
+    rng2 = np.random.default_rng(seed=1234)
+
+    result1 = sim._run(circuit, param_resolver=cirq.ParamResolver({}), repetitions=10, rng=rng1)
+    result2 = sim._run(circuit, param_resolver=cirq.ParamResolver({}), repetitions=10, rng=rng2)
+    assert np.array_equal(result1['q(0)'], result2['q(0)'])
+
+    rng3 = np.random.default_rng(seed=5678)
+    result3 = sim._run(circuit, param_resolver=cirq.ParamResolver({}), repetitions=10, rng=rng3)
+    assert not np.array_equal(result1['q(0)'], result3['q(0)'])
+
+
+def test_run_with_explicit_rng_override():
+    sim1 = cirq.Simulator(seed=1234)
+    sim2 = cirq.Simulator(seed=5678)
+    circuit = cirq.Circuit(cirq.H(q0), cirq.measure(q0))
+    rng = np.random.default_rng(1234)
+
+    result1 = sim1._run(circuit, cirq.ParamResolver({}), repetitions=10)
+    result2 = sim2._run(circuit, cirq.ParamResolver({}), repetitions=10, rng=rng)
+    assert np.array_equal(result1['q(0)'], result2['q(0)'])