Added optional prng parameter to sampler and simulator methods

cirq-aqt/cirq_aqt/aqt_sampler.py

@@ -405,7 +405,11 @@ def _send_json(
         return measurements
 
     def run_sweep(
-        self, program: cirq.AbstractCircuit, params: cirq.Sweepable, repetitions: int = 1
+        self,
+        program: cirq.AbstractCircuit,
+        params: cirq.Sweepable,
+        repetitions: int = 1,
+        prng: np.random.Generator | None = None,
     ) -> Sequence[cirq.Result]:
         """Samples from the given Circuit.
 
@@ -417,11 +421,14 @@ def run_sweep(
             Should be generated using AQTSampler.generate_circuit_from_list
             params: Parameters to run with the program.
             repetitions: The number of repetitions to simulate.
+            prng: Not supported for this class as no client-side RNG, must be None.
 
         Returns:
             Result list for this run; one for each possible parameter
             resolver.
         """
+        if prng is not None:
+            raise ValueError("AQTSampler has no client-side RNG.")
         # TODO: Use measurement name from circuit.
         # Github issue: https://github.com/quantumlib/Cirq/issues/2199
         meas_name = 'm'

cirq-aqt/cirq_aqt/aqt_sampler_test.py

@@ -459,3 +459,12 @@ def json(self):
     assert workspaces[0]["resources"][0]["id"] == "rid"
     assert workspaces[0]["resources"][0]["name"] == "Resource"
     assert workspaces[0]["resources"][0]["type"] == "device"
+
+
+def test_aqt_sampler_rejects_prng() -> None:
+    a = cirq.LineQubit(0)
+    circuit = cirq.Circuit(cirq.H(a), cirq.measure(a, key='m'))
+
+    sampler = AQTSampler("default", "test", "testkey")
+    with pytest.raises(ValueError, match='RNG'):
+        sampler.run_sweep(circuit, None, 1, prng=np.random.default_rng(0))

cirq-core/cirq/contrib/custom_simulators/custom_state_simulator.py

@@ -84,6 +84,7 @@ def _create_partial_simulation_state(
         initial_state: Any,
         qubits: Sequence[cirq.Qid],
         classical_data: cirq.ClassicalDataStore,
+        prng: np.random.Generator | None = None,
     ) -> TSimulationState:
         return self.state_type(
             initial_state=initial_state, qubits=qubits, classical_data=classical_data

cirq-core/cirq/contrib/quimb/mps_simulator.py

@@ -90,6 +90,7 @@ def _create_partial_simulation_state(
         initial_state: int | MPSState,
         qubits: Sequence[cirq.Qid],
         classical_data: cirq.ClassicalDataStore,
+        prng: np.random.Generator | None = None,
     ) -> MPSState:
         """Creates MPSState args for simulating the Circuit.
 
@@ -101,6 +102,8 @@ def _create_partial_simulation_state(
                 ordering of the computational basis states.
             classical_data: The shared classical data container for this
                 simulation.
+            prng: An `np.random.Generator` to draw from for this call
+                instead of the internal random state.
 
         Returns:
             MPSState args for simulating the Circuit.
@@ -110,7 +113,7 @@ def _create_partial_simulation_state(
 
         return MPSState(
             qubits=qubits,
-            prng=self._prng,
+            prng=prng if prng is not None else self._prng,
             simulation_options=self.simulation_options,
             grouping=self.grouping,
             initial_state=initial_state,
@@ -382,12 +385,16 @@ def to_numpy(self) -> np.ndarray:
         """An alias for the state vector."""
         return self.state_vector()
 
-    def apply_op(self, op: Any, axes: Sequence[int], prng: np.random.RandomState):
+    def apply_op(
+        self, op: Any, axes: Sequence[int], prng: np.random.RandomState | np.random.Generator
+    ):
         """Applies a unitary operation, mutating the object to represent the new state.
 
         op:
             The operation that mutates the object. Note that currently, only 1-
             and 2- qubit operations are currently supported.
+        prng:
+            An `np.random.Generator` to draw from for this call instead of the internal random state.
         """
 
         old_inds = tuple(map(self.i_str, axes))
@@ -484,7 +491,10 @@ def estimation_stats(self):  # pragma: no cover
         }
 
     def _measure(
-        self, axes: Sequence[int], prng: np.random.RandomState, collapse_state_vector=True
+        self,
+        axes: Sequence[int],
+        prng: np.random.RandomState | np.random.Generator,
+        collapse_state_vector=True,
     ) -> list[int]:
         results: list[int] = []
 
@@ -565,7 +575,7 @@ def __init__(
         self,
         *,
         qubits: Sequence[cirq.Qid],
-        prng: np.random.RandomState,
+        prng: np.random.RandomState | np.random.Generator,
         simulation_options: MPSOptions = MPSOptions(),
         grouping: dict[cirq.Qid, int] | None = None,
         initial_state: int = 0,

cirq-core/cirq/experiments/random_quantum_circuit_generation.py

@@ -173,7 +173,7 @@ def random_rotations_between_two_qubit_circuit(
     q1: cirq.Qid,
     depth: int,
     two_qubit_op_factory: Callable[
-        [cirq.Qid, cirq.Qid, np.random.RandomState], cirq.OP_TREE
+        [cirq.Qid, cirq.Qid, np.random.RandomState | np.random.Generator], cirq.OP_TREE
     ] = lambda a, b, _: ops.CZPowGate()(a, b),
     single_qubit_gates: Sequence[cirq.Gate] = (
         ops.X**0.5,
@@ -339,11 +339,13 @@ def _get_random_combinations(
         returned list can be provided to `sample_2q_xeb_circuits` to efficiently
         sample parallel XEB circuits.
     """
-    rs = value.parse_random_state(random_state)
+    parsed_rs = value.parse_random_state(random_state)
 
     combinations_by_layer = []
     for pairs, layer in pair_gen:
-        combinations = rs.randint(0, n_library_circuits, size=(n_combinations, len(pairs)))
+        combinations = value.get_random_int(
+            parsed_rs, 0, n_library_circuits, size=(n_combinations, len(pairs))
+        )
         combinations_by_layer.append(
             CircuitLibraryCombination(layer=layer, combinations=combinations, pairs=pairs)
         )
@@ -541,7 +543,7 @@ def random_rotations_between_grid_interaction_layers_circuit(
     *,  # forces keyword arguments
     device_graph: nx.Graph | None = None,
     two_qubit_op_factory: Callable[
-        [cirq.GridQubit, cirq.GridQubit, np.random.RandomState], cirq.OP_TREE
+        [cirq.GridQubit, cirq.GridQubit, np.random.RandomState | np.random.Generator], cirq.OP_TREE
     ] = lambda a, b, _: ops.CZPowGate()(a, b),
     pattern: Sequence[GridInteractionLayer] = GRID_STAGGERED_PATTERN,
     single_qubit_gates: Sequence[cirq.Gate] = (
@@ -642,7 +644,7 @@ def __init__(
         self,
         qubits: Sequence[cirq.Qid],
         single_qubit_gates: Sequence[cirq.Gate],
-        prng: np.random.RandomState,
+        prng: np.random.RandomState | np.random.Generator,
     ) -> None:
         self.qubits = qubits
         self.single_qubit_gates = single_qubit_gates
@@ -652,9 +654,13 @@ def new_layer(self, previous_single_qubit_layer: cirq.Moment) -> cirq.Moment:
         def random_gate(qubit: cirq.Qid) -> cirq.Gate:
             excluded_op = previous_single_qubit_layer.operation_at(qubit)
             excluded_gate = excluded_op.gate if excluded_op is not None else None
-            g = self.single_qubit_gates[self.prng.randint(0, len(self.single_qubit_gates))]
+            g = self.single_qubit_gates[
+                value.get_random_int(self.prng, 0, len(self.single_qubit_gates))
+            ]
             while g is excluded_gate:
-                g = self.single_qubit_gates[self.prng.randint(0, len(self.single_qubit_gates))]
+                g = self.single_qubit_gates[
+                    value.get_random_int(self.prng, 0, len(self.single_qubit_gates))
+                ]
             return g
 
         return circuits.Moment(random_gate(q).on(q) for q in self.qubits)
@@ -672,7 +678,9 @@ def new_layer(self, previous_single_qubit_layer: cirq.Moment) -> cirq.Moment:
 
 
 def _single_qubit_gates_arg_to_factory(
-    single_qubit_gates: Sequence[cirq.Gate], qubits: Sequence[cirq.Qid], prng: np.random.RandomState
+    single_qubit_gates: Sequence[cirq.Gate],
+    qubits: Sequence[cirq.Qid],
+    prng: np.random.RandomState | np.random.Generator,
 ) -> _SingleQubitLayerFactory:
     """Parse the `single_qubit_gates` argument for circuit generation functions.
 
@@ -689,10 +697,10 @@ def _single_qubit_gates_arg_to_factory(
 def _two_qubit_layer(
     coupled_qubit_pairs: list[GridQubitPairT],
     two_qubit_op_factory: Callable[
-        [cirq.GridQubit, cirq.GridQubit, np.random.RandomState], cirq.OP_TREE
+        [cirq.GridQubit, cirq.GridQubit, np.random.RandomState | np.random.Generator], cirq.OP_TREE
     ],
     layer: GridInteractionLayer,
-    prng: np.random.RandomState,
+    prng: np.random.RandomState | np.random.Generator,
 ) -> Iterator[cirq.OP_TREE]:
     for a, b in coupled_qubit_pairs:
         if (a, b) in layer or (b, a) in layer:

cirq-core/cirq/experiments/random_quantum_circuit_generation_test.py

@@ -247,7 +247,7 @@ def test_random_combinations_layer_circuit_vs_device() -> None:
 
 
 def _cz_with_adjacent_z_rotations(
-    a: cirq.GridQubit, b: cirq.GridQubit, prng: np.random.RandomState
+    a: cirq.GridQubit, b: cirq.GridQubit, prng: np.random.RandomState | np.random.Generator
 ):
     z_exponents = [prng.uniform(0, 1) for _ in range(4)]
     yield cirq.Z(a) ** z_exponents[0]
@@ -370,7 +370,7 @@ def test_random_rotations_between_grid_interaction_layers(
     qubits: Iterable[cirq.GridQubit],
     depth: int,
     two_qubit_op_factory: Callable[
-        [cirq.GridQubit, cirq.GridQubit, np.random.RandomState], cirq.OP_TREE
+        [cirq.GridQubit, cirq.GridQubit, np.random.RandomState | np.random.Generator], cirq.OP_TREE
     ],
     pattern: Sequence[GridInteractionLayer],
     single_qubit_gates: Sequence[cirq.Gate],

cirq-core/cirq/experiments/single_qubit_readout_calibration_test.py

@@ -46,7 +46,7 @@ def __init__(self, p0: float, p1: float, seed: cirq.RANDOM_STATE_OR_SEED_LIKE =
         self.prng = cirq.value.parse_random_state(seed)
         self.simulator = cirq.Simulator(seed=self.prng, split_untangled_states=False)
 
-    def run_sweep(
+    def run_sweep(  # type: ignore[override]
         self, program: cirq.AbstractCircuit, params: cirq.Sweepable, repetitions: int = 1
     ) -> Sequence[cirq.Result]:
         results = self.simulator.run_sweep(program, params, repetitions)

cirq-core/cirq/linalg/decompositions_test.py

@@ -590,7 +590,7 @@ def _random_two_qubit_unitaries(num_samples: int, random_state: cirq.RANDOM_STAT
 
     prng = value.parse_random_state(random_state)
     # Generate the non-local part by explicit matrix exponentiation.
-    kak_vecs = prng.rand(num_samples, 3) * np.pi
+    kak_vecs = value.get_random_array(prng, (num_samples, 3)) * np.pi
     gens = np.einsum('...a,abc->...bc', kak_vecs, _kak_gens)
     evals, evecs = np.linalg.eigh(gens)
     A = np.einsum('...ab,...b,...cb', evecs, np.exp(1j * evals), evecs.conj())

cirq-core/cirq/qis/clifford_tableau.py

@@ -519,7 +519,7 @@ def destabilizers(self) -> list[cirq.DensePauliString]:
         generators above generate the full Pauli group on n qubits."""
         return [self._row_to_dense_pauli(i) for i in range(self.n)]
 
-    def _measure(self, q, prng: np.random.RandomState) -> int:
+    def _measure(self, q, prng: np.random.RandomState | np.random.Generator) -> int:
         """Performs a projective measurement on the q'th qubit.
 
         Returns: the result (0 or 1) of the measurement.
@@ -554,7 +554,7 @@ def _measure(self, q, prng: np.random.RandomState) -> int:
 
         self.zs[p, q] = True
 
-        self.rs[p] = bool(prng.randint(2))
+        self.rs[p] = bool(random_state.get_random_int(prng, 2))
 
         return int(self.rs[p])
 

cirq-core/cirq/sim/classical_simulator.py

@@ -252,6 +252,7 @@ def _create_partial_simulation_state(
         initial_state: Any,
         qubits: Sequence[cirq.Qid],
         classical_data: cirq.ClassicalDataStore,
+        prng: np.random.Generator | None = None,
     ) -> ClassicalBasisSimState:
         """Creates a partial simulation state for the simulator.
 

cirq-core/cirq/sim/clifford/clifford_simulator.py

@@ -76,6 +76,7 @@ def _create_partial_simulation_state(
         initial_state: int | cirq.StabilizerChFormSimulationState,
         qubits: Sequence[cirq.Qid],
         classical_data: cirq.ClassicalDataStore,
+        prng: np.random.Generator | None = None,
     ) -> cirq.StabilizerChFormSimulationState:
         """Creates the StabilizerChFormSimulationState for a circuit.
 
@@ -97,7 +98,7 @@ def _create_partial_simulation_state(
             return initial_state  # pragma: no cover
 
         return clifford.StabilizerChFormSimulationState(
-            prng=self._prng,
+            prng=prng if prng is not None else self._prng,
             classical_data=classical_data,
             qubits=qubits,
             initial_state=initial_state,
@@ -257,7 +258,7 @@ def apply_measurement(
         self,
         op: cirq.Operation,
         measurements: dict[str, list[int]],
-        prng: np.random.RandomState,
+        prng: np.random.RandomState | np.random.Generator,
         collapse_state_vector=True,
     ) -> None:
         if not isinstance(op.gate, cirq.MeasurementGate):

cirq-core/cirq/sim/clifford/clifford_tableau_simulation_state.py

@@ -35,7 +35,7 @@ class CliffordTableauSimulationState(StabilizerSimulationState[clifford_tableau.
     def __init__(
         self,
         tableau: cirq.CliffordTableau,
-        prng: np.random.RandomState | None = None,
+        prng: np.random.RandomState | np.random.Generator | None = None,
         qubits: Sequence[cirq.Qid] | None = None,
         classical_data: cirq.ClassicalDataStore | None = None,
     ):

cirq-core/cirq/sim/clifford/stabilizer_ch_form_simulation_state.py

@@ -35,7 +35,7 @@ class StabilizerChFormSimulationState(
     def __init__(
         self,
         *,
-        prng: np.random.RandomState | None = None,
+        prng: np.random.RandomState | np.random.Generator | None = None,
         qubits: Sequence[cirq.Qid] | None = None,
         initial_state: int | cirq.StabilizerStateChForm = 0,
         classical_data: cirq.ClassicalDataStore | None = None,

cirq-core/cirq/sim/clifford/stabilizer_sampler.py

@@ -38,16 +38,25 @@ def __init__(self, *, seed: cirq.RANDOM_STATE_OR_SEED_LIKE = None):
         self._prng = value.parse_random_state(seed)
 
     def run_sweep(
-        self, program: cirq.AbstractCircuit, params: cirq.Sweepable, repetitions: int = 1
+        self,
+        program: cirq.AbstractCircuit,
+        params: cirq.Sweepable,
+        repetitions: int = 1,
+        prng: np.random.Generator | None = None,
     ) -> Sequence[cirq.Result]:
         results: list[cirq.Result] = []
         for param_resolver in cirq.to_resolvers(params):
             resolved_circuit = cirq.resolve_parameters(program, param_resolver)
-            measurements = self._run(resolved_circuit, repetitions=repetitions)
+            measurements = self._run(resolved_circuit, repetitions=repetitions, prng=prng)
             results.append(cirq.ResultDict(params=param_resolver, measurements=measurements))
         return results
 
-    def _run(self, circuit: cirq.AbstractCircuit, repetitions: int) -> dict[str, np.ndarray]:
+    def _run(
+        self,
+        circuit: cirq.AbstractCircuit,
+        repetitions: int,
+        prng: np.random.Generator | None = None,
+    ) -> dict[str, np.ndarray]:
 
         measurements: dict[str, list[np.ndarray]] = {
             key: [] for key in protocols.measurement_key_names(circuit)
@@ -56,7 +65,9 @@ def _run(self, circuit: cirq.AbstractCircuit, repetitions: int) -> dict[str, np.
 
         for _ in range(repetitions):
             state = CliffordTableauSimulationState(
-                CliffordTableau(num_qubits=len(qubits)), qubits=list(qubits), prng=self._prng
+                CliffordTableau(num_qubits=len(qubits)),
+                qubits=list(qubits),
+                prng=prng if prng is not None else self._prng,
             )
             for op in circuit.all_operations():
                 protocols.act_on(op, state)

cirq-core/cirq/sim/clifford/stabilizer_simulation_state.py

@@ -45,7 +45,7 @@ def __init__(
         self,
         *,
         state: TStabilizerState,
-        prng: np.random.RandomState | None = None,
+        prng: np.random.RandomState | np.random.Generator | None = None,
         qubits: Sequence[cirq.Qid] | None = None,
         classical_data: cirq.ClassicalDataStore | None = None,
     ):

cirq-core/cirq/sim/clifford/stabilizer_state_ch_form.py

@@ -240,7 +240,7 @@ def to_state_vector(self) -> np.ndarray:
 
         return arr
 
-    def _measure(self, q, prng: np.random.RandomState) -> int:
+    def _measure(self, q, prng: np.random.RandomState | np.random.Generator) -> int:
         """Measures the q'th qubit.
 
         Reference: Section 4.1 "Simulating measurements"
@@ -250,7 +250,7 @@ def _measure(self, q, prng: np.random.RandomState) -> int:
         w = self.s.copy()
         for i, v_i in enumerate(self.v):
             if v_i == 1:
-                w[i] = bool(prng.randint(2))
+                w[i] = bool(random_state.get_random_int(prng, 2))
         x_i = sum(w & self.G[q, :]) % 2
         # Project the state to the above measurement outcome.
         self.project_Z(q, x_i)

cirq-core/cirq/sim/density_matrix_simulation_state.py

@@ -247,7 +247,7 @@ def __init__(
         self,
         *,
         available_buffer: list[np.ndarray] | None = None,
-        prng: np.random.RandomState | None = None,
+        prng: np.random.RandomState | np.random.Generator | None = None,
         qubits: Sequence[cirq.Qid] | None = None,
         initial_state: np.ndarray | cirq.STATE_VECTOR_LIKE = 0,
         dtype: type[np.complexfloating] = np.complex64,

cirq-core/cirq/sim/density_matrix_simulator.py

@@ -155,6 +155,7 @@ def _create_partial_simulation_state(
         initial_state: np.ndarray | cirq.STATE_VECTOR_LIKE | cirq.DensityMatrixSimulationState,
         qubits: Sequence[cirq.Qid],
         classical_data: cirq.ClassicalDataStore,
+        prng: np.random.Generator | None = None,
     ) -> cirq.DensityMatrixSimulationState:
         """Creates the DensityMatrixSimulationState for a circuit.
 
@@ -176,7 +177,7 @@ def _create_partial_simulation_state(
 
         return density_matrix_simulation_state.DensityMatrixSimulationState(
             qubits=qubits,
-            prng=self._prng,
+            prng=prng if prng is not None else self._prng,
             classical_data=classical_data,
             initial_state=initial_state,
             dtype=self._dtype,