use RANDOM_STATE_LIKE in simulators (#2531)

quantumlib · Nov 13, 2019 · d7927cb · d7927cb
1 parent 15c142d
commit d7927cb
Show file tree

Hide file tree

Showing 7 changed files with 41 additions and 89 deletions.
diff --git a/cirq/sim/clifford/clifford_simulator.py b/cirq/sim/clifford/clifford_simulator.py
@@ -29,15 +29,15 @@
     to wavefunction amplitudes.
 """
 
-from typing import Dict, List, Iterator, Union, Optional, Sequence
+from typing import Dict, List, Iterator, Sequence
 import collections
 import numpy as np
 import cirq
 from cirq.sim import simulator
 from cirq.sim.clifford import clifford_tableau, stabilizer_state_ch_form
 from cirq.ops import raw_types
 from cirq.ops.dense_pauli_string import DensePauliString
-from cirq import circuits, study, ops, protocols
+from cirq import circuits, study, ops, protocols, value
 
 
 class CliffordSimulator(simulator.SimulatesSamples,
@@ -204,8 +204,7 @@ def _simulator_state(self):
     def sample(self,
                qubits: List[ops.Qid],
                repetitions: int = 1,
-               seed: Optional[Union[int, np.random.RandomState]] = None
-              ) -> np.ndarray:
+               seed: value.RANDOM_STATE_LIKE = None) -> np.ndarray:
 
         measurements = []
 

diff --git a/cirq/sim/density_matrix_simulator.py b/cirq/sim/density_matrix_simulator.py
@@ -16,7 +16,7 @@
 
 import collections
 
-from typing import Dict, Iterator, List, Optional, Type, Union, TYPE_CHECKING
+from typing import Dict, Iterator, List, Type, Union, TYPE_CHECKING
 
 import numpy as np
 
@@ -120,7 +120,7 @@ def __init__(self,
                  *,
                  dtype: Type[np.number] = np.complex64,
                  noise: 'cirq.NOISE_MODEL_LIKE' = None,
-                 seed: Optional[Union[int, np.random.RandomState]] = None):
+                 seed: value.RANDOM_STATE_LIKE = None):
         """Density matrix simulator.
 
          Args:
@@ -134,15 +134,9 @@ def __init__(self,
                 'dtype must be complex64 or complex128, was {}'.format(dtype))
 
         self._dtype = dtype
+        self._prng = value.parse_random_state(seed)
         self.noise = devices.NoiseModel.from_noise_model_like(noise)
 
-        if seed is None:
-            self.prng = None
-        elif isinstance(seed, np.random.RandomState):
-            self.prng = seed
-        else:
-            self.prng = np.random.RandomState(seed)
-
     def _run(self, circuit: circuits.Circuit,
              param_resolver: study.ParamResolver,
              repetitions: int) -> Dict[str, np.ndarray]:
@@ -170,7 +164,7 @@ def _run_sweep_sample(self,
                                ops.MeasurementGate)]
         return step_result.sample_measurement_ops(measurement_ops,
                                                   repetitions,
-                                                  seed=self.prng)
+                                                  seed=self._prng)
 
     def _run_sweep_repeat(self,
                           circuit: circuits.Circuit,
@@ -285,7 +279,7 @@ def keep(potential_op: ops.Operation) -> bool:
                             indices,
                             qid_shape=qid_shape,
                             out=state.tensor,
-                            seed=self.prng)
+                            seed=self._prng)
                         corrected = [
                             bit ^ (bit < 2 and mask)
                             for bit, mask in zip(bits, invert_mask)
@@ -418,8 +412,7 @@ def density_matrix(self):
     def sample(self,
                qubits: List[ops.Qid],
                repetitions: int = 1,
-               seed: Optional[Union[int, np.random.RandomState]] = None
-              ) -> np.ndarray:
+               seed: value.RANDOM_STATE_LIKE = None) -> np.ndarray:
         indices = [self._qubit_map[q] for q in qubits]
         return density_matrix_utils.sample_density_matrix(
             self._simulator_state().density_matrix,

diff --git a/cirq/sim/density_matrix_utils.py b/cirq/sim/density_matrix_utils.py
@@ -91,7 +91,7 @@ def sample_density_matrix(
         *,  # Force keyword arguments
         qid_shape: Optional[Tuple[int, ...]] = None,
         repetitions: int = 1,
-        seed: Optional[Union[int, np.random.RandomState]] = None) -> np.ndarray:
+        seed: value.RANDOM_STATE_LIKE = None) -> np.ndarray:
     """Samples repeatedly from measurements in the computational basis.
 
     Note that this does not modify the density_matrix.
@@ -136,12 +136,7 @@ def sample_density_matrix(
     if repetitions == 0 or len(indices) == 0:
         return np.zeros(shape=(repetitions, len(indices)), dtype=np.int8)
 
-    if seed is None:
-        prng = np.random
-    elif isinstance(seed, np.random.RandomState):
-        prng = seed
-    else:
-        prng = np.random.RandomState(seed)
+    prng = value.parse_random_state(seed)
 
     # Calculate the measurement probabilities.
     probs = _probs(density_matrix, indices, qid_shape)
@@ -158,13 +153,12 @@ def sample_density_matrix(
                     dtype=np.int8)
 
 
-def measure_density_matrix(
-        density_matrix: np.ndarray,
-        indices: List[int],
-        qid_shape: Optional[Tuple[int, ...]] = None,
-        out: np.ndarray = None,
-        seed: Optional[Union[int, np.random.RandomState]] = None
-) -> Tuple[List[int], np.ndarray]:
+def measure_density_matrix(density_matrix: np.ndarray,
+                           indices: List[int],
+                           qid_shape: Optional[Tuple[int, ...]] = None,
+                           out: np.ndarray = None,
+                           seed: value.RANDOM_STATE_LIKE = None
+                          ) -> Tuple[List[int], np.ndarray]:
     """Performs a measurement of the density matrix in the computational basis.
 
     This does not modify `density_matrix` unless the optional `out` is
@@ -218,12 +212,7 @@ def measure_density_matrix(
         return ([], out)
         # Final else: if out is matrix then matrix will be modified in place.
 
-    if seed is None:
-        prng = np.random
-    elif isinstance(seed, np.random.RandomState):
-        prng = seed
-    else:
-        prng = np.random.RandomState(seed)
+    prng = value.parse_random_state(seed)
 
     # Cache initial shape.
     initial_shape = density_matrix.shape

diff --git a/cirq/sim/mux.py b/cirq/sim/mux.py
@@ -21,7 +21,7 @@
 
 import numpy as np
 
-from cirq import circuits, protocols, study, schedules, devices, ops
+from cirq import circuits, protocols, study, schedules, devices, ops, value
 from cirq.sim import sparse_simulator, density_matrix_simulator
 
 if TYPE_CHECKING:
@@ -34,8 +34,7 @@ def sample(program: Union[circuits.Circuit, schedules.Schedule],
            param_resolver: Optional[study.ParamResolver] = None,
            repetitions: int = 1,
            dtype: Type[np.number] = np.complex64,
-           seed: Optional[Union[int, np.random.RandomState]] = None
-          ) -> study.TrialResult:
+           seed: value.RANDOM_STATE_LIKE = None) -> study.TrialResult:
     """Simulates sampling from the given circuit or schedule.
 
     Args:
@@ -73,8 +72,7 @@ def final_wavefunction(
         param_resolver: study.ParamResolverOrSimilarType = None,
         qubit_order: ops.QubitOrderOrList = ops.QubitOrder.DEFAULT,
         dtype: Type[np.number] = np.complex64,
-        seed: Optional[Union[int, np.random.RandomState]] = None
-) -> 'np.ndarray':
+        seed: value.RANDOM_STATE_LIKE = None) -> 'np.ndarray':
     """Returns the state vector resulting from acting operations on a state.
 
     By default the input state is the computational basis zero state, in which
@@ -141,7 +139,7 @@ def sample_sweep(program: Union[circuits.Circuit, schedules.Schedule],
                  noise: 'cirq.NOISE_MODEL_LIKE' = None,
                  repetitions: int = 1,
                  dtype: Type[np.number] = np.complex64,
-                 seed: Optional[Union[int, np.random.RandomState]] = None
+                 seed: value.RANDOM_STATE_LIKE = None
                 ) -> List[study.TrialResult]:
     """Runs the supplied Circuit or Schedule, mimicking quantum hardware.
 
@@ -163,12 +161,7 @@ def sample_sweep(program: Union[circuits.Circuit, schedules.Schedule],
         TrialResult list for this run; one for each possible parameter
         resolver.
     """
-    if seed is None:
-        prng = None
-    elif isinstance(seed, np.random.RandomState):
-        prng = seed
-    else:
-        prng = np.random.RandomState(seed)
+    prng = value.parse_random_state(seed)
 
     circuit = (program.to_circuit()
                if isinstance(program, schedules.Schedule) else program)

diff --git a/cirq/sim/simulator.py b/cirq/sim/simulator.py
@@ -420,8 +420,7 @@ def _simulator_state(self) -> Any:
     def sample(self,
                qubits: List[ops.Qid],
                repetitions: int = 1,
-               seed: Optional[Union[int, np.random.RandomState]] = None
-              ) -> np.ndarray:
+               seed: value.RANDOM_STATE_LIKE = None) -> np.ndarray:
         """Samples from the system at this point in the computation.
 
         Note that this does not collapse the wave function.
@@ -440,12 +439,11 @@ def sample(self,
         """
         raise NotImplementedError()
 
-    def sample_measurement_ops(
-            self,
-            measurement_ops: List[ops.GateOperation],
-            repetitions: int = 1,
-            seed: Optional[Union[int, np.random.RandomState]] = None
-    ) -> Dict[str, np.ndarray]:
+    def sample_measurement_ops(self,
+                               measurement_ops: List[ops.GateOperation],
+                               repetitions: int = 1,
+                               seed: value.RANDOM_STATE_LIKE = None
+                              ) -> Dict[str, np.ndarray]:
         """Samples from the system at this point in the computation.
 
         Note that this does not collapse the wave function.

diff --git a/cirq/sim/sparse_simulator.py b/cirq/sim/sparse_simulator.py
@@ -16,12 +16,11 @@
 
 import collections
 
-from typing import Dict, Iterator, List, Optional, Tuple, Type, Union, \
-    TYPE_CHECKING
+from typing import Dict, Iterator, List, Tuple, Type, TYPE_CHECKING
 
 import numpy as np
 
-from cirq import circuits, linalg, ops, protocols, study
+from cirq import circuits, linalg, ops, protocols, study, value
 from cirq.sim import simulator, wave_function, wave_function_simulator
 
 if TYPE_CHECKING:
@@ -139,7 +138,7 @@ class Simulator(simulator.SimulatesSamples,
     def __init__(self,
                  *,
                  dtype: Type[np.number] = np.complex64,
-                 seed: Optional[Union[int, np.random.RandomState]] = None):
+                 seed: value.RANDOM_STATE_LIKE = None):
         """A sparse matrix simulator.
 
         Args:
@@ -151,13 +150,7 @@ def __init__(self,
             raise ValueError(
                 'dtype must be a complex type but was {}'.format(dtype))
         self._dtype = dtype
-
-        if seed is None:
-            self.prng = None
-        elif isinstance(seed, np.random.RandomState):
-            self.prng = seed
-        else:
-            self.prng = np.random.RandomState(seed)
+        self._prng = value.parse_random_state(seed)
 
     def _run(
         self,
@@ -192,7 +185,7 @@ def _run_sweep_sample(
                                    ops.MeasurementGate)]
         return step_result.sample_measurement_ops(measurement_ops,
                                                   repetitions,
-                                                  seed=self.prng)
+                                                  seed=self._prng)
 
     def _run_sweep_repeat(
         self,
@@ -344,7 +337,7 @@ def _simulate_measurement(self, op: ops.Operation, data: _StateAndBuffer,
                 indices,
                 out=data.state,
                 qid_shape=data.state.shape,
-                seed=self.prng)
+                seed=self._prng)
             corrected = [
                 bit ^ (bit < 2 and mask)
                 for bit, mask in zip(bits, invert_mask)
@@ -359,8 +352,7 @@ def _simulate_mixture(self, op: ops.Operation, data: _StateAndBuffer,
         # We work around numpy barfing on choosing from a list of
         # numpy arrays (which is not `one-dimensional`) by selecting
         # the index of the unitary.
-        prng = self.prng or np.random
-        index = prng.choice(range(len(unitaries)), p=probs)
+        index = self._prng.choice(range(len(unitaries)), p=probs)
         shape = protocols.qid_shape(op) * 2
         unitary = unitaries[index].astype(self._dtype).reshape(shape)
         result = linalg.targeted_left_multiply(unitary, data.state, indices,
@@ -445,8 +437,7 @@ def set_state_vector(self, state: 'cirq.STATE_VECTOR_LIKE'):
     def sample(self,
                qubits: List[ops.Qid],
                repetitions: int = 1,
-               seed: Optional[Union[int, np.random.RandomState]] = None
-              ) -> np.ndarray:
+               seed: value.RANDOM_STATE_LIKE = None) -> np.ndarray:
         indices = [self.qubit_map[qubit] for qubit in qubits]
         return wave_function.sample_state_vector(self._state_vector,
                                                  indices,

diff --git a/cirq/sim/wave_function.py b/cirq/sim/wave_function.py
@@ -527,7 +527,7 @@ def sample_state_vector(
         *,  # Force keyword args
         qid_shape: Optional[Tuple[int, ...]] = None,
         repetitions: int = 1,
-        seed: Optional[Union[int, np.random.RandomState]] = None) -> np.ndarray:
+        seed: value.RANDOM_STATE_LIKE = None) -> np.ndarray:
     """Samples repeatedly from measurements in the computational basis.
 
     Note that this does not modify the passed in state.
@@ -567,12 +567,7 @@ def sample_state_vector(
     if repetitions == 0 or len(indices) == 0:
         return np.zeros(shape=(repetitions, len(indices)), dtype=np.uint8)
 
-    if seed is None:
-        prng = np.random
-    elif isinstance(seed, np.random.RandomState):
-        prng = seed
-    else:
-        prng = np.random.RandomState(seed)
+    prng = value.parse_random_state(seed)
 
     # Calculate the measurement probabilities.
     probs = _probs(state, indices, qid_shape)
@@ -596,8 +591,7 @@ def measure_state_vector(
         *,  # Force keyword args
         qid_shape: Optional[Tuple[int, ...]] = None,
         out: np.ndarray = None,
-        seed: Optional[Union[int, np.random.RandomState]] = None
-) -> Tuple[List[int], np.ndarray]:
+        seed: value.RANDOM_STATE_LIKE = None) -> Tuple[List[int], np.ndarray]:
     """Performs a measurement of the state in the computational basis.
 
     This does not modify `state` unless the optional `out` is `state`.
@@ -643,12 +637,7 @@ def measure_state_vector(
         # Final else: if out is state then state will be modified in place.
         return ([], out)
 
-    if seed is None:
-        prng = np.random
-    elif isinstance(seed, np.random.RandomState):
-        prng = seed
-    else:
-        prng = np.random.RandomState(seed)
+    prng = value.parse_random_state(seed)
 
     # Cache initial shape.
     initial_shape = state.shape