Track atom loss per qubit rather than globally

Author: david-pl

src/bloqade/analysis/fidelity/analysis.py

@@ -1,3 +1,4 @@
+from typing import Any
 from dataclasses import field
 
 from kirin import ir
@@ -6,6 +7,8 @@
 from kirin.interp.value import Successor
 from kirin.analysis.forward import ForwardFrame
 
+from ..address import AddressAnalysis
+
 
 class FidelityAnalysis(Forward):
     """
@@ -25,23 +28,42 @@ class FidelityAnalysis(Forward):
     """
     The probability that all atoms survive for the entirety of the analysed program. Decreases whenever an atomic loss channel is encountered.
     """
-    atom_survival_probability: float = 1.0
+    atom_survival_probability: list[float] = field(init=False)
+
+    _current_atom_survival_probability: list[float] = field(init=False)
 
-    _current_atom_survival_probability: float = field(init=False)
+    addr_frame: ForwardFrame = field(init=False)
 
     def initialize(self):
         super().initialize()
         self._current_gate_fidelity = 1.0
-        self._current_atom_survival_probability = 1.0
+        self._current_atom_survival_probability = [
+            1.0 for _ in range(len(self.atom_survival_probability))
+        ]
         return self
 
     def posthook_succ(self, frame: ForwardFrame, succ: Successor):
         self.gate_fidelity *= self._current_gate_fidelity
-        self.atom_survival_probability *= self._current_atom_survival_probability
+        for i, _current_survival in enumerate(self._current_atom_survival_probability):
+            self.atom_survival_probability[i] *= _current_survival
 
     def eval_stmt_fallback(self, frame: ForwardFrame, stmt: ir.Statement):
         # NOTE: default is to conserve fidelity, so do nothing here
         return
 
     def run_method(self, method: ir.Method, args: tuple[EmptyLattice, ...]):
         return self.run_callable(method.code, (self.lattice.bottom(),) + args)
+
+    def run_analysis(
+        self, method: ir.Method, args: tuple | None = None, *, no_raise: bool = True
+    ) -> tuple[ForwardFrame, Any]:
+        self._run_address_analysis(method, no_raise=no_raise)
+        return super().run_analysis(method, args, no_raise=no_raise)
+
+    def _run_address_analysis(self, method: ir.Method, no_raise: bool):
+        addr_analysis = AddressAnalysis(self.dialects)
+        addr_frame, _ = addr_analysis.run_analysis(method=method, no_raise=no_raise)
+        self.addr_frame = addr_frame
+
+        # NOTE: make sure we have as many probabilities as we have addresses
+        self.atom_survival_probability = [1.0] * addr_analysis.qubit_count

src/bloqade/noise/fidelity.py

@@ -5,6 +5,7 @@
 
 from .native import dialect as native
 from .native.stmts import PauliChannel, CZPauliChannel, AtomLossChannel
+from ..analysis.address import AddressQubit, AddressTuple
 
 
 @native.register(key="circuit.fidelity")
@@ -40,4 +41,11 @@ def atom_loss(
         frame: interp.Frame[EmptyLattice],
         stmt: AtomLossChannel,
     ):
-        interp._current_atom_survival_probability *= 1 - stmt.prob
+        # NOTE: since AtomLossChannel acts on IList[Qubit], we know the assigned address is a tuple
+        addresses: AddressTuple = interp.addr_frame.get(stmt.qargs)
+
+        # NOTE: get the corresponding index and reduce survival probability accordingly
+        for qbit_address in addresses.data:
+            assert isinstance(qbit_address, AddressQubit)
+            index = qbit_address.data
+            interp._current_atom_survival_probability[index] *= 1 - stmt.prob

test/analysis/fidelity/test_fidelity.py

@@ -59,8 +59,9 @@ def main():
         == (1 - p_noise)
     )
 
-    assert 0.9 < fid_analysis.atom_survival_probability < 1
-    assert fid_analysis.atom_survival_probability == 1 - noise_params.local_loss_prob
+    assert 0.9 < fid_analysis.atom_survival_probability[0] < 1
+    assert fid_analysis.atom_survival_probability[0] == 1 - noise_params.local_loss_prob
+    assert fid_analysis.atom_survival_probability[1] == 1
 
 
 def test_c_noise():
@@ -108,10 +109,10 @@ def main():
     assert fid_analysis.gate_fidelity == fid_analysis._current_gate_fidelity
     assert math.isclose(fid_cz, fid_analysis.gate_fidelity, abs_tol=1e-14)
 
-    assert 0.9 < fid_analysis.atom_survival_probability < 1
-    assert fid_analysis.atom_survival_probability == (
+    assert 0.9 < fid_analysis.atom_survival_probability[0] < 1
+    assert fid_analysis.atom_survival_probability[0] == (
         1 - noise_params.cz_gate_loss_prob
-    ) ** 2 * (1 - p_loss)
+    ) * (1 - p_loss)
 
 
 @pytest.mark.xfail
@@ -169,8 +170,8 @@ def main_if():
     assert 0 < fid_if_analysis.gate_fidelity == fid_analysis.gate_fidelity < 1
     assert (
         0
-        < fid_if_analysis.atom_survival_probability
-        == fid_analysis.atom_survival_probability
+        < fid_if_analysis.atom_survival_probability[0]
+        == fid_analysis.atom_survival_probability[0]
         < 1
     )
 
@@ -239,7 +240,7 @@ def main_for():
     assert 0 < fid_for_analysis.gate_fidelity == fid_analysis.gate_fidelity < 1
     assert (
         0
-        < fid_for_analysis.atom_survival_probability
-        == fid_analysis.atom_survival_probability
+        < fid_for_analysis.atom_survival_probability[0]
+        == fid_analysis.atom_survival_probability[0]
         < 1
     )