Split the simulators' creation of ActOnArgs and iteration

cirq/contrib/quimb/mps_simulator.py

@@ -17,10 +17,10 @@
 https://arxiv.org/abs/2002.07730
 """
 
+import dataclasses
 import math
-from typing import Any, Dict, List, Iterator, Optional, Sequence, Set, TYPE_CHECKING, Iterable
+from typing import Any, Dict, List, Optional, Sequence, Set, TYPE_CHECKING, Iterable, Union
 
-import dataclasses
 import numpy as np
 import quimb.tensor as qtn
 
@@ -54,7 +54,9 @@ class MPSOptions:
 
 class MPSSimulator(
     simulator.SimulatesSamples,
-    simulator.SimulatesIntermediateState['MPSSimulatorStepResult', 'MPSTrialResult', 'MPSState'],
+    simulator.SimulatesIntermediateState[
+        'MPSSimulatorStepResult', 'MPSTrialResult', 'MPSState', 'MPSState'
+    ],
 ):
     """An efficient simulator for MPS circuits."""
 
@@ -82,58 +84,68 @@ def __init__(
         self.simulation_options = simulation_options
         self.grouping = grouping
 
-    def _base_iterator(
-        self, circuit: circuits.Circuit, qubit_order: ops.QubitOrderOrList, initial_state: int
-    ) -> Iterator['MPSSimulatorStepResult']:
-        """Iterator over MPSSimulatorStepResult from Moments of a Circuit
+    def create_act_on_args(
+        self,
+        initial_state: Union[int, 'MPSState'],
+        qubits: Sequence['cirq.Qid'],
+    ) -> 'MPSState':
+        """Creates MPSState args for simulating the Circuit.
 
         Args:
-            circuit: The circuit to simulate.
-            qubit_order: Determines the canonical ordering of the qubits. This
-                is often used in specifying the initial state, i.e. the
-                ordering of the computational basis states.
             initial_state: The initial state for the simulation in the
                 computational basis. Represented as a big endian int.
+            qubits: Determines the canonical ordering of the qubits. This
+                is often used in specifying the initial state, i.e. the
+                ordering of the computational basis states.
 
-        Yields:
-            MPSStepResult from simulating a Moment of the Circuit.
+        Returns:
+            MPSState args for simulating the Circuit.
         """
-        qubits = ops.QubitOrder.as_qubit_order(qubit_order).order_for(circuit.all_qubits())
+        if isinstance(initial_state, MPSState):
+            return initial_state
+
+        return MPSState(
+            qubits=qubits,
+            prng=self.prng,
+            simulation_options=self.simulation_options,
+            grouping=self.grouping,
+            initial_state=initial_state,
+        )
 
-        qubit_map = {q: i for i, q in enumerate(qubits)}
+    def _core_iterator(
+        self,
+        circuit: circuits.Circuit,
+        sim_state: 'MPSState',
+    ):
+        """Iterator over MPSSimulatorStepResult from Moments of a Circuit
+
+        Args:
+            circuit: The circuit to simulate.
+            sim_state: The initial state args for the simulation in the
+                computational basis.
 
+        Yields:
+            MPSStepResult from simulating a Moment of the Circuit.
+        """
         if len(circuit) == 0:
             yield MPSSimulatorStepResult(
-                measurements={},
-                state=MPSState(
-                    qubit_map,
-                    self.prng,
-                    self.simulation_options,
-                    self.grouping,
-                    initial_state=initial_state,
-                ),
+                measurements=sim_state.log_of_measurement_results, state=sim_state
             )
             return
 
-        state = MPSState(
-            qubit_map,
-            self.prng,
-            self.simulation_options,
-            self.grouping,
-            initial_state=initial_state,
-        )
-
         noisy_moments = self.noise.noisy_moments(circuit, sorted(circuit.all_qubits()))
         for op_tree in noisy_moments:
             for op in flatten_to_ops(op_tree):
                 if protocols.is_measurement(op) or protocols.has_mixture(op):
-                    state.axes = tuple(qubit_map[qubit] for qubit in op.qubits)
-                    protocols.act_on(op, state)
+                    sim_state.axes = tuple(sim_state.qubit_map[qubit] for qubit in op.qubits)
+                    protocols.act_on(op, sim_state)
                 else:
                     raise NotImplementedError(f"Unrecognized operation: {op!r}")
 
-            yield MPSSimulatorStepResult(measurements=state.log_of_measurement_results, state=state)
-            state.log_of_measurement_results.clear()
+            yield MPSSimulatorStepResult(
+                measurements=sim_state.log_of_measurement_results, state=sim_state
+            )
+            sim_state.log_of_measurement_results.clear()
 
     def _create_simulator_trial_result(
         self,
@@ -286,7 +298,7 @@ class MPSState(ActOnArgs):
 
     def __init__(
         self,
-        qubit_map: Dict['cirq.Qid', int],
+        qubits: Sequence['cirq.Qid'],
         prng: np.random.RandomState,
         simulation_options: MPSOptions = MPSOptions(),
         grouping: Optional[Dict['cirq.Qid', int]] = None,
@@ -297,7 +309,9 @@ def __init__(
         """Creates and MPSState
 
         Args:
-            qubit_map: A map from Qid to an integer that uniquely identifies it.
+            qubits: Determines the canonical ordering of the qubits. This
+                is often used in specifying the initial state, i.e. the
+                ordering of the computational basis states.
             prng: A random number generator, used to simulate measurements.
             simulation_options: Numerical options for the simulation.
             grouping: How to group qubits together, if None all are individual.
@@ -307,8 +321,8 @@ def __init__(
             log_of_measurement_results: A mutable object that measurements are
                 being recorded into.
         """
-        super().__init__(prng, axes, log_of_measurement_results)
-        self.qubit_map = qubit_map
+        super().__init__(prng, qubits, axes, log_of_measurement_results)
+        qubit_map = self.qubit_map
         self.grouping = qubit_map if grouping is None else grouping
         if self.grouping.keys() != self.qubit_map.keys():
             raise ValueError('Grouping must cover exactly the qubits.')
@@ -364,10 +378,10 @@ def _value_equality_values_(self) -> Any:
 
     def copy(self) -> 'MPSState':
         state = MPSState(
-            self.qubit_map,
-            self.prng,
-            self.simulation_options,
-            self.grouping,
+            qubits=self.qubits,
+            prng=self.prng,
+            simulation_options=self.simulation_options,
+            grouping=self.grouping,
         )
         state.M = [x.copy() for x in self.M]
         state.estimated_gate_error_list = self.estimated_gate_error_list
@@ -584,6 +598,5 @@ def perform_measurement(
 
     def _perform_measurement(self) -> List[int]:
         """Measures the axes specified by the simulator."""
-        qubit_map_inv = {v: k for k, v in self.qubit_map.items()}
-        qubits = [qubit_map_inv[key] for key in self.axes]
+        qubits = [self.qubits[key] for key in self.axes]
         return self.perform_measurement(qubits, self.prng)

cirq/contrib/quimb/mps_simulator_test.py

@@ -205,6 +205,23 @@ def test_cnot_flipped():
         )
 
 
+def test_act_on_args():
+    q0, q1 = qubit_order = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit(cirq.CNOT(q1, q0))
+    mps_simulator = ccq.mps_simulator.MPSSimulator()
+    ref_simulator = cirq.Simulator()
+    for initial_state in range(4):
+        args = mps_simulator.create_act_on_args(initial_state=initial_state, qubits=(q0, q1))
+        actual = mps_simulator.simulate(circuit, qubit_order=qubit_order, initial_state=args)
+        expected = ref_simulator.simulate(
+            circuit, qubit_order=qubit_order, initial_state=initial_state
+        )
+        np.testing.assert_allclose(
+            actual.final_state.to_numpy(), expected.final_state_vector, atol=1e-4
+        )
+        assert len(actual.measurements) == 0
+
+
 def test_three_qubits():
     q0, q1, q2 = cirq.LineQubit.range(3)
     circuit = cirq.Circuit(cirq.CCX(q0, q1, q2))
@@ -257,7 +274,7 @@ def test_measurement_str():
 def test_trial_result_str():
     q0 = cirq.LineQubit(0)
     final_simulator_state = ccq.mps_simulator.MPSState(
-        qubit_map={q0: 0},
+        qubits=(q0,),
         prng=value.parse_random_state(0),
         simulation_options=ccq.mps_simulator.MPSOptions(),
     )
@@ -278,7 +295,7 @@ def test_trial_result_str():
 
 def test_empty_step_result():
     q0 = cirq.LineQubit(0)
-    state = ccq.mps_simulator.MPSState(qubit_map={q0: 0}, prng=value.parse_random_state(0))
+    state = ccq.mps_simulator.MPSState(qubits=(q0,), prng=value.parse_random_state(0))
     step_result = ccq.mps_simulator.MPSSimulatorStepResult(state, measurements={'0': [1]})
     assert (
         str(step_result)
@@ -292,17 +309,17 @@ def test_empty_step_result():
 def test_state_equal():
     q0, q1 = cirq.LineQubit.range(2)
     state0 = ccq.mps_simulator.MPSState(
-        qubit_map={q0: 0},
+        qubits=(q0,),
         prng=value.parse_random_state(0),
         simulation_options=ccq.mps_simulator.MPSOptions(cutoff=1e-3, sum_prob_atol=1e-3),
     )
     state1a = ccq.mps_simulator.MPSState(
-        qubit_map={q1: 0},
+        qubits=(q1,),
         prng=value.parse_random_state(0),
         simulation_options=ccq.mps_simulator.MPSOptions(cutoff=1e-3, sum_prob_atol=1e-3),
     )
     state1b = ccq.mps_simulator.MPSState(
-        qubit_map={q1: 0},
+        qubits=(q1,),
         prng=value.parse_random_state(0),
         simulation_options=ccq.mps_simulator.MPSOptions(cutoff=1729.0, sum_prob_atol=1e-3),
     )
@@ -500,7 +517,7 @@ def test_state_copy():
 
 def test_state_act_on_args_initializer():
     s = ccq.mps_simulator.MPSState(
-        qubit_map={cirq.LineQubit(0): 0},
+        qubits=(cirq.LineQubit(0),),
         prng=np.random.RandomState(0),
         axes=[2],
         log_of_measurement_results={'test': 4},

cirq/google/calibration/engine_simulator.py

@@ -3,7 +3,6 @@
     Callable,
     Dict,
     Iterable,
-    Iterator,
     List,
     Optional,
     Sequence,
@@ -14,38 +13,35 @@
 import numpy as np
 
 from cirq.circuits import Circuit, PointOptimizer, PointOptimizationSummary
+from cirq.google.calibration.phased_fsim import (
+    FloquetPhasedFSimCalibrationRequest,
+    PhaseCalibratedFSimGate,
+    IncompatibleMomentError,
+    PhasedFSimCalibrationRequest,
+    PhasedFSimCalibrationResult,
+    PhasedFSimCharacterization,
+    SQRT_ISWAP_PARAMETERS,
+    try_convert_sqrt_iswap_to_fsim,
+)
 from cirq.ops import (
     FSimGate,
     Gate,
     MeasurementGate,
     Operation,
     PhasedFSimGate,
     Qid,
-    QubitOrderOrList,
     SingleQubitGate,
     WaitGate,
 )
 from cirq.sim import (
     Simulator,
     SimulatesSamples,
     SimulatesIntermediateStateVector,
-    StateVectorStepResult,
+    ActOnStateVectorArgs,
 )
 from cirq.study import ParamResolver
 from cirq.value import RANDOM_STATE_OR_SEED_LIKE, parse_random_state
 
-from cirq.google.calibration.phased_fsim import (
-    FloquetPhasedFSimCalibrationRequest,
-    PhaseCalibratedFSimGate,
-    IncompatibleMomentError,
-    PhasedFSimCalibrationRequest,
-    PhasedFSimCalibrationResult,
-    PhasedFSimCharacterization,
-    SQRT_ISWAP_PARAMETERS,
-    try_convert_sqrt_iswap_to_fsim,
-)
-
-
 ParametersDriftGenerator = Callable[[Qid, Qid, FSimGate], PhasedFSimCharacterization]
 PhasedFsimDictParameters = Dict[
     Tuple[Qid, Qid], Union[Dict[str, float], PhasedFSimCharacterization]
@@ -395,14 +391,20 @@ def _run(
         converted = _convert_to_circuit_with_drift(self, circuit)
         return self._simulator._run(converted, param_resolver, repetitions)
 
-    def _base_iterator(
+    def _core_iterator(
         self,
         circuit: Circuit,
-        qubit_order: QubitOrderOrList,
-        initial_state: Any,
-    ) -> Iterator[StateVectorStepResult]:
+        sim_state: Any,
+    ):
         converted = _convert_to_circuit_with_drift(self, circuit)
-        return self._simulator._base_iterator(converted, qubit_order, initial_state)
+        return self._simulator._core_iterator(converted, sim_state)
+
+    def create_act_on_args(
+        self,
+        initial_state: Union[int, ActOnStateVectorArgs],
+        qubits: Sequence[Qid],
+    ):
+        return self._simulator.create_act_on_args(initial_state, qubits)
 
 
 class _PhasedFSimConverter(PointOptimizer):