Deduplicate shared preamble in expectation-value sweeps

cirq-core/cirq/sim/density_matrix_simulator.py

@@ -200,7 +200,6 @@ def _create_simulator_trial_result(
             params=params, measurements=measurements, final_simulator_state=final_simulator_state
         )
 
-    # TODO(#4209): Deduplicate with identical code in sparse_simulator.
     def simulate_expectation_values_sweep(
         self,
         program: cirq.AbstractCircuit,
@@ -210,18 +209,10 @@ def simulate_expectation_values_sweep(
         initial_state: Any = None,
         permit_terminal_measurements: bool = False,
     ) -> list[list[float]]:
-        if not permit_terminal_measurements and program.are_any_measurements_terminal():
-            raise ValueError(
-                'Provided circuit has terminal measurements, which may '
-                'skew expectation values. If this is intentional, set '
-                'permit_terminal_measurements=True.'
-            )
+        qubit_order, qmap, pslist = self._qubit_map_and_pauli_sums(
+            program, observables, qubit_order, permit_terminal_measurements
+        )
         swept_evs = []
-        qubit_order = ops.QubitOrder.as_qubit_order(qubit_order)
-        qmap = {q: i for i, q in enumerate(qubit_order.order_for(program.all_qubits()))}
-        if not isinstance(observables, list):
-            observables = [observables]
-        pslist = [ops.PauliSum.wrap(pslike) for pslike in observables]
         for param_resolver in study.to_resolvers(params):
             result = self.simulate(
                 program, param_resolver, qubit_order=qubit_order, initial_state=initial_state

cirq-core/cirq/sim/simulator.py

@@ -439,6 +439,48 @@ def simulate_expectation_values_sweep_iter(
         """
         raise NotImplementedError
 
+    def _qubit_map_and_pauli_sums(
+        self,
+        program: cirq.AbstractCircuit,
+        observables: cirq.PauliSumLike | list[cirq.PauliSumLike],
+        qubit_order: cirq.QubitOrderOrList,
+        permit_terminal_measurements: bool,
+    ) -> tuple[cirq.QubitOrder, dict[cirq.Qid, int], list[cirq.PauliSum]]:
+        """Validates inputs and builds the data shared by expectation-value sweeps.
+
+        Runs the common preamble for `simulate_expectation_values_sweep` and
+        `simulate_expectation_values_sweep_iter` implementations: it rejects
+        terminal measurements (unless permitted), resolves the qubit order, and
+        wraps the observables as `cirq.PauliSum`s.
+
+        Args:
+            program: The circuit whose qubits define the ordering.
+            observables: An observable or list of observables.
+            qubit_order: Determines the canonical ordering of the qubits.
+            permit_terminal_measurements: If False, raises when `program` ends
+                with measurement(s).
+
+        Returns:
+            A tuple of the resolved `cirq.QubitOrder`, the qubit-to-index map, and
+            the observables wrapped as a list of `cirq.PauliSum`.
+
+        Raises:
+            ValueError: If `program` has terminal measurement(s) and
+                `permit_terminal_measurements` is False.
+        """
+        if not permit_terminal_measurements and program.are_any_measurements_terminal():
+            raise ValueError(
+                'Provided circuit has terminal measurements, which may '
+                'skew expectation values. If this is intentional, set '
+                'permit_terminal_measurements=True.'
+            )
+        qubit_order = ops.QubitOrder.as_qubit_order(qubit_order)
+        qmap = {q: i for i, q in enumerate(qubit_order.order_for(program.all_qubits()))}
+        if not isinstance(observables, list):
+            observables = [observables]
+        pslist = [ops.PauliSum.wrap(pslike) for pslike in observables]
+        return qubit_order, qmap, pslist
+
 
 class SimulatesFinalState(
     Generic[TSimulationTrialResult], metaclass=value.ABCMetaImplementAnyOneOf

cirq-core/cirq/sim/sparse_simulator.py

@@ -199,17 +199,9 @@ def simulate_expectation_values_sweep_iter(
         initial_state: Any = None,
         permit_terminal_measurements: bool = False,
     ) -> Iterator[list[float]]:
-        if not permit_terminal_measurements and program.are_any_measurements_terminal():
-            raise ValueError(
-                'Provided circuit has terminal measurements, which may '
-                'skew expectation values. If this is intentional, set '
-                'permit_terminal_measurements=True.'
-            )
-        qubit_order = ops.QubitOrder.as_qubit_order(qubit_order)
-        qmap = {q: i for i, q in enumerate(qubit_order.order_for(program.all_qubits()))}
-        if not isinstance(observables, list):
-            observables = [observables]
-        pslist = [ops.PauliSum.wrap(pslike) for pslike in observables]
+        qubit_order, qmap, pslist = self._qubit_map_and_pauli_sums(
+            program, observables, qubit_order, permit_terminal_measurements
+        )
         yield from (
             [obs.expectation_from_state_vector(result.final_state_vector, qmap) for obs in pslist]
             for result in self.simulate_sweep_iter(