Introduce FidelityRange data class to store values

Author: david-pl

src/bloqade/analysis/fidelity/__init__.py

@@ -1,2 +1,5 @@
 from . import impls as impls
-from .analysis import FidelityAnalysis as FidelityAnalysis
+from .analysis import (
+    FidelityRange as FidelityRange,
+    FidelityAnalysis as FidelityAnalysis,
+)

src/bloqade/analysis/fidelity/analysis.py

@@ -3,6 +3,14 @@
 from ..address import Address, AddressReg, AddressAnalysis
 
 
+@dataclass
+class FidelityRange:
+    """Range of fidelity for a qubit as pair of (min, max) values"""
+
+    min: float
+    max: float
+
+
 @dataclass
 class FidelityAnalysis(AddressAnalysis):
     """
@@ -37,10 +45,10 @@ def main():
     keys = ("circuit.fidelity", "qubit.address")
     lattice = Address
 
-    gate_fidelities: list[list[float]] = field(init=False, default_factory=list)
+    gate_fidelities: list[FidelityRange] = field(init=False, default_factory=list)
     """Gate fidelities of each qubit as (min, max) pairs to provide a range"""
 
-    qubit_survival_fidelities: list[list[float]] = field(
+    qubit_survival_fidelities: list[FidelityRange] = field(
         init=False, default_factory=list
     )
     """Qubit survival fidelity given as (min, max) pairs"""
@@ -59,30 +67,34 @@ def extend_fidelities(self):
         self.extend_fidelity(self.gate_fidelities)
         self.extend_fidelity(self.qubit_survival_fidelities)
 
-    def extend_fidelity(self, fidelities: list[list[float]]):
+    def extend_fidelity(self, fidelities: list[FidelityRange]):
         n = self.qubit_count
-        fidelities.extend([[1.0, 1.0] for _ in range(n - len(fidelities))])
+        fidelities.extend([FidelityRange(1.0, 1.0) for _ in range(n - len(fidelities))])
 
     def reset_fidelities(self):
-        self.gate_fidelities = [[1.0, 1.0] for _ in range(self.qubit_count)]
-        self.qubit_survival_fidelities = [[1.0, 1.0] for _ in range(self.qubit_count)]
+        self.gate_fidelities = [
+            FidelityRange(1.0, 1.0) for _ in range(self.qubit_count)
+        ]
+        self.qubit_survival_fidelities = [
+            FidelityRange(1.0, 1.0) for _ in range(self.qubit_count)
+        ]
 
     @staticmethod
     def update_fidelities(
-        fidelities: list[list[float]], fidelity: float, addresses: AddressReg
+        fidelities: list[FidelityRange], fidelity: float, addresses: AddressReg
     ):
         """short-hand to update both (min, max) values"""
 
         for idx in addresses.data:
-            fidelities[idx][0] *= fidelity
-            fidelities[idx][1] *= fidelity
+            fidelities[idx].min *= fidelity
+            fidelities[idx].max *= fidelity
 
     def update_branched_fidelities(
         self,
-        fidelities: list[list[float]],
-        current_fidelities: list[list[float]],
-        then_fidelities: list[list[float]],
-        else_fidelities: list[list[float]],
+        fidelities: list[FidelityRange],
+        current_fidelities: list[FidelityRange],
+        then_fidelities: list[FidelityRange],
+        else_fidelities: list[FidelityRange],
     ):
         # NOTE: make sure they are all of the same length
         map(
@@ -94,8 +106,8 @@ def update_branched_fidelities(
         for fid, current_fid, then_fid, else_fid in zip(
             fidelities, current_fidelities, then_fidelities, else_fidelities
         ):
-            fid[0] = current_fid[0] * min(then_fid[0], else_fid[0])
-            fid[1] = current_fid[1] * max(then_fid[1], else_fid[1])
+            fid.min = current_fid.min * min(then_fid.min, else_fid.min)
+            fid.max = current_fid.max * max(then_fid.max, else_fid.max)
 
     def initialize(self):
         super().initialize()

test/analysis/fidelity/test_fidelity.py

@@ -4,7 +4,7 @@
 
 from bloqade import qasm2, squin
 from bloqade.qasm2 import noise
-from bloqade.analysis.fidelity import FidelityAnalysis
+from bloqade.analysis.fidelity import FidelityRange, FidelityAnalysis
 from bloqade.qasm2.passes.noise import NoisePass
 
 
@@ -216,11 +216,14 @@ def main():
     print(fid_analysis.gate_fidelities)
 
     assert len(fid_analysis.gate_fidelities) == 2
-    assert math.isclose(fid_analysis.gate_fidelities[0][0], 0.4)
-    assert math.isclose(fid_analysis.gate_fidelities[0][1], 0.4)
-    assert fid_analysis.gate_fidelities[1] == [1.0, 1.0]
+    assert math.isclose(fid_analysis.gate_fidelities[0].max, 0.4)
+    assert math.isclose(fid_analysis.gate_fidelities[0].min, 0.4)
+    assert fid_analysis.gate_fidelities[1] == FidelityRange(1.0, 1.0)
 
-    assert fid_analysis.qubit_survival_fidelities == [[1.0, 1.0], [0.9, 0.9]]
+    assert fid_analysis.qubit_survival_fidelities == [
+        FidelityRange(1.0, 1.0),
+        FidelityRange(0.9, 0.9),
+    ]
 
 
 def test_squin_if():
@@ -242,8 +245,14 @@ def main():
     fidelity_analysis = FidelityAnalysis(main.dialects)
     frame, _ = fidelity_analysis.run(main)
 
-    assert fidelity_analysis.gate_fidelities == [[0.9, 1.0], [0.8, 1.0]]
-    assert fidelity_analysis.qubit_survival_fidelities == [[0.85, 1.0], [0.75, 1.0]]
+    assert fidelity_analysis.gate_fidelities == [
+        FidelityRange(0.9, 1.0),
+        FidelityRange(0.8, 1.0),
+    ]
+    assert fidelity_analysis.qubit_survival_fidelities == [
+        FidelityRange(0.85, 1.0),
+        FidelityRange(0.75, 1.0),
+    ]
 
 
 def test_squin_for():
@@ -258,7 +267,7 @@ def main():
     frame, _ = fidelity_analysis.run(main)
 
     assert fidelity_analysis.gate_fidelities == [
-        [1.0 - i * 0.01, 1.0 - i * 0.01] for i in range(4)
+        FidelityRange(1.0 - i * 0.01, 1.0 - i * 0.01) for i in range(4)
     ]
 
 
@@ -302,19 +311,25 @@ def main():
     frame, _ = fidelity_analysis.run(main)
 
     assert fidelity_analysis.gate_fidelities == [
-        [0.4, 0.4],  # squin.single_qubit_pauli_channel(0.15, 0.2, 0.25, q[0])
-        [0.8, 0.8],  # squin.depolarize(0.2, q[1])
-        [1 - 12 * 0.01, 1 - 12 * 0.01],  # squin.two_qubit_pauli_channel(..., q[2])
-        [1 - 12 * 0.01, 1 - 12 * 0.01],  # squin.two_qubit_pauli_channel(..., q[3])
-        [0.88, 0.88],  # squin.depolarize2(0.15, q[4])
-        [0.88, 0.88],  # squin.depolarize2(0.15, q[5])
+        FidelityRange(
+            0.4, 0.4
+        ),  # squin.single_qubit_pauli_channel(0.15, 0.2, 0.25, q[0])
+        FidelityRange(0.8, 0.8),  # squin.depolarize(0.2, q[1])
+        FidelityRange(
+            1 - 12 * 0.01, 1 - 12 * 0.01
+        ),  # squin.two_qubit_pauli_channel(..., q[2])
+        FidelityRange(
+            1 - 12 * 0.01, 1 - 12 * 0.01
+        ),  # squin.two_qubit_pauli_channel(..., q[3])
+        FidelityRange(0.88, 0.88),  # squin.depolarize2(0.15, q[4])
+        FidelityRange(0.88, 0.88),  # squin.depolarize2(0.15, q[5])
     ]
 
     assert (
         fidelity_analysis.qubit_survival_fidelities
         == [
-            [0.9, 0.9],  # squin.qubit_loss(0.1, q[0])
-            [1.0, 1.0],
+            FidelityRange(0.9, 0.9),  # squin.qubit_loss(0.1, q[0])
+            FidelityRange(1.0, 1.0),
         ]
-        + [[0.87, 0.87]] * 4
+        + [FidelityRange(0.87, 0.87)] * 4
     )  # squin.correlated_qubit_loss