Make Collector and PauliSumCollector work with AbstractCircuit (#6109)

quantumlib · May 26, 2023 · 36d67c1 · 36d67c1
1 parent 16ed00d
commit 36d67c1
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 10 deletions.
diff --git a/cirq-core/cirq/work/collector.py b/cirq-core/cirq/work/collector.py
@@ -19,7 +19,7 @@
 import duet
 import numpy as np
 
-from cirq import circuits, study, value
+from cirq import study, value
 
 if TYPE_CHECKING:
     import cirq
@@ -29,7 +29,7 @@
 class CircuitSampleJob:
     """Describes a sampling task."""
 
-    def __init__(self, circuit: circuits.Circuit, *, repetitions: int, tag: Any = None):
+    def __init__(self, circuit: 'cirq.AbstractCircuit', *, repetitions: int, tag: Any = None):
         """Inits CircuitSampleJob.
 
         Args:

diff --git a/cirq-core/cirq/work/pauli_sum_collector.py b/cirq-core/cirq/work/pauli_sum_collector.py
@@ -17,7 +17,7 @@
 
 import numpy as np
 
-from cirq import circuits, ops
+from cirq import ops
 from cirq.work import collector
 
 if TYPE_CHECKING:
@@ -29,7 +29,7 @@ class PauliSumCollector(collector.Collector):
 
     def __init__(
         self,
-        circuit: 'cirq.Circuit',
+        circuit: 'cirq.AbstractCircuit',
         observable: 'cirq.PauliSumLike',
         *,
         samples_per_term: int,
@@ -99,11 +99,12 @@ def estimated_energy(self) -> Union[float, complex]:
 
 
 def _circuit_plus_pauli_string_measurements(
-    circuit: 'cirq.Circuit', pauli_string: 'cirq.PauliString'
-) -> 'cirq.Circuit':
+    circuit: 'cirq.AbstractCircuit', pauli_string: 'cirq.PauliString'
+) -> 'cirq.AbstractCircuit':
     """A circuit measuring the given observable at the end of the given circuit."""
     assert pauli_string
-    circuit = circuit.copy()
-    circuit.append(circuits.Moment(pauli_string.to_z_basis_ops()))
-    circuit.append(circuits.Moment([ops.measure(*sorted(pauli_string.keys()), key='out')]))
-    return circuit
+    return circuit.from_moments(
+        *circuit,
+        pauli_string.to_z_basis_ops(),
+        ops.measure(*sorted(pauli_string.keys()), key='out'),
+    )