Fidelity calculation for IfElse

Author: david-pl

src/bloqade/analysis/fidelity/analysis.py

@@ -18,6 +18,7 @@ class FidelityAnalysis(Forward):
     # TODO: this should be a tuple[float, float] = (mean, max)
     current_fidelity: float = field(init=False)
     global_fidelity: float = 1.0
+    # TODO: atom loss
 
     def initialize(self):
         super().initialize()
@@ -28,12 +29,12 @@ def posthook_succ(self, frame: ForwardFrame, succ: Successor):
         self.global_fidelity *= self.current_fidelity
 
     def eval_stmt_fallback(self, frame: ForwardFrame, stmt: ir.Statement):
-        print(
-            "no implementation for stmt "
-            + stmt.print_str(end="")
-            + " from "
-            + str(type(self))
-        )
+        # print(
+        #     "no implementation for stmt "
+        #     + stmt.print_str(end="")
+        #     + " from "
+        #     + str(type(self))
+        # )
         return
 
     def run_method(self, method: ir.Method, args: tuple[EmptyLattice, ...]):

src/bloqade/noise/fidelity.py

@@ -1,22 +1,24 @@
 from kirin import interp
 from kirin.lattice import EmptyLattice
+from kirin.dialects.scf import dialect as scf
+from kirin.dialects.scf.stmts import IfElse
 
 from bloqade.analysis.fidelity import FidelityAnalysis
 
-from .native import dialect
+from .native import dialect as native
 from .native.stmts import PauliChannel, CZPauliChannel
 
 
-@dialect.register(key="circuit.fidelity")
+@native.register(key="circuit.fidelity")
 class FidelityMethodTable(interp.MethodTable):
 
     @interp.impl(PauliChannel)
     @interp.impl(CZPauliChannel)
-    def single_qubit_gate(
+    def pauli_channel(
         self,
         interp: FidelityAnalysis,
         frame: interp.Frame[EmptyLattice],
-        stmt: PauliChannel,
+        stmt: PauliChannel | CZPauliChannel,
     ):
         probs = stmt.probabilities
         try:
@@ -32,3 +34,43 @@ def single_qubit_gate(
         fid = (1 - p) * (1 - p_ctrl)
 
         interp.current_fidelity *= fid
+
+
+@scf.register(key="circuit.fidelity")
+class ScfFidelityMethodTable(FidelityMethodTable):
+
+    @interp.impl(IfElse)
+    def if_else(
+        self,
+        interp: FidelityAnalysis,
+        frame: interp.Frame[EmptyLattice],
+        stmt: IfElse,
+    ):
+        # NOTE: store current fidelity for later
+        current_fidelity = interp.current_fidelity
+
+        for s in stmt.then_body.stmts():
+            stmt_impl = interp.lookup_registry(frame=frame, stmt=s)
+            if stmt_impl is None:
+                continue
+
+            stmt_impl(interp=interp, frame=frame, stmt=s)
+
+        then_fidelity = interp.current_fidelity
+
+        # NOTE: reset fidelity of interp to check if the else body results in a worse fidelity
+        interp.current_fidelity = current_fidelity
+
+        for s in stmt.else_body.stmts():
+            stmt_impl = interp.lookup_registry(frame=frame, stmt=s)
+            if stmt_impl is None:
+                continue
+
+            stmt_impl(interp=interp, frame=frame, stmt=s)
+
+        else_fidelity = interp.current_fidelity
+
+        if then_fidelity < else_fidelity:
+            interp.current_fidelity = then_fidelity
+        else:
+            interp.current_fidelity = else_fidelity

test/analysis/fidelity/test_fidelity.py

@@ -11,44 +11,98 @@ def parallel_cz_errors(self, ctrls, qargs, rest):
         return {(0.01, 0.01, 0.01, 0.01): ctrls + qargs + rest}
 
 
-@noise_main
-def main():
-    q = qasm2.qreg(2)
-    qasm2.x(q[0])
-    return q
+def test_basic_noise():
 
+    @noise_main
+    def main():
+        q = qasm2.qreg(2)
+        qasm2.x(q[0])
+        return q
 
-main.print()
+    main.print()
 
-fid_analysis = FidelityAnalysis(main.dialects)
-fid_analysis.run_analysis(main)
+    fid_analysis = FidelityAnalysis(main.dialects)
+    fid_analysis.run_analysis(main)
 
-assert fid_analysis.global_fidelity == fid_analysis.current_fidelity == 1
+    assert fid_analysis.global_fidelity == fid_analysis.current_fidelity == 1
 
+    px = 0.01
+    py = 0.01
+    pz = 0.01
+    p_loss = 0.01
 
-px = 0.01
-py = 0.01
-pz = 0.01
-p_loss = 0.01
+    noise_params = native.GateNoiseParams(
+        global_loss_prob=p_loss,
+        global_px=px,
+        global_py=py,
+        global_pz=pz,
+        local_px=0.002,
+    )
 
-noise_params = native.GateNoiseParams(
-    global_loss_prob=p_loss,
-    global_px=px,
-    global_py=py,
-    global_pz=pz,
-    local_px=0.002,
-)
+    model = NoiseTestModel()
 
-model = NoiseTestModel()
+    NoisePass(main.dialects, noise_model=model, gate_noise_params=noise_params)(main)
 
+    main.print()
 
-NoisePass(main.dialects, noise_model=model, gate_noise_params=noise_params)(main)
+    fid_analysis = FidelityAnalysis(main.dialects)
+    fid_analysis.run_analysis(main, no_raise=False)
 
+    p_noise = noise_params.local_px + noise_params.local_py + noise_params.local_pz
+    assert (
+        fid_analysis.global_fidelity == fid_analysis.current_fidelity == (1 - p_noise)
+    )
 
-main.print()
 
-fid_analysis = FidelityAnalysis(main.dialects)
-fid_analysis.run_analysis(main, no_raise=False)
+def test_if():
 
-p_noise = noise_params.local_px + noise_params.local_py + noise_params.local_pz
-assert fid_analysis.global_fidelity == fid_analysis.current_fidelity == (1 - p_noise)
+    @noise_main
+    def main():
+        q = qasm2.qreg(1)
+        c = qasm2.creg(1)
+        qasm2.h(q[0])
+        qasm2.measure(q, c)
+        qasm2.x(q[0])
+        qasm2.measure(q, c)
+
+        return c
+
+    @noise_main
+    def main_if():
+        q = qasm2.qreg(1)
+        c = qasm2.creg(1)
+        qasm2.h(q[0])
+        qasm2.measure(q, c)
+
+        if c[0] == 0:
+            qasm2.x(q[0])
+
+        qasm2.measure(q, c)
+        return c
+
+    px = 0.01
+    py = 0.01
+    pz = 0.01
+    p_loss = 0.01
+
+    noise_params = native.GateNoiseParams(
+        global_loss_prob=p_loss,
+        global_px=px,
+        global_py=py,
+        global_pz=pz,
+        local_px=0.002,
+    )
+
+    model = NoiseTestModel()
+    NoisePass(main.dialects, noise_model=model, gate_noise_params=noise_params)(main)
+    fid_analysis = FidelityAnalysis(main.dialects)
+    fid_analysis.run_analysis(main, no_raise=False)
+
+    model = NoiseTestModel()
+    NoisePass(main_if.dialects, noise_model=model, gate_noise_params=noise_params)(
+        main_if
+    )
+    fid_if_analysis = FidelityAnalysis(main_if.dialects)
+    fid_if_analysis.run_analysis(main_if, no_raise=False)
+
+    assert 0 < fid_if_analysis.global_fidelity == fid_analysis.global_fidelity < 1