Revert "remove tests for measurement/reset and pyqrack."

This reverts commit 3c21aae.

Author: david-pl

src/bloqade/cirq_utils/noise/model.py

@@ -57,8 +57,8 @@ class GeminiNoiseModelABC(cirq.NoiseModel, MoveNoiseModelABC):
 
     def __post_init__(self):
         if (
-                self.cz_paired_correlated_rates is None
-                and self.cz_paired_error_probabilities is None
+            self.cz_paired_correlated_rates is None
+            and self.cz_paired_error_probabilities is None
         ):
             # NOTE: no input, set to default value; weird setattr for frozen dataclass
             object.__setattr__(
@@ -67,8 +67,8 @@ def __post_init__(self):
                 _default_cz_paired_correlated_rates(),
             )
         elif (
-                self.cz_paired_correlated_rates is not None
-                and self.cz_paired_error_probabilities is None
+            self.cz_paired_correlated_rates is not None
+            and self.cz_paired_error_probabilities is None
         ):
 
             if self.cz_paired_correlated_rates.shape != (4, 4):
@@ -83,8 +83,8 @@ def __post_init__(self):
                 correlated_noise_array_to_dict(self.cz_paired_correlated_rates),
             )
         elif (
-                self.cz_paired_correlated_rates is not None
-                and self.cz_paired_error_probabilities is not None
+            self.cz_paired_correlated_rates is not None
+            and self.cz_paired_error_probabilities is not None
         ):
             raise ValueError(
                 "Received both `cz_paired_correlated_rates` and `cz_paired_correlated_rates` as input. This is ambiguous, please only set one."
@@ -93,7 +93,9 @@ def __post_init__(self):
     @staticmethod
     def validate_moments(moments: Iterable[cirq.Moment]):
         reset_family = cirq.GateFamily(gate=cirq.ResetChannel, ignore_global_phase=True)
-        allowed_target_gates: frozenset[cirq.GateFamily] = cirq.CZTargetGateset(additional_gates=[reset_family]).gates
+        allowed_target_gates: frozenset[cirq.GateFamily] = cirq.CZTargetGateset(
+            additional_gates=[reset_family]
+        ).gates
         # allowed_target_gates: frozenset[cirq.GateFamily] = cirq.CZTargetGateset().gates
 
         for moment in moments:
@@ -113,7 +115,7 @@ def validate_moments(moments: Iterable[cirq.Moment]):
                     )
 
     def parallel_cz_errors(
-            self, ctrls: list[int], qargs: list[int], rest: list[int]
+        self, ctrls: list[int], qargs: list[int], rest: list[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."
@@ -179,7 +181,7 @@ class GeminiOneZoneNoiseModel(GeminiNoiseModelABC):
     parallelize_circuit: bool = False
 
     def _single_qubit_moment_noise_ops(
-            self, moment: cirq.Moment, system_qubits: Sequence[cirq.Qid]
+        self, moment: cirq.Moment, system_qubits: Sequence[cirq.Qid]
     ) -> tuple[list, list]:
         """
         Helper function to determine the noise operations for a single qubit moment.
@@ -211,7 +213,7 @@ def _single_qubit_moment_noise_ops(
             op.qubits[0]
             for op in moment.operations
             if not (
-                    np.isclose(op.gate.x_exponent, 0) and np.isclose(op.gate.z_exponent, 0)
+                np.isclose(op.gate.x_exponent, 0) and np.isclose(op.gate.z_exponent, 0)
             )
         ]
 
@@ -247,9 +249,11 @@ def noisy_moment(self, moment, system_qubits):
             gate_noise_ops = []
         # Check if the moment contains 1-qubit gates or 2-qubit gates
         elif len(moment.operations[0].qubits) == 1:
-            if (isinstance(moment.operations[0].gate, cirq.ResetChannel)) or (
-                    cirq.is_measurement(moment.operations[0])) or (
-            isinstance(moment.operations[0].gate, cirq.BitFlipChannel)):
+            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:
@@ -303,7 +307,7 @@ def noisy_moment(self, moment, system_qubits):
         ]
 
     def noisy_moments(
-            self, moments: Iterable[cirq.Moment], system_qubits: Sequence[cirq.Qid]
+        self, moments: Iterable[cirq.Moment], system_qubits: Sequence[cirq.Qid]
     ) -> Sequence[cirq.OP_TREE]:
         """Adds possibly stateful noise to a series of moments.
 
@@ -321,18 +325,37 @@ 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) and (not cirq.is_measurement(op)) and (
-                not isinstance(op.gate, cirq.ResetChannel))])
+            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) and (not cirq.is_measurement(op))])
+            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))])
+            cirq.Moment(
+                [
+                    op
+                    for op in moment.operations
+                    if (cirq.is_measurement(op))
+                    or (isinstance(op.gate, cirq.ResetChannel))
+                ]
+            )
             for moment in moments
         ]
 
@@ -356,7 +379,9 @@ def count_remaining_cz_moments(moments_2q):
         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
+        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)
@@ -369,8 +394,15 @@ def count_remaining_cz_moments(moments_2q):
                 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)))
+            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)
@@ -412,9 +444,11 @@ def noisy_moment(self, moment, system_qubits):
             gate_noise_ops = []
         # Check if the moment contains 1-qubit gates or 2-qubit gates
         elif len(moment.operations[0].qubits) == 1:
-            if (isinstance(moment.operations[0].gate, cirq.ResetChannel)) or (
-                    cirq.is_measurement(moment.operations[0])) or (
-            isinstance(moment.operations[0].gate, cirq.BitFlipChannel)):
+            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(
@@ -484,7 +518,7 @@ def noisy_moment(self, moment, system_qubits):
 @dataclass(frozen=True)
 class GeminiTwoZoneNoiseModel(GeminiNoiseModelABC):
     def noisy_moments(
-            self, moments: Iterable[cirq.Moment], system_qubits: Sequence[cirq.Qid]
+        self, moments: Iterable[cirq.Moment], system_qubits: Sequence[cirq.Qid]
     ) -> Sequence[cirq.OP_TREE]:
         """Adds possibly stateful noise to a series of moments.
 
@@ -516,12 +550,12 @@ def noisy_moments(
                 [
                     moment
                     for moment in _two_zone_utils.get_move_error_channel_two_zoned(
-                    moments[i],
-                    prev_moment,
-                    np.array(self.mover_pauli_rates),
-                    np.array(self.sitter_pauli_rates),
-                    nqubs,
-                ).moments
+                        moments[i],
+                        prev_moment,
+                        np.array(self.mover_pauli_rates),
+                        np.array(self.sitter_pauli_rates),
+                        nqubs,
+                    ).moments
                     if len(moment) > 0
                 ]
             )
@@ -532,13 +566,13 @@ def noisy_moments(
                 [
                     moment
                     for moment in _two_zone_utils.get_gate_error_channel(
-                    moments[i],
-                    np.array(self.local_pauli_rates),
-                    np.array(self.global_pauli_rates),
-                    self.two_qubit_pauli,
-                    np.array(self.cz_unpaired_pauli_rates),
-                    nqubs,
-                ).moments
+                        moments[i],
+                        np.array(self.local_pauli_rates),
+                        np.array(self.global_pauli_rates),
+                        self.two_qubit_pauli,
+                        np.array(self.cz_unpaired_pauli_rates),
+                        nqubs,
+                    ).moments
                     if len(moment) > 0
                 ]
             )

test/cirq_utils/noise/test_noise_models.py

@@ -40,16 +40,16 @@ def create_ghz_circuit(qubits, measurements: bool = False):
     [
         (GeminiOneZoneNoiseModel(), None, False),
         (
-                GeminiOneZoneNoiseModelConflictGraphMoves(),
-                cirq.GridQubit.rect(rows=1, cols=2),
-                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
+            GeminiOneZoneNoiseModelConflictGraphMoves(),
+            cirq.GridQubit.rect(rows=1, cols=2),
+            True,
         ),
         (GeminiTwoZoneNoiseModel(), None, True),
     ],
@@ -76,22 +76,21 @@ def test_simple_model(model: cirq.NoiseModel, qubits, measurements: bool):
     dm = cirq_sim.simulate(noisy_circuit).final_density_matrix
     pops_cirq = np.real(np.diag(dm))
 
-    if not measurements:
-        kernel = load_circuit(noisy_circuit)
-        pyqrack_sim = StackMemorySimulator(
-            min_qubits=2, rng_state=np.random.default_rng(1234)
-        )
+    kernel = load_circuit(noisy_circuit)
+    pyqrack_sim = StackMemorySimulator(
+        min_qubits=2, rng_state=np.random.default_rng(1234)
+    )
 
-        pops_bloqade = [0.0] * 4
+    pops_bloqade = [0.0] * 4
 
-        nshots = 500
-        for _ in range(nshots):
-            ket = pyqrack_sim.state_vector(kernel)
-            for i in range(4):
-                pops_bloqade[i] += abs(ket[i]) ** 2 / nshots
+    nshots = 500
+    for _ in range(nshots):
+        ket = pyqrack_sim.state_vector(kernel)
+        for i in range(4):
+            pops_bloqade[i] += abs(ket[i]) ** 2 / nshots
 
     if measurements is True:
-        for pops in [pops_cirq]:
+        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)

test/cirq_utils/test_parallelize.py

@@ -27,7 +27,7 @@ def test1():
 
     circuit_m, _ = moment_similarity(circuit, weight=1.0)
     circuit_b, _ = block_similarity(circuit, weight=1.0, block_id=1)
-    circuit_m2 = remove_tags(circuit_m)
+    remove_tags(circuit_m)
     circuit2 = parallelize(circuit)
     assert len(circuit2.moments) == 7
 
@@ -49,16 +49,23 @@ def test_measurement_and_reset():
 
     circuit_m, _ = moment_similarity(circuit, weight=1.0)
     circuit_b, _ = block_similarity(circuit, weight=1.0, block_id=1)
-    circuit_m2 = remove_tags(circuit_m)
+    remove_tags(circuit_m)
 
     parallelized_circuit = parallelize(circuit)
 
+    assert len(parallelized_circuit.moments) == 11
+
     # 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)))))
+        assert np.all(
+            np.isclose(
+                np.abs(state_vector),
+                np.concatenate((np.array([1]), np.zeros(2**4 - 1))),
+            )
+        )
 
 
 def test_nonunitary_error_gate():