Fix qasm2 impls

Author: david-pl

src/bloqade/qasm2/dialects/noise/fidelity.py

@@ -1,6 +1,7 @@
 from kirin import interp
-from kirin.lattice import EmptyLattice
+from kirin.analysis import ForwardFrame
 
+from bloqade.analysis.address import Address, AddressReg
 from bloqade.analysis.fidelity import FidelityAnalysis
 
 from .stmts import PauliChannel, CZPauliChannel, AtomLossChannel
@@ -11,37 +12,68 @@
 class FidelityMethodTable(interp.MethodTable):
 
     @interp.impl(PauliChannel)
-    @interp.impl(CZPauliChannel)
     def pauli_channel(
         self,
-        interp: FidelityAnalysis,
-        frame: interp.Frame[EmptyLattice],
-        stmt: PauliChannel | CZPauliChannel,
+        interp_: FidelityAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: PauliChannel,
+    ):
+        (ps,) = stmt.probabilities
+        fidelity = 1 - sum(ps)
+
+        addresses = frame.get(stmt.qargs)
+
+        if not isinstance(addresses, AddressReg):
+            return ()
+
+        interp_.update_fidelities(interp_.gate_fidelities, fidelity, addresses)
+
+        return ()
+
+    @interp.impl(CZPauliChannel)
+    def cz_pauli_channel(
+        self,
+        interp_: FidelityAnalysis,
+        frame: ForwardFrame[Address],
+        stmt: CZPauliChannel,
     ):
-        probs = stmt.probabilities
-        try:
-            ps, ps_ctrl = probs
-        except ValueError:
-            (ps,) = probs
-            ps_ctrl = ()
+        ps_ctrl, ps_target = stmt.probabilities
+
+        fidelity_ctrl = 1 - sum(ps_ctrl)
+        fidelity_target = 1 - sum(ps_target)
+
+        addresses_ctrl = frame.get(stmt.ctrls)
+        addresses_target = frame.get(stmt.qargs)
 
-        p = sum(ps)
-        p_ctrl = sum(ps_ctrl)
+        if not isinstance(addresses_ctrl, AddressReg) or not isinstance(
+            addresses_target, AddressReg
+        ):
+            return ()
 
-        # NOTE: fidelity is just the inverse probability of any noise to occur
-        fid = (1 - p) * (1 - p_ctrl)
+        interp_.update_fidelities(
+            interp_.gate_fidelities, fidelity_ctrl, addresses_ctrl
+        )
+        interp_.update_fidelities(
+            interp_.gate_fidelities, fidelity_target, addresses_target
+        )
 
-        interp.gate_fidelity *= fid
+        return ()
 
     @interp.impl(AtomLossChannel)
     def atom_loss(
         self,
-        interp: FidelityAnalysis,
-        frame: interp.Frame[EmptyLattice],
+        interp_: FidelityAnalysis,
+        frame: ForwardFrame[Address],
         stmt: AtomLossChannel,
     ):
-        # NOTE: since AtomLossChannel acts on IList[Qubit], we know the assigned address is a tuple
-        addresses = interp.addr_frame.get(stmt.qargs)
-        # NOTE: get the corresponding index and reduce survival probability accordingly
-        for index in addresses.data:
-            interp.atom_survival_probability[index] *= 1 - stmt.prob
+        addresses = frame.get(stmt.qargs)
+
+        if not isinstance(addresses, AddressReg):
+            return ()
+
+        fidelity = 1 - stmt.prob
+        interp_.update_fidelities(
+            interp_.qubit_survival_fidelities, fidelity, addresses
+        )
+
+        return ()

test/analysis/fidelity/test_fidelity.py

@@ -21,9 +21,11 @@ def main():
     fid_analysis = FidelityAnalysis(main.dialects)
     fid_analysis.run(main)
 
-    assert fid_analysis.gate_fidelity == fid_analysis._current_gate_fidelity == 1
-    assert fid_analysis.atom_survival_probability[0] == 1 - p_loss
-    assert fid_analysis.atom_survival_probability[1] == 1
+    assert fid_analysis.gate_fidelities == [FidelityRange(1.0, 1.0)] * 2
+    assert fid_analysis.qubit_survival_fidelities == [
+        FidelityRange(1 - p_loss, 1 - p_loss),
+        FidelityRange(1.0, 1.0),
+    ]
 
 
 def test_cz_noise():
@@ -51,9 +53,12 @@ def main():
     fid_analysis = FidelityAnalysis(main.dialects)
     fid_analysis.run(main)
 
-    expected_fidelity = (1 - 3 * p_ch) ** 2
+    expected_fidelity = 1 - 3 * p_ch
 
-    assert math.isclose(fid_analysis.gate_fidelity, expected_fidelity)
+    assert (
+        fid_analysis.gate_fidelities
+        == [FidelityRange(expected_fidelity, expected_fidelity)] * 2
+    )
 
 
 def test_single_qubit_noise():
@@ -72,7 +77,10 @@ def main():
 
     expected_fidelity = 1 - 3 * p_ch
 
-    assert math.isclose(fid_analysis.gate_fidelity, expected_fidelity)
+    assert fid_analysis.gate_fidelities == [
+        FidelityRange(expected_fidelity, expected_fidelity),
+        FidelityRange(1.0, 1.0),
+    ]
 
 
 class NoiseTestModel(noise.MoveNoiseModelABC):
@@ -113,6 +121,7 @@ def main_if():
 
     model = NoiseTestModel(
         global_loss_prob=p_loss,
+        local_loss_prob=p_loss,
         global_px=px,
         global_py=py,
         global_pz=pz,
@@ -122,18 +131,21 @@ def main_if():
     fid_analysis = FidelityAnalysis(main.dialects)
     fid_analysis.run(main)
 
-    model = NoiseTestModel()
     NoisePass(main_if.dialects, noise_model=model)(main_if)
     fid_if_analysis = FidelityAnalysis(main_if.dialects)
     fid_if_analysis.run(main_if)
 
-    assert 0 < fid_if_analysis.gate_fidelity == fid_analysis.gate_fidelity < 1
-    assert (
-        0
-        < fid_if_analysis.atom_survival_probability[0]
-        == fid_analysis.atom_survival_probability[0]
-        < 1
-    )
+    main.print()
+    main_if.print()
+
+    fidelity_if = fid_if_analysis.gate_fidelities[0]
+    fidelity = fid_analysis.gate_fidelities[0]
+
+    survival = fid_analysis.qubit_survival_fidelities[0]
+    survival_if = fid_if_analysis.qubit_survival_fidelities[0]
+
+    assert 0 < fidelity_if.min == fidelity.min == fidelity.max < fidelity_if.max < 1
+    assert 0 < survival_if.min == survival.min == survival.max < survival_if.max < 1
 
 
 def test_for():
@@ -184,21 +196,20 @@ def main_for():
     fid_analysis = FidelityAnalysis(main.dialects)
     fid_analysis.run(main)
 
-    model = NoiseTestModel()
     NoisePass(main_for.dialects, noise_model=model)(main_for)
 
     main_for.print()
 
     fid_for_analysis = FidelityAnalysis(main_for.dialects)
     fid_for_analysis.run(main_for)
 
-    assert 0 < fid_for_analysis.gate_fidelity == fid_analysis.gate_fidelity < 1
-    assert (
-        0
-        < fid_for_analysis.atom_survival_probability[0]
-        == fid_analysis.atom_survival_probability[0]
-        < 1
-    )
+    fid = fid_analysis.gate_fidelities[0]
+    fid_for = fid_for_analysis.gate_fidelities[0]
+    survival = fid_analysis.qubit_survival_fidelities[0]
+    survival_for = fid_for_analysis.qubit_survival_fidelities[0]
+
+    assert 0 < fid.min == fid.max == fid_for.min == fid_for.max < 1
+    assert 0 < survival.min == survival.max == survival_for.min == survival_for.max < 1
 
 
 def test_stdlib_call():