Share check method between classes and simplify test

Author: david-pl

src/bloqade/noise/native/stmts.py

@@ -1,3 +1,5 @@
+from typing import Tuple
+
 from kirin import ir, types, lowering
 from kirin.decl import info, statement
 from kirin.dialects import ilist
@@ -7,30 +9,43 @@
 from ._dialect import dialect
 
 
-@statement(dialect=dialect)
-class PauliChannel(ir.Statement):
-
+@statement
+class NativeNoiseStmt(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
 
+    @property
+    def probabilities(self) -> Tuple[Tuple[float, ...], ...]: ...
+
+    def check(self):
+        for probs in self.probabilities:
+            self.check_probability(sum(probs))
+            for p in probs:
+                self.check_probability(p)
+
+    def check_probability(self, p: float):
+        if not 0 <= p <= 1:
+            raise ValueError(
+                f"Invalid noise probability encountered in {type(self).__name__}: {p}"
+            )
+
+
+@statement(dialect=dialect)
+class PauliChannel(NativeNoiseStmt):
     px: float = info.attribute(types.Float)
     py: float = info.attribute(types.Float)
     pz: float = info.attribute(types.Float)
     qargs: ir.SSAValue = info.argument(ilist.IListType[QubitType])
 
-    def check(self):
-        probs = (self.px, self.py, self.pz)
-        if not all(0 <= p <= 1 for p in probs) or not 0 <= sum(probs) <= 1:
-            raise ValueError(f"Invalid Pauli error probabilities (px, py, pz): {probs}")
+    @property
+    def probabilities(self) -> Tuple[Tuple[float, ...], ...]:
+        return ((self.px, self.py, self.pz),)
 
 
 NumQubits = types.TypeVar("NumQubits")
 
 
 @statement(dialect=dialect)
-class CZPauliChannel(ir.Statement):
-
-    traits = frozenset({lowering.FromPythonCall()})
-
+class CZPauliChannel(NativeNoiseStmt):
     paired: bool = info.attribute(types.Bool)
     px_ctrl: float = info.attribute(types.Float)
     py_ctrl: float = info.attribute(types.Float)
@@ -41,32 +56,19 @@ class CZPauliChannel(ir.Statement):
     ctrls: ir.SSAValue = info.argument(ilist.IListType[QubitType, NumQubits])
     qargs: ir.SSAValue = info.argument(ilist.IListType[QubitType, NumQubits])
 
-    def check(self):
-        probs_ctrl = (self.px_ctrl, self.py_ctrl, self.pz_ctrl)
-
-        def check_prob(p: float) -> bool:
-            return 0 <= p <= 1
-
-        if not map(check_prob, probs_ctrl) or not check_prob(sum(probs_ctrl)):
-            raise ValueError(
-                f"Invalid control probabilities for CZ Pauli channel (px_ctrl, py_ctrl, pz_ctrl): {probs_ctrl}"
-            )
-
-        probs_qarg = (self.px_qarg, self.py_qarg, self.pz_qarg)
-        if not map(check_prob, probs_qarg) or not check_prob(sum(probs_qarg)):
-            raise ValueError(
-                f"Invalid probabilities for CZ Pauli channel (px_qarg, py_qarg, pz_qarg): {probs_qarg}"
-            )
+    @property
+    def probabilities(self) -> Tuple[Tuple[float, ...], ...]:
+        return (
+            (self.px_ctrl, self.py_ctrl, self.pz_ctrl),
+            (self.px_qarg, self.py_qarg, self.pz_qarg),
+        )
 
 
 @statement(dialect=dialect)
-class AtomLossChannel(ir.Statement):
-
-    traits = frozenset({lowering.FromPythonCall()})
-
+class AtomLossChannel(NativeNoiseStmt):
     prob: float = info.attribute(types.Float)
     qargs: ir.SSAValue = info.argument(ilist.IListType[QubitType])
 
-    def check(self):
-        if not 0 <= self.prob <= 1:
-            raise ValueError(f"Invalid atom loss probability {self.prob}")
+    @property
+    def probabilities(self) -> Tuple[Tuple[float, ...], ...]:
+        return ((self.prob,),)

test/pyqrack/runtime/noise/native/test_loss.py

@@ -1,17 +1,13 @@
 from typing import Literal
-import textwrap
 from unittest.mock import Mock
 
-import pytest
 from kirin import ir
 from kirin.dialects import ilist
 
 from bloqade import qasm2
 from bloqade.noise import native
 from bloqade.pyqrack import PyQrackQubit, PyQrackInterpreter, reg
 from bloqade.pyqrack.base import MockMemory
-from bloqade.qasm2.passes import QASM2Py, NoisePass, QASM2Fold
-from bloqade.qasm2.parse.lowering import QASM2
 
 simulation = qasm2.extended.add(native)
 
@@ -47,97 +43,3 @@ def test_atom_loss(c: qasm2.CReg):
     assert result[0].state is reg.QubitState.Lost
     assert result[1].state is reg.QubitState.Active
     assert input[0] is reg.Measurement.One
-
-
-@pytest.mark.xfail
-def test_noise_probs():
-    test_qasm = textwrap.dedent(
-        """
-    OPENQASM 2.0;
-    include "qelib1.inc";
-
-    // Qubits: [q_0, q_1, q_2, q_3, q_4, q_5]
-    qreg q[6];
-
-
-    u3(pi*0.9999896015,pi*1.8867094803,pi*0.1132905197) q[2];
-    u3(pi*1.499959526,pi*1.2634437582,pi*0.7365562418) q[3];
-    u3(pi*1.4998447568,pi*1.8205928898,pi*0.1794071102) q[4];
-    u3(pi*1.4998052589,pi*1.5780611154,pi*0.4219388846) q[5];
-    u3(pi*0.4920440401,pi*1.287644074,pi*0.712355926) q[0];
-    u3(pi*1.0012473155,pi*1.3019213156,pi*0.6980786844) q[1];
-    cz q[1],q[2];
-    cz q[1],q[2];
-    cz q[2],q[3];
-    cz q[4],q[5];
-    cz q[2],q[3];
-    cz q[4],q[5];
-    cz q[2],q[3];
-    cz q[4],q[5];
-    cz q[2],q[3];
-    cz q[4],q[5];
-    cz q[2],q[3];
-    cz q[4],q[5];
-    u3(pi*1.0,pi*1.5687764466,pi*0.4312235534) q[2];
-    u3(pi*0.5,0,pi*1.7365086077) q[3];
-    u3(pi*0.5,pi*1.0,pi*0.6112880576) q[4];
-    u3(pi*0.1388474164,pi*1.7687898606,pi*1.2425564668) q[5];
-    """
-    )
-
-    entry = QASM2(qasm2.main.add(qasm2.inline_)).loads(test_qasm, "entry", returns="q")
-    QASM2Py(entry.dialects)(entry)
-    entry = entry.similar(qasm2.extended.add(native))
-    QASM2Fold(entry.dialects).fixpoint(entry)
-
-    # Noise parameters
-    gate_noise_value = 1e-3
-    move_noise_value = 0.5
-
-    gate_noise_params = native.GateNoiseParams(
-        local_px=gate_noise_value,
-        local_py=gate_noise_value,
-        local_pz=gate_noise_value,
-        local_loss_prob=gate_noise_value,
-        #
-        global_px=gate_noise_value,
-        global_py=gate_noise_value,
-        global_pz=gate_noise_value,
-        global_loss_prob=gate_noise_value,
-        #
-        cz_paired_gate_px=gate_noise_value,
-        cz_paired_gate_py=gate_noise_value,
-        cz_paired_gate_pz=gate_noise_value,
-        cz_gate_loss_prob=gate_noise_value,
-        #
-        cz_unpaired_gate_px=gate_noise_value,
-        cz_unpaired_gate_py=gate_noise_value,
-        cz_unpaired_gate_pz=gate_noise_value,
-        cz_unpaired_loss_prob=gate_noise_value,
-    )
-
-    move_noise_params = native.model.MoveNoiseParams(
-        idle_px_rate=move_noise_value,
-        idle_py_rate=move_noise_value,
-        idle_pz_rate=move_noise_value,
-        idle_loss_rate=move_noise_value,
-        move_px_rate=move_noise_value,
-        move_py_rate=move_noise_value,
-        move_pz_rate=move_noise_value,
-        move_loss_rate=move_noise_value,
-        #
-        pick_px=move_noise_value,
-        pick_py=move_noise_value,
-        pick_pz=move_noise_value,
-        pick_loss_prob=move_noise_value,
-        #
-        move_speed=5e-1,  # default 5e-1
-        storage_spacing=4.0,  # default 4.0
-    )
-
-    with pytest.raises(ir.ValidationError):
-        NoisePass(
-            entry.dialects,
-            gate_noise_params=gate_noise_params,
-            noise_model=native.TwoRowZoneModel(params=move_noise_params),
-        )(entry)

test/pyqrack/runtime/noise/native/test_pauli.py

@@ -1,5 +1,6 @@
 from unittest.mock import Mock, call
 
+import pytest
 from kirin import ir
 
 from bloqade import qasm2
@@ -41,6 +42,23 @@ def test_atom_loss():
     sim_reg.assert_has_calls([call.y(0)])
 
 
+@pytest.mark.xfail
+def test_pauli_probs_check():
+    @simulation
+    def test_atom_loss():
+        q = qasm2.qreg(2)
+        native.pauli_channel(
+            [q[0]],
+            px=0.1,
+            py=0.4,
+            pz=1.3,
+        )
+        return q
+
+    with pytest.raises(ir.ValidationError):
+        test_atom_loss.verify()
+
+
 def test_cz_pauli_channel_false():
     @simulation
     def test_atom_loss():
@@ -126,5 +144,6 @@ def test_atom_loss():
 
 if __name__ == "__main__":
     test_pauli_channel()
+    test_pauli_probs_check()
     test_cz_pauli_channel_false()
     test_cz_pauli_channel_true()