Provide sweep-iterator API. (#4094)

Fixes #4064.

This PR provides `*_iter` methods for all simulator "sweep" methods. External simulators which implement the old list-based APIs will still function properly (through the magic of `@value.alternative`), but these simulators are advised to switch to the iterator-based API for better performance.

Note that the list-based API is _not_ deprecated, and we have no plans to remove it at this time.

Author: 95-martin-orion

cirq-core/cirq/__init__.py

@@ -429,6 +429,7 @@
     chosen_angle_to_half_turns,
     Duration,
     DURATION_LIKE,
+    GenericMetaImplementAnyOneOf,
     LinearDict,
     MEASUREMENT_KEY_SEPARATOR,
     MeasurementKey,

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

@@ -119,6 +119,10 @@
         'ThreeQubitGate',
         'TwoQubitGate',
         'ABCMetaImplementAnyOneOf',
+        'GenericMetaImplementAnyOneOf',
+        'SimulatesAmplitudes',
+        'SimulatesExpectationValues',
+        'SimulatesFinalState',
         # protocols:
         'SupportsActOn',
         'SupportsApplyChannel',

cirq-core/cirq/sim/simulator.py

@@ -73,6 +73,14 @@ def run_sweep(
         params: study.Sweepable,
         repetitions: int = 1,
     ) -> List[study.Result]:
+        return list(self.run_sweep_iter(program, params, repetitions))
+
+    def run_sweep_iter(
+        self,
+        program: 'cirq.Circuit',
+        params: study.Sweepable,
+        repetitions: int = 1,
+    ) -> Iterator[study.Result]:
         """Runs the supplied Circuit, mimicking quantum hardware.
 
         In contrast to run, this allows for sweeping over different parameter
@@ -92,7 +100,6 @@ def run_sweep(
 
         _verify_unique_measurement_keys(program)
 
-        trial_results = []  # type: List[study.Result]
         for param_resolver in study.to_resolvers(params):
             measurements = {}
             if repetitions == 0:
@@ -102,12 +109,9 @@ def run_sweep(
                 measurements = self._run(
                     circuit=program, param_resolver=param_resolver, repetitions=repetitions
                 )
-            trial_results.append(
-                study.Result.from_single_parameter_set(
-                    params=param_resolver, measurements=measurements
-                )
+            yield study.Result.from_single_parameter_set(
+                params=param_resolver, measurements=measurements
             )
-        return trial_results
 
     @abc.abstractmethod
     def _run(
@@ -131,13 +135,13 @@ def _run(
         raise NotImplementedError()
 
 
-class SimulatesAmplitudes(metaclass=abc.ABCMeta):
+class SimulatesAmplitudes(metaclass=value.ABCMetaImplementAnyOneOf):
     """Simulator that computes final amplitudes of given bitstrings.
 
     Given a circuit and a list of bitstrings, computes the amplitudes
     of the given bitstrings in the state obtained by applying the circuit
     to the all zeros state. Implementors of this interface should implement
-    the compute_amplitudes_sweep method.
+    the compute_amplitudes_sweep_iter method.
     """
 
     def compute_amplitudes(
@@ -169,14 +173,42 @@ def compute_amplitudes(
             program, bitstrings, study.ParamResolver(param_resolver), qubit_order
         )[0]
 
-    @abc.abstractmethod
     def compute_amplitudes_sweep(
         self,
         program: 'cirq.Circuit',
         bitstrings: Sequence[int],
         params: study.Sweepable,
         qubit_order: ops.QubitOrderOrList = ops.QubitOrder.DEFAULT,
     ) -> Sequence[Sequence[complex]]:
+        """Wraps computed amplitudes in a list.
+
+        Prefer overriding `compute_amplitudes_sweep_iter`.
+        """
+        return list(self.compute_amplitudes_sweep_iter(program, bitstrings, params, qubit_order))
+
+    def _compute_amplitudes_sweep_to_iter(
+        self,
+        program: 'cirq.Circuit',
+        bitstrings: Sequence[int],
+        params: study.Sweepable,
+        qubit_order: ops.QubitOrderOrList = ops.QubitOrder.DEFAULT,
+    ) -> Iterator[Sequence[complex]]:
+        if type(self).compute_amplitudes_sweep == SimulatesAmplitudes.compute_amplitudes_sweep:
+            raise RecursionError(
+                "Must define either compute_amplitudes_sweep or compute_amplitudes_sweep_iter."
+            )
+        yield from self.compute_amplitudes_sweep(program, bitstrings, params, qubit_order)
+
+    @value.alternative(
+        requires='compute_amplitudes_sweep', implementation=_compute_amplitudes_sweep_to_iter
+    )
+    def compute_amplitudes_sweep_iter(
+        self,
+        program: 'cirq.Circuit',
+        bitstrings: Sequence[int],
+        params: study.Sweepable,
+        qubit_order: ops.QubitOrderOrList = ops.QubitOrder.DEFAULT,
+    ) -> Iterator[Sequence[complex]]:
         """Computes the desired amplitudes.
 
         The initial state is assumed to be the all zeros state.
@@ -193,20 +225,20 @@ def compute_amplitudes_sweep(
                 ordering of the computational basis states.
 
         Returns:
-            List of lists of amplitudes. The outer dimension indexes the
-            circuit parameters and the inner dimension indexes the bitstrings.
+            An Iterator over lists of amplitudes. The outer dimension indexes
+            the circuit parameters and the inner dimension indexes bitstrings.
         """
         raise NotImplementedError()
 
 
-class SimulatesExpectationValues(metaclass=abc.ABCMeta):
+class SimulatesExpectationValues(metaclass=value.ABCMetaImplementAnyOneOf):
     """Simulator that computes exact expectation values of observables.
 
     Given a circuit and an observable map, computes exact (to float precision)
     expectation values for each observable at the end of the circuit.
 
     Implementors of this interface should implement the
-    simulate_expectation_values_sweep method.
+    simulate_expectation_values_sweep_iter method.
     """
 
     def simulate_expectation_values(
@@ -257,7 +289,6 @@ def simulate_expectation_values(
             permit_terminal_measurements,
         )[0]
 
-    @abc.abstractmethod
     def simulate_expectation_values_sweep(
         self,
         program: 'cirq.Circuit',
@@ -267,6 +298,60 @@ def simulate_expectation_values_sweep(
         initial_state: Any = None,
         permit_terminal_measurements: bool = False,
     ) -> List[List[float]]:
+        """Wraps computed expectation values in a list.
+
+        Prefer overriding `simulate_expectation_values_sweep_iter`.
+        """
+        return list(
+            self.simulate_expectation_values_sweep_iter(
+                program,
+                observables,
+                params,
+                qubit_order,
+                initial_state,
+                permit_terminal_measurements,
+            )
+        )
+
+    def _simulate_expectation_values_sweep_to_iter(
+        self,
+        program: 'cirq.Circuit',
+        observables: Union['cirq.PauliSumLike', List['cirq.PauliSumLike']],
+        params: 'study.Sweepable',
+        qubit_order: ops.QubitOrderOrList = ops.QubitOrder.DEFAULT,
+        initial_state: Any = None,
+        permit_terminal_measurements: bool = False,
+    ) -> Iterator[List[float]]:
+        if (
+            type(self).simulate_expectation_values_sweep
+            == SimulatesExpectationValues.simulate_expectation_values_sweep
+        ):
+            raise RecursionError(
+                "Must define either simulate_expectation_values_sweep or "
+                "simulate_expectation_values_sweep_iter."
+            )
+        yield from self.simulate_expectation_values_sweep(
+            program,
+            observables,
+            params,
+            qubit_order,
+            initial_state,
+            permit_terminal_measurements,
+        )
+
+    @value.alternative(
+        requires='simulate_expectation_values_sweep',
+        implementation=_simulate_expectation_values_sweep_to_iter,
+    )
+    def simulate_expectation_values_sweep_iter(
+        self,
+        program: 'cirq.Circuit',
+        observables: Union['cirq.PauliSumLike', List['cirq.PauliSumLike']],
+        params: 'study.Sweepable',
+        qubit_order: ops.QubitOrderOrList = ops.QubitOrder.DEFAULT,
+        initial_state: Any = None,
+        permit_terminal_measurements: bool = False,
+    ) -> Iterator[List[float]]:
         """Simulates the supplied circuit and calculates exact expectation
         values for the given observables on its final state, sweeping over the
         given params.
@@ -291,21 +376,23 @@ def simulate_expectation_values_sweep(
                 ruining expectation value calculations.
 
         Returns:
-            A list of expectation-value lists. The outer index determines the
-            sweep, and the inner index determines the observable. For instance,
-            results[1][3] would select the fourth observable measured in the
-            second sweep.
+            An Iterator over expectation-value lists. The outer index determines
+            the sweep, and the inner index determines the observable. For
+            instance, results[1][3] would select the fourth observable measured
+            in the second sweep.
 
         Raises:
             ValueError if 'program' has terminal measurement(s) and
             'permit_terminal_measurements' is False.
         """
 
 
-class SimulatesFinalState(Generic[TSimulationTrialResult], metaclass=abc.ABCMeta):
+class SimulatesFinalState(
+    Generic[TSimulationTrialResult], metaclass=value.GenericMetaImplementAnyOneOf
+):
     """Simulator that allows access to the simulator's final state.
 
-    Implementors of this interface should implement the simulate_sweep
+    Implementors of this interface should implement the simulate_sweep_iter
     method. This simulator only returns the state of the quantum system
     for the final step of a simulation. This simulator state may be a state
     vector, the density matrix, or another representation, depending on the
@@ -342,14 +429,38 @@ def simulate(
             program, study.ParamResolver(param_resolver), qubit_order, initial_state
         )[0]
 
-    @abc.abstractmethod
     def simulate_sweep(
         self,
         program: 'cirq.Circuit',
         params: study.Sweepable,
         qubit_order: ops.QubitOrderOrList = ops.QubitOrder.DEFAULT,
         initial_state: Any = None,
     ) -> List[TSimulationTrialResult]:
+        """Wraps computed states in a list.
+
+        Prefer overriding `simulate_sweep_iter`.
+        """
+        return list(self.simulate_sweep_iter(program, params, qubit_order, initial_state))
+
+    def _simulate_sweep_to_iter(
+        self,
+        program: 'cirq.Circuit',
+        params: study.Sweepable,
+        qubit_order: ops.QubitOrderOrList = ops.QubitOrder.DEFAULT,
+        initial_state: Any = None,
+    ) -> Iterator[TSimulationTrialResult]:
+        if type(self).simulate_sweep == SimulatesFinalState.simulate_sweep:
+            raise RecursionError("Must define either simulate_sweep or simulate_sweep_iter.")
+        yield from self.simulate_sweep(program, params, qubit_order, initial_state)
+
+    @value.alternative(requires='simulate_sweep', implementation=_simulate_sweep_to_iter)
+    def simulate_sweep_iter(
+        self,
+        program: 'cirq.Circuit',
+        params: study.Sweepable,
+        qubit_order: ops.QubitOrderOrList = ops.QubitOrder.DEFAULT,
+        initial_state: Any = None,
+    ) -> Iterator[TSimulationTrialResult]:
         """Simulates the supplied Circuit.
 
         This method returns a result which allows access to the entire final
@@ -367,7 +478,7 @@ def simulate_sweep(
                 documentation of the implementing class for details.
 
         Returns:
-            List of SimulationTrialResults for this run, one for each
+            Iterator over SimulationTrialResults for this run, one for each
             possible parameter resolver.
         """
         raise NotImplementedError()
@@ -391,13 +502,13 @@ class SimulatesIntermediateState(
     a state vector.
     """
 
-    def simulate_sweep(
+    def simulate_sweep_iter(
         self,
         program: 'cirq.Circuit',
         params: study.Sweepable,
         qubit_order: ops.QubitOrderOrList = ops.QubitOrder.DEFAULT,
         initial_state: Any = None,
-    ) -> List[TSimulationTrialResult]:
+    ) -> Iterator[TSimulationTrialResult]:
         """Simulates the supplied Circuit.
 
         This method returns a result which allows access to the entire
@@ -419,7 +530,6 @@ def simulate_sweep(
             List of SimulationTrialResults for this run, one for each
             possible parameter resolver.
         """
-        trial_results = []
         qubit_order = ops.QubitOrder.as_qubit_order(qubit_order)
         for param_resolver in study.to_resolvers(params):
             all_step_results = self.simulate_moment_steps(
@@ -429,14 +539,11 @@ def simulate_sweep(
             for step_result in all_step_results:
                 for k, v in step_result.measurements.items():
                     measurements[k] = np.array(v, dtype=np.uint8)
-            trial_results.append(
-                self._create_simulator_trial_result(
-                    params=param_resolver,
-                    measurements=measurements,
-                    final_simulator_state=step_result._simulator_state(),
-                )
+            yield self._create_simulator_trial_result(
+                params=param_resolver,
+                measurements=measurements,
+                final_simulator_state=step_result._simulator_state(),
             )
-        return trial_results
 
     def simulate_moment_steps(
         self,