Allow measurements and resets into Gemini noise model.

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."
@@ -246,14 +249,22 @@ def noisy_moment(self, moment, system_qubits):
         original_moment = moment
 
         # Check if the moment is empty
-        if len(moment.operations) == 0:
+        if len(moment.operations) == 0 or cirq.is_measurement(moment.operations[0]):
             move_noise_ops = []
             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, move_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))
+            ):
+                move_noise_ops = []
+                gate_noise_ops = []
+            else:
+                gate_noise_ops, move_noise_ops = self._single_qubit_moment_noise_ops(
+                    moment, system_qubits
+                )
         elif len(moment.operations[0].qubits) == 2:
             control_qubits = [op.qubits[0] for op in moment.operations]
             target_qubits = [op.qubits[1] for op in moment.operations]
@@ -319,20 +330,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
+        ]
+
+        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)
+        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 +443,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)