Allow measurements and resets into Gemini noise model. (#583)

To simulate some operations, for example multiple rounds of syndrome
extraction, it is useful to be able to measure and reuse the measurement
qubits in the cirq simulation to save on RAM (even if on Gemini we will
actually use different measurement atoms for each round of syndrome
extraction).

Here is a minimum fix to allow for measurement and reset gates in the
noise models. The next step would be to add error in the measurement
that is proportional to the number of gate operations left in the
circuit, since in Gemini we will have to wait until the end of the
circuit to measure any atom.

@david-pl what do you think is the right course of action?

---------

Co-authored-by: David Plankensteiner <[email protected]>
Co-authored-by: David Plankensteiner <[email protected]>

Author: tcochran-quera

src/bloqade/cirq_utils/noise/model.py

@@ -98,7 +98,10 @@ def __post_init__(self):
 
     @staticmethod
     def validate_moments(moments: Iterable[cirq.Moment]):
-        allowed_target_gates: frozenset[cirq.GateFamily] = cirq.CZTargetGateset().gates
+        reset_family = cirq.GateFamily(gate=cirq.ResetChannel, ignore_global_phase=True)
+        allowed_target_gates: frozenset[cirq.GateFamily] = cirq.CZTargetGateset(
+            additional_gates=[reset_family]
+        ).gates
 
         for moment in moments:
             for operation in moment:
@@ -117,7 +120,7 @@ def validate_moments(moments: Iterable[cirq.Moment]):
                     )
 
     def parallel_cz_errors(
-        self, ctrls: list[int], qargs: list[int], rest: list[int]
+        self, ctrls: Sequence[int], qargs: Sequence[int], rest: Sequence[int]
     ) -> dict[tuple[float, float, float, float], list[int]]:
         raise NotImplementedError(
             "This noise model doesn't support rewrites on bloqade kernels, but should be used with cirq."
@@ -245,8 +248,15 @@ def noisy_moment(self, moment, system_qubits):
         # Moment with original ops
         original_moment = moment
 
+        no_noise_condition = (
+            len(moment.operations) == 0
+            or cirq.is_measurement(moment.operations[0])
+            or isinstance(moment.operations[0].gate, cirq.ResetChannel)
+            or isinstance(moment.operations[0].gate, cirq.BitFlipChannel)
+        )
+
         # Check if the moment is empty
-        if len(moment.operations) == 0:
+        if no_noise_condition:
             move_noise_ops = []
             gate_noise_ops = []
         # Check if the moment contains 1-qubit gates or 2-qubit gates
@@ -319,20 +329,84 @@ def noisy_moments(
 
         # Split into moments with only 1Q and 2Q gates
         moments_1q = [
-            cirq.Moment([op for op in moment.operations if len(op.qubits) == 1])
+            cirq.Moment(
+                [
+                    op
+                    for op in moment.operations
+                    if (len(op.qubits) == 1)
+                    and (not cirq.is_measurement(op))
+                    and (not isinstance(op.gate, cirq.ResetChannel))
+                ]
+            )
             for moment in moments
         ]
         moments_2q = [
-            cirq.Moment([op for op in moment.operations if len(op.qubits) == 2])
+            cirq.Moment(
+                [
+                    op
+                    for op in moment.operations
+                    if (len(op.qubits) == 2) and (not cirq.is_measurement(op))
+                ]
+            )
             for moment in moments
         ]
 
-        assert len(moments_1q) == len(moments_2q)
+        moments_measurement = [
+            cirq.Moment(
+                [
+                    op
+                    for op in moment.operations
+                    if (cirq.is_measurement(op))
+                    or (isinstance(op.gate, cirq.ResetChannel))
+                ]
+            )
+            for moment in moments
+        ]
+
+        assert len(moments_1q) == len(moments_2q) == len(moments_measurement)
 
         interleaved_moments = []
+
+        def count_remaining_cz_moments(moments_2q):
+            remaining_cz_counts = []
+            count = 0
+            for m in moments_2q[::-1]:
+                if any(isinstance(op.gate, cirq.CZPowGate) for op in m.operations):
+                    count += 1
+                remaining_cz_counts = [count] + remaining_cz_counts
+            return remaining_cz_counts
+
+        remaining_cz_moments = count_remaining_cz_moments(moments_2q)
+
+        pm = 2 * self.sitter_pauli_rates[0]
+        ps = 2 * self.cz_unpaired_pauli_rates[0]
+
+        # probability of a bitflip error for a sitting, unpaired qubit during a move/cz/move cycle.
+        heuristic_1step_bitflip_error: float = (
+            2 * pm * (1 - ps) * (1 - pm) + (1 - pm) ** 2 * ps + pm**2 * ps
+        )
+
         for idx, moment in enumerate(moments_1q):
             interleaved_moments.append(moment)
             interleaved_moments.append(moments_2q[idx])
+            # Measurements on Gemini will be at the end, so for circuits with mid-circuit measurements we will insert a
+            # bitflip error proportional to the number of moments left in the circuit to account for the decoherence
+            # that will happen before the final terminal measurement.
+            measured_qubits = []
+            for op in moments_measurement[idx].operations:
+                if cirq.is_measurement(op):
+                    measured_qubits += list(op.qubits)
+            # probability of a bitflip error should be Binomial(moments_left,heuristic_1step_bitflip_error)
+            delayed_measurement_error = (
+                1
+                - (1 - 2 * heuristic_1step_bitflip_error) ** (remaining_cz_moments[idx])
+            ) / 2
+            interleaved_moments.append(
+                cirq.Moment(
+                    cirq.bit_flip(delayed_measurement_error).on_each(measured_qubits)
+                )
+            )
+            interleaved_moments.append(moments_measurement[idx])
 
         interleaved_circuit = cirq.Circuit.from_moments(*interleaved_moments)
 
@@ -368,14 +442,21 @@ def noisy_moment(self, moment, system_qubits):
             "all qubits in the circuit must be defined as cirq.GridQubit objects."
         )
         # Check if the moment is empty
-        if len(moment.operations) == 0:
+        if len(moment.operations) == 0 or cirq.is_measurement(moment.operations[0]):
             move_moments = []
             gate_noise_ops = []
         # Check if the moment contains 1-qubit gates or 2-qubit gates
         elif len(moment.operations[0].qubits) == 1:
-            gate_noise_ops, _ = self._single_qubit_moment_noise_ops(
-                moment, system_qubits
-            )
+            if (
+                (isinstance(moment.operations[0].gate, cirq.ResetChannel))
+                or (cirq.is_measurement(moment.operations[0]))
+                or (isinstance(moment.operations[0].gate, cirq.BitFlipChannel))
+            ):
+                gate_noise_ops = []
+            else:
+                gate_noise_ops, _ = self._single_qubit_moment_noise_ops(
+                    moment, system_qubits
+                )
             move_moments = []
         elif len(moment.operations[0].qubits) == 2:
             cg = OneZoneConflictGraph(moment)

src/bloqade/cirq_utils/parallelize.py

@@ -119,7 +119,11 @@ def auto_similarity(
     flattened_circuit: list[GateOperation] = list(cirq.flatten_op_tree(circuit))
     weights = {}
     for i in range(len(flattened_circuit)):
+        if not cirq.has_unitary(flattened_circuit[i]):
+            continue
         for j in range(i + 1, len(flattened_circuit)):
+            if not cirq.has_unitary(flattened_circuit[j]):
+                continue
             op1 = flattened_circuit[i]
             op2 = flattened_circuit[j]
             if can_be_parallel(op1, op2):
@@ -297,14 +301,20 @@ def colorize(
     for epoch in epochs:
         oneq_gates = []
         twoq_gates = []
+        nonunitary_gates = []
         for gate in epoch:
-            if len(gate.val.qubits) == 1:
+            if not cirq.has_unitary(gate.val):
+                nonunitary_gates.append(gate.val)
+            elif len(gate.val.qubits) == 1:
                 oneq_gates.append(gate.val)
             elif len(gate.val.qubits) == 2:
                 twoq_gates.append(gate.val)
             else:
                 raise RuntimeError("Unsupported gate type")
 
+        if len(nonunitary_gates) > 0:
+            yield nonunitary_gates
+
         if len(oneq_gates) > 0:
             yield oneq_gates
 

test/cirq_utils/noise/test_noise_models.py

@@ -14,7 +14,7 @@
 )
 
 
-def create_ghz_circuit(qubits):
+def create_ghz_circuit(qubits, measurements: bool = False):
     n = len(qubits)
     circuit = cirq.Circuit()
 
@@ -24,26 +24,41 @@ def create_ghz_circuit(qubits):
     # Step 2: CNOT chain from qubit i to i+1
     for i in range(n - 1):
         circuit.append(cirq.CNOT(qubits[i], qubits[i + 1]))
+        if measurements:
+            circuit.append(cirq.measure(qubits[i]))
+            circuit.append(cirq.reset(qubits[i]))
+
+    if measurements:
+        circuit.append(cirq.measure(qubits[-1]))
+        circuit.append(cirq.reset(qubits[-1]))
 
     return circuit
 
 
 @pytest.mark.parametrize(
-    "model,qubits",
+    "model,qubits,measurements",
     [
-        (GeminiOneZoneNoiseModel(), None),
+        (GeminiOneZoneNoiseModel(), None, False),
+        (
+            GeminiOneZoneNoiseModelConflictGraphMoves(),
+            cirq.GridQubit.rect(rows=1, cols=2),
+            False,
+        ),
+        (GeminiTwoZoneNoiseModel(), None, False),
+        (GeminiOneZoneNoiseModel(), None, True),
         (
             GeminiOneZoneNoiseModelConflictGraphMoves(),
             cirq.GridQubit.rect(rows=1, cols=2),
+            True,
         ),
-        (GeminiTwoZoneNoiseModel(), None),
+        (GeminiTwoZoneNoiseModel(), None, True),
     ],
 )
-def test_simple_model(model: cirq.NoiseModel, qubits):
+def test_simple_model(model: cirq.NoiseModel, qubits, measurements: bool):
     if qubits is None:
         qubits = cirq.LineQubit.range(2)
 
-    circuit = create_ghz_circuit(qubits)
+    circuit = create_ghz_circuit(qubits, measurements=measurements)
 
     with pytest.raises(ValueError):
         # make sure only native gate set is supported
@@ -74,13 +89,25 @@ def test_simple_model(model: cirq.NoiseModel, qubits):
         for i in range(4):
             pops_bloqade[i] += abs(ket[i]) ** 2 / nshots
 
-    for pops in (pops_bloqade, pops_cirq):
-        assert math.isclose(pops[0], 0.5, abs_tol=1e-1)
-        assert math.isclose(pops[3], 0.5, abs_tol=1e-1)
-        assert math.isclose(pops[1], 0.0, abs_tol=1e-1)
-        assert math.isclose(pops[2], 0.0, abs_tol=1e-1)
-
-        assert pops[0] < 0.5001
-        assert pops[3] < 0.5001
-        assert pops[1] >= 0.0
-        assert pops[2] >= 0.0
+    if measurements is True:
+        for pops in (pops_bloqade, pops_cirq):
+            assert math.isclose(pops[0], 1.0, abs_tol=1e-1)
+            assert math.isclose(pops[3], 0.0, abs_tol=1e-1)
+            assert math.isclose(pops[1], 0.0, abs_tol=1e-1)
+            assert math.isclose(pops[2], 0.0, abs_tol=1e-1)
+
+            assert pops[0] > 0.99
+            assert pops[3] >= 0.0
+            assert pops[1] >= 0.0
+            assert pops[2] >= 0.0
+    else:
+        for pops in (pops_bloqade, pops_cirq):
+            assert math.isclose(pops[0], 0.5, abs_tol=1e-1)
+            assert math.isclose(pops[3], 0.5, abs_tol=1e-1)
+            assert math.isclose(pops[1], 0.0, abs_tol=1e-1)
+            assert math.isclose(pops[2], 0.0, abs_tol=1e-1)
+
+            assert pops[0] < 0.5001
+            assert pops[3] < 0.5001
+            assert pops[1] >= 0.0
+            assert pops[2] >= 0.0

test/cirq_utils/test_parallelize.py

@@ -26,15 +26,67 @@ def test1():
     )
 
     circuit_m, _ = moment_similarity(circuit, weight=1.0)
-    # print(circuit_m)
     circuit_b, _ = block_similarity(circuit, weight=1.0, block_id=1)
-    circuit_m2 = remove_tags(circuit_m)
-    print(circuit_m2)
+    remove_tags(circuit_m)
     circuit2 = parallelize(circuit)
-    # print(circuit2)
     assert len(circuit2.moments) == 7
 
 
+def test_measurement_and_reset():
+    qubits = cirq.LineQubit.range(4)
+    circuit = cirq.Circuit(
+        cirq.H(qubits[0]),
+        cirq.CX(qubits[0], qubits[1]),
+        cirq.measure(qubits[1]),
+        cirq.reset(qubits[1]),
+        cirq.CX(qubits[1], qubits[2]),
+        cirq.measure(qubits[2]),
+        cirq.reset(qubits[2]),
+        cirq.CX(qubits[2], qubits[3]),
+        cirq.measure(qubits[0]),
+        cirq.reset(qubits[0]),
+    )
+
+    circuit_m, _ = moment_similarity(circuit, weight=1.0)
+    circuit_b, _ = block_similarity(circuit, weight=1.0, block_id=1)
+    remove_tags(circuit_m)
+
+    parallelized_circuit = parallelize(circuit)
+
+    print(parallelized_circuit)
+
+    # NOTE: depending on hardware, cirq produces differing, but unitary equivalent
+    # native circuits; in some cases, there is a PhZX gate with a negative phase
+    # which cannot be combined with others in the parallelization leading to a longer circuit
+    assert len(parallelized_circuit.moments) in (11, 13)
+
+    # this circuit should deterministically return all qubits to |0>
+    # let's check:
+    simulator = cirq.Simulator()
+    for _ in range(20):  # one in a million chance we miss an error
+        state_vector = simulator.simulate(parallelized_circuit).state_vector()
+        assert np.all(
+            np.isclose(
+                np.abs(state_vector),
+                np.concatenate((np.array([1]), np.zeros(2**4 - 1))),
+            )
+        )
+
+
+def test_nonunitary_error_gate():
+    qubits = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit(
+        cirq.H(qubits[0]),
+        cirq.CX(qubits[0], qubits[1]),
+        cirq.amplitude_damp(0.5).on(qubits[1]),
+        cirq.CX(qubits[1], qubits[0]),
+    )
+
+    parallelized_circuit = parallelize(circuit)
+
+    assert len(parallelized_circuit.moments) == 7
+
+
 RNG_STATE = np.random.RandomState(1902833)