Repository: quantumlib/Cirq

Issue:
Optimize parameter sweeps to cache unparameterized prefixes
**Is your feature request related to a use case or problem? Please describe.**
Sweeps rerun the whole circuit from scratch for each parameter set. They could be made more efficient by caching the result of the unparameterized prefix and restarting each parameterized simulation from there.

**Describe the solution you'd like**
We do something similar in `run`, where we cache the result of the unitary prefix, and rerun repetitions starting from there. So it should be straightforward to add this same style of optimization for sweeps.

**Describe alternatives/workarounds you've considered**

* Could do nothing: I assume in practice, parameterized circuits usually have *some* parameters at or near the beginning. Also once we merge the feature that simulates unentangled qubits independently, the first part of the circuit should be a tiny fraction of the simulation time.
* Going the other direction, this could be extended to where if the first N parameters are the same in some parameter sets, it can start each run from the Nth (factoring the circuit would be a fun challenge).
* Tangentially, I wonder if this logic really belongs in the simulator class itself. Seems these could be module-level functions since they just call public simulator methods. This would allow us to create different sweep functions for different optimization styles without having to create a specific simulator class for each.


**What is the urgency from your perspective for this issue? Is it blocking important work?**
P3 - I'm not really blocked by it, it is an idea I'd like to discuss / suggestion based on principle 


Discussion:

Resolving patch:
diff --git a/cirq-core/cirq/sim/simulator.py b/cirq-core/cirq/sim/simulator.py
--- a/cirq-core/cirq/sim/simulator.py
+++ b/cirq-core/cirq/sim/simulator.py
@@ -50,6 +50,7 @@
 
 from cirq import circuits, ops, protocols, study, value, work
 from cirq.sim.act_on_args import ActOnArgs
+from cirq.sim.operation_target import OperationTarget
 
 if TYPE_CHECKING:
     import cirq
@@ -538,8 +539,13 @@ def simulate_sweep_iter(
         """
         qubit_order = ops.QubitOrder.as_qubit_order(qubit_order)
         for param_resolver in study.to_resolvers(params):
+            state = (
+                initial_state.copy()
+                if isinstance(initial_state, OperationTarget)
+                else initial_state
+            )
             all_step_results = self.simulate_moment_steps(
-                program, param_resolver, qubit_order, initial_state
+                program, param_resolver, qubit_order, state
             )
             measurements = {}  # type: Dict[str, np.ndarray]
             for step_result in all_step_results:
diff --git a/cirq-core/cirq/sim/simulator_base.py b/cirq-core/cirq/sim/simulator_base.py
--- a/cirq-core/cirq/sim/simulator_base.py
+++ b/cirq-core/cirq/sim/simulator_base.py
@@ -281,6 +281,55 @@ def _run(
                 measurements[k].append(np.array(v, dtype=np.uint8))
         return {k: np.array(v) for k, v in measurements.items()}
 
+    def simulate_sweep_iter(
+        self,
+        program: 'cirq.AbstractCircuit',
+        params: study.Sweepable,
+        qubit_order: ops.QubitOrderOrList = ops.QubitOrder.DEFAULT,
+        initial_state: Any = None,
+    ) -> Iterator[TSimulationTrialResult]:
+        """Simulates the supplied Circuit.
+
+        This particular implementation overrides the base implementation such
+        that an unparameterized prefix circuit is simulated and fed into the
+        parameterized suffix circuit.
+
+        Args:
+            program: The circuit to simulate.
+            params: Parameters to run with the program.
+            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. This can be
+                either a raw state or an `OperationTarget`. The form of the
+                raw state depends on the simulation implementation. See
+                documentation of the implementing class for details.
+
+        Returns:
+            List of SimulationTrialResults for this run, one for each
+            possible parameter resolver.
+        """
+
+        def sweep_prefixable(op: 'cirq.Operation'):
+            return self._can_be_in_run_prefix(op) and not protocols.is_parameterized(op)
+
+        qubits = ops.QubitOrder.as_qubit_order(qubit_order).order_for(program.all_qubits())
+        initial_state = 0 if initial_state is None else initial_state
+        sim_state = self._create_act_on_args(initial_state, qubits)
+        prefix, suffix = (
+            split_into_matching_protocol_then_general(program, sweep_prefixable)
+            if self._can_be_in_run_prefix(self.noise)
+            else (program[0:0], program)
+        )
+        step_result = None
+        for step_result in self._core_iterator(
+            circuit=prefix,
+            sim_state=sim_state,
+        ):
+            pass
+        sim_state = step_result._sim_state
+        yield from super().simulate_sweep_iter(suffix, params, qubit_order, sim_state)
+
     def _create_act_on_args(
         self,
         initial_state: Any,


Test patch:
diff --git a/cirq-core/cirq/sim/density_matrix_simulator_test.py b/cirq-core/cirq/sim/density_matrix_simulator_test.py
--- a/cirq-core/cirq/sim/density_matrix_simulator_test.py
+++ b/cirq-core/cirq/sim/density_matrix_simulator_test.py
@@ -1557,3 +1557,34 @@ def test_large_untangled_okay():
     assert len(result.measurements) == 59
     assert len(result.measurements['0']) == 1000
     assert (result.measurements['0'] == np.full(1000, 1)).all()
+
+
+def test_sweep_unparameterized_prefix_not_repeated_even_non_unitaries():
+    q = cirq.LineQubit(0)
+
+    class NonUnitaryOp(cirq.Operation):
+        count = 0
+
+        def _act_on_(self, args):
+            self.count += 1
+            return True
+
+        def with_qubits(self, qubits):
+            pass
+
+        @property
+        def qubits(self):
+            return (q,)
+
+    simulator = cirq.DensityMatrixSimulator()
+    params = [
+        cirq.ParamResolver({'a': 0}),
+        cirq.ParamResolver({'a': 1}),
+    ]
+
+    op1 = NonUnitaryOp()
+    op2 = NonUnitaryOp()
+    circuit = cirq.Circuit(op1, cirq.XPowGate(exponent=sympy.Symbol('a'))(q), op2)
+    simulator.simulate_sweep(program=circuit, params=params)
+    assert op1.count == 1
+    assert op2.count == 2
diff --git a/cirq-core/cirq/sim/simulator_base_test.py b/cirq-core/cirq/sim/simulator_base_test.py
--- a/cirq-core/cirq/sim/simulator_base_test.py
+++ b/cirq-core/cirq/sim/simulator_base_test.py
@@ -16,6 +16,7 @@
 
 import numpy as np
 import pytest
+import sympy
 
 import cirq
 
@@ -356,3 +357,46 @@ def test_measurements_retained_in_step_results():
     assert next(iterator).measurements.keys() == {'a', 'b'}
     assert next(iterator).measurements.keys() == {'a', 'b', 'c'}
     assert not any(iterator)
+
+
+def test_sweep_unparameterized_prefix_not_repeated_iff_unitary():
+    q = cirq.LineQubit(0)
+
+    class TestOp(cirq.Operation):
+        def __init__(self, *, has_unitary: bool):
+            self.count = 0
+            self.has_unitary = has_unitary
+
+        def _act_on_(self, args):
+            self.count += 1
+            return True
+
+        def with_qubits(self, qubits):
+            pass
+
+        @property
+        def qubits(self):
+            return (q,)
+
+        def _has_unitary_(self):
+            return self.has_unitary
+
+    simulator = CountingSimulator()
+    params = [
+        cirq.ParamResolver({'a': 0}),
+        cirq.ParamResolver({'a': 1}),
+    ]
+
+    op1 = TestOp(has_unitary=True)
+    op2 = TestOp(has_unitary=True)
+    circuit = cirq.Circuit(op1, cirq.XPowGate(exponent=sympy.Symbol('a'))(q), op2)
+    simulator.simulate_sweep(program=circuit, params=params)
+    assert op1.count == 1
+    assert op2.count == 2
+
+    op1 = TestOp(has_unitary=False)
+    op2 = TestOp(has_unitary=False)
+    circuit = cirq.Circuit(op1, cirq.XPowGate(exponent=sympy.Symbol('a'))(q), op2)
+    simulator.simulate_sweep(program=circuit, params=params)
+    assert op1.count == 2
+    assert op2.count == 2
diff --git a/cirq-core/cirq/sim/sparse_simulator_test.py b/cirq-core/cirq/sim/sparse_simulator_test.py
--- a/cirq-core/cirq/sim/sparse_simulator_test.py
+++ b/cirq-core/cirq/sim/sparse_simulator_test.py
@@ -1331,4 +1331,4 @@ def test_noise_model():
     simulator = cirq.Simulator(noise=noise_model)
     result = simulator.run(circuit, repetitions=100)
 
-    assert 40 <= sum(result.measurements['0'])[0] < 60
+    assert 20 <= sum(result.measurements['0'])[0] < 80