bloqade-circuit 0.6.2__py3-none-any.whl → 0.9.1__py3-none-any.whl

bloqade/pyqrack/squin/noise/native.py

@@ -1,72 +1,138 @@
-import random
-import typing
-from dataclasses import dataclass
-
 from kirin import interp
-from kirin.dialects import ilist
 
-from bloqade.pyqrack import QubitState, PyQrackQubit, PyQrackInterpreter
-from bloqade.squin.noise.stmts import QubitLoss, StochasticUnitaryChannel
+from bloqade.pyqrack import PyQrackQubit, PyQrackInterpreter
+from bloqade.squin.noise.stmts import (
+    QubitLoss,
+    Depolarize,
+    Depolarize2,
+    CorrelatedQubitLoss,
+    TwoQubitPauliChannel,
+    SingleQubitPauliChannel,
+)
 from bloqade.squin.noise._dialect import dialect as squin_noise_dialect
 
-from ..runtime import OperatorRuntimeABC
-
-
-@dataclass(frozen=True)
-class StochasticUnitaryChannelRuntime(OperatorRuntimeABC):
-    operators: ilist.IList[OperatorRuntimeABC, typing.Any]
-    probabilities: ilist.IList[float, typing.Any]
-
-    @property
-    def n_sites(self) -> int:
-        n = self.operators[0].n_sites
-        for op in self.operators[1:]:
-            assert (
-                op.n_sites == n
-            ), "Encountered a stochastic unitary channel with operators of different size!"
-        return n
-
-    def apply(self, *qubits: PyQrackQubit, adjoint: bool = False) -> None:
-        # NOTE: probabilities don't necessarily sum to 1; could be no noise event should occur
-        p_no_op = 1 - sum(self.probabilities)
-        if random.uniform(0.0, 1.0) < p_no_op:
-            return
-
-        selected_ops = random.choices(self.operators, weights=self.probabilities)
-        for op in selected_ops:
-            op.apply(*qubits, adjoint=adjoint)
-
 
-@dataclass(frozen=True)
-class QubitLossRuntime(OperatorRuntimeABC):
-    p: float
-
-    @property
-    def n_sites(self) -> int:
-        return 1
+@squin_noise_dialect.register(key="pyqrack")
+class PyQrackMethods(interp.MethodTable):
 
-    def apply(self, qubit: PyQrackQubit, adjoint: bool = False) -> None:
-        if random.uniform(0.0, 1.0) < self.p:
-            qubit.state = QubitState.Lost
+    single_pauli_choices = ("i", "x", "y", "z")
+    two_pauli_choices = (
+        "ii",
+        "ix",
+        "iy",
+        "iz",
+        "xi",
+        "xx",
+        "xy",
+        "xz",
+        "yi",
+        "yx",
+        "yy",
+        "yz",
+        "zi",
+        "zx",
+        "zy",
+        "zz",
+    )
+
+    @interp.impl(Depolarize)
+    def depolarize(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: Depolarize
+    ):
+        p = frame.get(stmt.p)
+        ps = [p / 3.0] * 3
+        qubits = frame.get(stmt.qubits)
+        self.apply_single_qubit_pauli_error(interp, ps, qubits)
 
+    @interp.impl(Depolarize2)
+    def depolarize2(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: Depolarize2
+    ):
+        p = frame.get(stmt.p)
+        ps = [p / 15.0] * 15
+        controls = frame.get(stmt.controls)
+        targets = frame.get(stmt.targets)
+        self.apply_two_qubit_pauli_error(interp, ps, controls, targets)
 
-@squin_noise_dialect.register(key="pyqrack")
-class PyQrackMethods(interp.MethodTable):
-    @interp.impl(StochasticUnitaryChannel)
-    def stochastic_unitary_channel(
+    @interp.impl(SingleQubitPauliChannel)
+    def single_qubit_pauli_channel(
         self,
         interp: PyQrackInterpreter,
         frame: interp.Frame,
-        stmt: StochasticUnitaryChannel,
+        stmt: SingleQubitPauliChannel,
     ):
-        operators = frame.get(stmt.operators)
-        probabilities = frame.get(stmt.probabilities)
-
-        return (StochasticUnitaryChannelRuntime(operators, probabilities),)
+        px = frame.get(stmt.px)
+        py = frame.get(stmt.py)
+        pz = frame.get(stmt.pz)
+        qubits = frame.get(stmt.qubits)
+        self.apply_single_qubit_pauli_error(interp, [px, py, pz], qubits)
+
+    @interp.impl(TwoQubitPauliChannel)
+    def two_qubit_pauli_channel(
+        self,
+        interp: PyQrackInterpreter,
+        frame: interp.Frame,
+        stmt: TwoQubitPauliChannel,
+    ):
+        ps = frame.get(stmt.probabilities)
+        controls = frame.get(stmt.controls)
+        targets = frame.get(stmt.targets)
+        self.apply_two_qubit_pauli_error(interp, ps, controls, targets)
 
     @interp.impl(QubitLoss)
     def qubit_loss(
         self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: QubitLoss
     ):
         p = frame.get(stmt.p)
-        return (QubitLossRuntime(p),)
+        qubits: list[PyQrackQubit] = frame.get(stmt.qubits)
+        for qbit in qubits:
+            if interp.rng_state.uniform(0.0, 1.0) <= p:
+                qbit.drop()
+
+    @interp.impl(CorrelatedQubitLoss)
+    def correlated_qubit_loss(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: CorrelatedQubitLoss
+    ):
+        p = frame.get(stmt.p)
+        qubits: list[list[PyQrackQubit]] = frame.get(stmt.qubits)
+        for qubit_group in qubits:
+            if interp.rng_state.uniform(0.0, 1.0) <= p:
+                for qbit in qubit_group:
+                    qbit.drop()
+
+    def apply_single_qubit_pauli_error(
+        self,
+        interp: PyQrackInterpreter,
+        ps: list[float],
+        qubits: list[PyQrackQubit],
+    ):
+        pi = 1 - sum(ps)
+        probs = [pi] + ps
+
+        assert all(0 <= x <= 1 for x in probs), "Invalid Pauli error probabilities"
+
+        for qbit in qubits:
+            which = interp.rng_state.choice(self.single_pauli_choices, p=probs)
+            self.apply_pauli_error(which, qbit)
+
+    def apply_two_qubit_pauli_error(
+        self,
+        interp: PyQrackInterpreter,
+        ps: list[float],
+        controls: list[PyQrackQubit],
+        targets: list[PyQrackQubit],
+    ):
+        pii = 1 - sum(ps)
+        probs = [pii] + ps
+        assert all(0 <= x <= 1 for x in probs), "Invalid Pauli error probabilities"
+
+        for control, target in zip(controls, targets):
+            which = interp.rng_state.choice(self.two_pauli_choices, p=probs)
+            self.apply_pauli_error(which[0], control)
+            self.apply_pauli_error(which[1], target)
+
+    def apply_pauli_error(self, which: str, qbit: PyQrackQubit):
+        if not qbit.is_active() or which == "i":
+            return
+
+        getattr(qbit.sim_reg, which)(qbit.addr)

bloqade/pyqrack/squin/qubit.py

@@ -3,61 +3,64 @@ from typing import Any
 from kirin import interp
 from kirin.dialects import ilist
 
-from bloqade.squin import qubit
-from bloqade.pyqrack.reg import QubitState, PyQrackQubit
+from bloqade.qubit import stmts as qubit
+from bloqade.pyqrack.reg import QubitState, Measurement, PyQrackQubit
 from bloqade.pyqrack.base import PyQrackInterpreter
 
-from .runtime import OperatorRuntimeABC
-
 
 @qubit.dialect.register(key="pyqrack")
 class PyQrackMethods(interp.MethodTable):
     @interp.impl(qubit.New)
-    def new(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.New):
-        n_qubits: int = frame.get(stmt.n_qubits)
-        qreg = ilist.IList(
-            [
-                PyQrackQubit(i, interp.memory.sim_reg, QubitState.Active)
-                for i in interp.memory.allocate(n_qubits=n_qubits)
-            ]
-        )
-        return (qreg,)
-
-    @interp.impl(qubit.Apply)
-    def apply(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.Apply):
-        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
-        operator: OperatorRuntimeABC = frame.get(stmt.operator)
-        operator.apply(*qubits)
-
-    @interp.impl(qubit.Broadcast)
-    def broadcast(
-        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.Broadcast
+    def new_qubit(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.New
     ):
-        operator: OperatorRuntimeABC = frame.get(stmt.operator)
-        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
-        operator.broadcast_apply(qubits)
+        (addr,) = interp.memory.allocate(1)
+        qb = PyQrackQubit(addr, interp.memory.sim_reg, QubitState.Active)
+        return (qb,)
 
     def _measure_qubit(self, qbit: PyQrackQubit, interp: PyQrackInterpreter):
         if qbit.is_active():
-            return bool(qbit.sim_reg.m(qbit.addr))
+            m = Measurement(bool(qbit.sim_reg.m(qbit.addr)))
         else:
-            return interp.loss_m_result
+            m = Measurement(interp.loss_m_result)
+
+        interp.set_global_measurement_id(m)
+        return m
 
-    @interp.impl(qubit.MeasureQubitList)
+    @interp.impl(qubit.Measure)
     def measure_qubit_list(
         self,
         interp: PyQrackInterpreter,
         frame: interp.Frame,
-        stmt: qubit.MeasureQubitList,
+        stmt: qubit.Measure,
     ):
         qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
         result = ilist.IList([self._measure_qubit(qbit, interp) for qbit in qubits])
         return (result,)
 
-    @interp.impl(qubit.MeasureQubit)
-    def measure_qubit(
-        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.MeasureQubit
+    @interp.impl(qubit.QubitId)
+    def qubit_id(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.QubitId
     ):
-        qbit: PyQrackQubit = frame.get(stmt.qubit)
-        result = self._measure_qubit(qbit, interp)
-        return (result,)
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+        ids = ilist.IList([qbit.addr for qbit in qubits])
+        return (ids,)
+
+    @interp.impl(qubit.MeasurementId)
+    def measurement_id(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.MeasurementId
+    ):
+        measurements: ilist.IList[Measurement, Any] = frame.get(stmt.measurements)
+        ids = ilist.IList([measurement.measurement_id for measurement in measurements])
+        return (ids,)
+
+    @interp.impl(qubit.Reset)
+    def reset(self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: qubit.Reset):
+        qubits: ilist.IList[PyQrackQubit, Any] = frame.get(stmt.qubits)
+        for qbit in qubits:
+            if not qbit.is_active():
+                continue
+
+            m = qbit.sim_reg.m(qbit.addr)
+            if m == Measurement.One:
+                qbit.sim_reg.x(qbit.addr)

bloqade/pyqrack/target.py

@@ -12,7 +12,7 @@ from bloqade.pyqrack.base import (
     PyQrackInterpreter,
     _default_pyqrack_args,
 )
-from bloqade.analysis.address import AnyAddress, AddressAnalysis
+from bloqade.analysis.address import UnknownQubit, AddressAnalysis
 
 Params = ParamSpec("Params")
 RetType = TypeVar("RetType")
@@ -51,9 +51,9 @@ class PyQrack:
             return PyQrackInterpreter(mt.dialects, memory=DynamicMemory(options))
         else:
             address_analysis = AddressAnalysis(mt.dialects)
-            frame, _ = address_analysis.run_analysis(mt)
+            frame, _ = address_analysis.run(mt)
             if self.min_qubits == 0 and any(
-                isinstance(a, AnyAddress) for a in frame.entries.values()
+                isinstance(a, UnknownQubit) for a in frame.entries.values()
             ):
                 raise ValueError(
                     "All addresses must be resolved. Or set min_qubits to a positive integer."
@@ -87,7 +87,8 @@ class PyQrack:
         """
         fold = Fold(mt.dialects)
         fold(mt)
-        return self._get_interp(mt).run(mt, args, kwargs)
+        _, ret = self._get_interp(mt).run(mt, *args, **kwargs)
+        return ret
 
     def multi_run(
         self,

bloqade/pyqrack/task.py

@@ -1,7 +1,12 @@
 from typing import TypeVar, ParamSpec, cast
+from collections import Counter
 from dataclasses import dataclass
 
+import numpy as np
+from kirin.dialects.ilist import IList
+
 from bloqade.task import AbstractSimulatorTask
+from bloqade.pyqrack.reg import QubitState, PyQrackQubit
 from bloqade.pyqrack.base import (
     MemoryABC,
     PyQrackInterpreter,
@@ -19,14 +24,12 @@ class PyQrackSimulatorTask(AbstractSimulatorTask[Param, RetType, MemoryType]):
     pyqrack_interp: PyQrackInterpreter[MemoryType]
 
     def run(self) -> RetType:
-        return cast(
-            RetType,
-            self.pyqrack_interp.run(
-                self.kernel,
-                args=self.args,
-                kwargs=self.kwargs,
-            ),
+        _, ret = self.pyqrack_interp.run(
+            self.kernel,
+            *self.args,
+            **self.kwargs,
         )
+        return cast(RetType, ret)
 
     @property
     def state(self) -> MemoryType:
@@ -36,3 +39,107 @@ class PyQrackSimulatorTask(AbstractSimulatorTask[Param, RetType, MemoryType]):
         """Returns the state vector of the simulator."""
         self.run()
         return self.state.sim_reg.out_ket()
+
+    def qubits(self) -> list[PyQrackQubit]:
+        """Returns the qubits in the simulator."""
+        try:
+            N = self.state.sim_reg.num_qubits()
+            return [
+                PyQrackQubit(
+                    addr=i, sim_reg=self.state.sim_reg, state=QubitState.Active
+                )
+                for i in range(N)
+            ]
+        except AttributeError:
+            Warning("Task has not been run, there are no qubits!")
+            return []
+
+    def batch_run(self, shots: int = 1) -> dict[RetType, float]:
+        """
+        Repeatedly run the task to collect statistics on the shot outcomes.
+        The average is done over [shots] repetitions and thus is frequentist
+        and converges to exact only in the shots -> infinity limit.
+
+        Args:
+            shots (int):
+                the number of repetitions of the task
+        Returns:
+            dict[RetType, float]:
+                a dictionary mapping outcomes to their probabilities,
+                as estimated from counting the shot outcomes. RetType must be hashable.
+        """
+
+        results: list[RetType] = [self.run() for _ in range(shots)]
+
+        # Convert IList to tuple so that it is hashable by Counter
+        def convert(data):
+            if isinstance(data, (list, IList)):
+                return tuple(convert(item) for item in data)
+            return data
+
+        results = convert(results)
+
+        data = {
+            key: value / len(results) for key, value in Counter(results).items()
+        }  # Normalize to probabilities
+        return data
+
+    def batch_state(
+        self, shots: int = 1, qubit_map: None = None
+    ) -> "QuantumState":  # noqa: F821
+        """
+        Repeatedly run the task to extract the averaged quantum state.
+        The average is done over [shots] repetitions and thus is frequentist
+        and converges to exact only in the shots -> infinity limit.
+
+        Args:
+            shots (int):
+                the number of repetitions of the task
+            qubit_map (callable | None):
+                an optional callable that takes the output of self.run() and extract
+                the [returned] qubits to be used for the quantum state.
+                If None, all qubits in the simulator are used, in the order set by the simulator.
+                If callable, qubit_map must have the signature
+                > qubit_map(output:RetType) -> list[PyQrackQubit]
+                and the averaged state is
+                > quantum_state(qubit_map(self.run())).
+                If qubit_map is not None, self.run() must return qubit(s).
+                Two common patterns here are:
+                 > qubit_map = lambda qubits: qubits
+                for the case where self.run() returns a list of qubits, or
+                 > qubit_map = lambda qubit: [qubits]
+                for the case where self.run() returns a single qubit.
+        Returns:
+            QuantumState:
+                the averaged quantum state as a density matrix,
+                represented in its eigenbasis.
+        """
+        # Import here to avoid circular dependencies.
+        from bloqade.pyqrack.device import QuantumState, PyQrackSimulatorBase
+
+        states: list[QuantumState] = []
+        for _ in range(shots):
+            res = self.run()
+            if callable(qubit_map):
+                qbs = qubit_map(res)
+            else:
+                qbs = self.qubits()
+            states.append(PyQrackSimulatorBase.quantum_state(qbs))
+
+        state = QuantumState(
+            eigenvectors=np.concatenate(
+                [state.eigenvectors for state in states], axis=1
+            ),
+            eigenvalues=np.concatenate([state.eigenvalues for state in states], axis=0)
+            / len(states),
+        )
+
+        # Canonicalize the state by orthoganalizing the basis vectors.
+        tol = 1e-7
+        s, v, d = np.linalg.svd(
+            state.eigenvectors * np.sqrt(state.eigenvalues), full_matrices=False
+        )
+        mask = v > tol
+        v = v[mask] ** 2
+        s = s[:, mask]
+        return QuantumState(eigenvalues=v, eigenvectors=s)

bloqade/qasm2/_qasm_loading.py

@@ -4,6 +4,7 @@ import pathlib
 from typing import Any
 
 from kirin import ir, lowering
+from kirin.types import MethodType
 from kirin.dialects import func
 
 from . import parse
@@ -82,11 +83,10 @@ def loads(
         body=body,
     )
 
+    body.blocks[0].args.append_from(MethodType, kernel_name + "_self")
+
     mt = ir.Method(
-        mod=None,
-        py_func=None,
         sym_name=kernel_name,
-        arg_names=[],
         dialects=qasm2_lowering.dialects,
         code=code,
     )

bloqade/qasm2/dialects/core/address.py

@@ -1,12 +1,14 @@
 from kirin import interp
+from kirin.analysis import const
 
 from bloqade.analysis.address import (
     Address,
-    NotQubit,
     AddressReg,
-    AddressQubit,
+    ConstResult,
+    UnknownQubit,
     AddressAnalysis,
 )
+from bloqade.analysis.address.lattice import UnknownReg
 
 from .stmts import QRegGet, QRegNew
 from ._dialect import dialect
@@ -22,17 +24,24 @@ class AddressMethodTable(interp.MethodTable):
         frame: interp.Frame[Address],
         stmt: QRegNew,
     ):
-        n_qubits = interp.get_const_value(int, stmt.n_qubits)
-        addr = AddressReg(range(interp.next_address, interp.next_address + n_qubits))
-        interp.next_address += n_qubits
-        return (addr,)
+        n_qubits = frame.get(stmt.n_qubits)
+        match n_qubits:
+            case ConstResult(const.Value(int() as n)):
+                addr = AddressReg(range(interp.next_address, interp.next_address + n))
+                interp.next_address += n
+                return (addr,)
+            case _:
+                return (UnknownReg(),)
 
     @interp.impl(QRegGet)
     def get(self, interp: AddressAnalysis, frame: interp.Frame[Address], stmt: QRegGet):
         addr = frame.get(stmt.reg)
-        pos = interp.get_const_value(int, stmt.idx)
-        if isinstance(addr, AddressReg):
-            global_idx = addr.data[pos]
-            return (AddressQubit(global_idx),)
-        else:  # this is not reachable
-            return (NotQubit(),)
+        idx = frame.get(stmt.idx)
+
+        typ, values = interp.unpack_iterable(addr)
+        idx_value = interp.get_const_value(idx, int)
+
+        if typ is not None and idx_value is not None:
+            return (values[idx_value],)
+
+        return (UnknownQubit(),)

bloqade/qasm2/dialects/expr/_emit.py

@@ -20,7 +20,10 @@ class EmitExpr(interp.MethodTable):
 
         args: list[ast.Node] = []
         cparams, qparams = [], []
-        for arg in stmt.body.blocks[0].args:
+        entry_args = stmt.body.blocks[0].args
+        user_args = entry_args[1:] if len(entry_args) > 0 else []
+
+        for arg in user_args:
             assert arg.name is not None
 
             args.append(ast.Name(id=arg.name))
@@ -29,14 +32,22 @@ class EmitExpr(interp.MethodTable):
             else:
                 cparams.append(arg.name)
 
-        emit.run_ssacfg_region(frame, stmt.body, tuple(args))
-        emit.output = ast.Gate(
-            name=stmt.sym_name,
-            cparams=cparams,
-            qparams=qparams,
-            body=frame.body,
+        frame.worklist.append(interp.Successor(stmt.body.blocks[0], *args))
+        if len(entry_args) > 0:
+            frame.set(entry_args[0], ast.Name(stmt.sym_name or "gate"))
+
+        while (succ := frame.worklist.pop()) is not None:
+            frame.set_values(succ.block.args[1:], succ.block_args)
+            block_header = emit.emit_block(frame, succ.block)
+            frame.block_ref[succ.block] = block_header
+        return (
+            ast.Gate(
+                name=stmt.sym_name,
+                cparams=cparams,
+                qparams=qparams,
+                body=frame.body,
+            ),
         )
-        return ()
 
     @interp.impl(stmts.ConstInt)
     @interp.impl(stmts.ConstFloat)

bloqade/qasm2/dialects/expr/stmts.py

@@ -87,7 +87,7 @@ class ConstPI(ir.Statement):
 
 
 # QASM 2.0 arithmetic operations
-PyNum = types.Union(types.Int, types.Float)
+PyNum = types.TypeVar("PyNum", bound=types.Union(types.Int, types.Float))
 
 
 @statement(dialect=dialect)
@@ -110,7 +110,7 @@ class Sin(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
     value: ir.SSAValue = info.argument(PyNum)
     """value (Union[int, float]): The number to take the sine of."""
-    result: ir.ResultValue = info.result(PyNum)
+    result: ir.ResultValue = info.result(types.Float)
     """result (float): The sine of the number."""
 
 
@@ -122,7 +122,7 @@ class Cos(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
     value: ir.SSAValue = info.argument(PyNum)
     """value (Union[int, float]): The number to take the cosine of."""
-    result: ir.ResultValue = info.result(PyNum)
+    result: ir.ResultValue = info.result(types.Float)
     """result (float): The cosine of the number."""
 
 
@@ -134,7 +134,7 @@ class Tan(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
     value: ir.SSAValue = info.argument(PyNum)
     """value (Union[int, float]): The number to take the tangent of."""
-    result: ir.ResultValue = info.result(PyNum)
+    result: ir.ResultValue = info.result(types.Float)
     """result (float): The tangent of the number."""
 
 
@@ -146,7 +146,7 @@ class Exp(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
     value: ir.SSAValue = info.argument(PyNum)
     """value (Union[int, float]): The number to take the exponential of."""
-    result: ir.ResultValue = info.result(PyNum)
+    result: ir.ResultValue = info.result(types.Float)
     """result (float): The exponential of the number."""
 
 
@@ -158,7 +158,7 @@ class Log(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
     value: ir.SSAValue = info.argument(PyNum)
     """value (Union[int, float]): The number to take the natural log of."""
-    result: ir.ResultValue = info.result(PyNum)
+    result: ir.ResultValue = info.result(types.Float)
     """result (float): The natural log of the number."""
 
 
@@ -170,7 +170,7 @@ class Sqrt(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
     value: ir.SSAValue = info.argument(PyNum)
     """value (Union[int, float]): The number to take the square root of."""
-    result: ir.ResultValue = info.result(PyNum)
+    result: ir.ResultValue = info.result(types.Float)
     """result (float): The square root of the number."""
 
 

bloqade/qasm2/dialects/noise/fidelity.py

@@ -1,7 +1,6 @@
 from kirin import interp
 from kirin.lattice import EmptyLattice
 
-from bloqade.analysis.address import AddressQubit, AddressTuple
 from bloqade.analysis.fidelity import FidelityAnalysis
 
 from .stmts import PauliChannel, CZPauliChannel, AtomLossChannel
@@ -32,7 +31,7 @@ class FidelityMethodTable(interp.MethodTable):
         # NOTE: fidelity is just the inverse probability of any noise to occur
         fid = (1 - p) * (1 - p_ctrl)
 
-        interp._current_gate_fidelity *= fid
+        interp.gate_fidelity *= fid
 
     @interp.impl(AtomLossChannel)
     def atom_loss(
@@ -42,10 +41,7 @@ class FidelityMethodTable(interp.MethodTable):
         stmt: AtomLossChannel,
     ):
         # NOTE: since AtomLossChannel acts on IList[Qubit], we know the assigned address is a tuple
-        addresses: AddressTuple = interp.addr_frame.get(stmt.qargs)
-
+        addresses = 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
+        for index in addresses.data:
+            interp.atom_survival_probability[index] *= 1 - stmt.prob