Add a simple circuit fidelity analysis pass

Author: david-pl

src/bloqade/analysis/fidelity/__init__.py

@@ -0,0 +1 @@
+from .analysis import FidelityAnalysis as FidelityAnalysis

src/bloqade/analysis/fidelity/analysis.py

@@ -0,0 +1,40 @@
+from dataclasses import field
+
+from kirin import ir
+from kirin.lattice import EmptyLattice
+from kirin.analysis import Forward
+from kirin.interp.value import Successor
+from kirin.analysis.forward import ForwardFrame
+
+
+class FidelityAnalysis(Forward):
+    """
+    This analysis pass can be used to track the global addresses of qubits and wires.
+    """
+
+    keys = ["circuit.fidelity"]
+    lattice = EmptyLattice
+
+    # TODO: this should be a tuple[float, float] = (mean, max)
+    current_fidelity: float = field(init=False)
+    global_fidelity: float = 1.0
+
+    def initialize(self):
+        super().initialize()
+        self.current_fidelity = 1.0
+        return self
+
+    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))
+        )
+        return
+
+    def run_method(self, method: ir.Method, args: tuple[EmptyLattice, ...]):
+        return self.run_callable(method.code, (self.lattice.bottom(),) + args)

src/bloqade/noise/__init__.py

@@ -1 +1,2 @@
-from . import native as native
+# NOTE: just to register methods
+from . import native as native, fidelity as fidelity

src/bloqade/noise/fidelity.py

@@ -0,0 +1,30 @@
+from kirin import interp
+from kirin.lattice import EmptyLattice
+
+from bloqade.analysis.fidelity import FidelityAnalysis
+
+from .native import dialect
+from .native.stmts import PauliChannel, CZPauliChannel
+
+
+@dialect.register(key="circuit.fidelity")
+class FidelityMethodTable(interp.MethodTable):
+
+    @interp.impl(PauliChannel)
+    @interp.impl(CZPauliChannel)
+    def single_qubit_gate(
+        self,
+        interp: FidelityAnalysis,
+        frame: interp.Frame[EmptyLattice],
+        stmt: PauliChannel,
+    ):
+        ps, ps_ctrl = stmt.probabilities
+        p = sum(ps)
+
+        # NOTE: will be 0 if no ctrls
+        p_ctrl = sum(ps_ctrl)
+
+        # NOTE: fidelity is just the inverse probability of any noise to occur
+        fid = (1 - p) * (1 - p_ctrl)
+
+        interp.current_fidelity *= fid

test/analysis/fidelity/test_fidelity.py

@@ -0,0 +1,48 @@
+from bloqade import qasm2
+from bloqade.noise import native
+from bloqade.analysis.fidelity import FidelityAnalysis
+from bloqade.qasm2.passes.noise import NoisePass
+
+
+class NoiseTestModel(native.MoveNoiseModelABC):
+    def parallel_cz_errors(self, ctrls, qargs, rest):
+        return {(0.01, 0.01, 0.01, 0.01): ctrls + qargs + rest}
+
+
+@qasm2.main.add(native.dialect)
+def main():
+    q = qasm2.qreg(2)
+    qasm2.x(q[0])
+    return q
+
+
+main.print()
+
+fid_analysis = FidelityAnalysis(main.dialects)
+fid_analysis.run_analysis(main)
+
+assert fid_analysis.global_fidelity == fid_analysis.current_fidelity == 1
+
+
+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
+)
+
+model = NoiseTestModel()
+
+
+NoisePass(main.dialects, noise_model=model, gate_noise_params=noise_params)(main)
+
+
+main.print()
+
+fid_analysis = FidelityAnalysis(main.dialects)
+fid_analysis.run_analysis(main)
+
+print(fid_analysis.global_fidelity)
+print(fid_analysis.current_fidelity)

test/qasm2/passes/test_heuristic_noise.py

@@ -12,9 +12,7 @@
 
 
 class NoiseTestModel(native.MoveNoiseModelABC):
-
-    @classmethod
-    def parallel_cz_errors(cls, ctrls, qargs, rest):
+    def parallel_cz_errors(self, ctrls, qargs, rest):
         return {(0.01, 0.01, 0.01, 0.01): ctrls + qargs + rest}