fix import

Author: kaihsin

src/bloqade/noise/fidelity.py

@@ -3,12 +3,12 @@
 
 from bloqade.analysis.fidelity import FidelityAnalysis
 
-from .native import dialect as native
+from .native import dialect 
 from .native.stmts import PauliChannel, CZPauliChannel, AtomLossChannel
-from ..analysis.address import AddressQubit, AddressTuple
+from bloqade.analysis.address import AddressQubit, AddressTuple
 
 
-@native.register(key="circuit.fidelity")
+@dialect.register(key="circuit.fidelity")
 class FidelityMethodTable(interp.MethodTable):
 
     @interp.impl(PauliChannel)

src/bloqade/noise/native/__init__.py

@@ -1,7 +1,6 @@
 """NOTE: This module is not guaranteed to be supported long-term in bloqade. We will be
 moving towards a more general approach to noise modeling in the future."""
 
-from ._emit import NativeNoise as NativeNoise
 from .model import (
     GateNoiseParams as GateNoiseParams,
     TwoRowZoneModel as TwoRowZoneModel,

src/bloqade/noise/native/_emit.py

@@ -4,7 +4,7 @@
 from kirin.decl import info, statement
 from kirin.dialects import ilist
 
-from bloqade.qasm2.types import QubitType
+from bloqade.types import QubitType
 
 from ._dialect import dialect
 

src/bloqade/noise/native/stmts.py

@@ -1,3 +1,4 @@
+from typing import Any
 from dataclasses import field, dataclass
 
 from kirin import ir, types, interp
@@ -7,6 +8,8 @@
 from bloqade.types import QubitType
 from bloqade.qasm2.parse import ast
 
+from bloqade.noise import native
+
 from .base import EmitError, EmitQASM2Base, EmitQASM2Frame
 
 
@@ -86,3 +89,84 @@ def emit_err(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt):
     @interp.impl(func.ConstantNone)
     def ignore(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt):
         return ()
+
+
+@native.dialect.register(key="emit.qasm2.gate")
+class NativeNoise(interp.MethodTable):
+
+    def _convert(self, node: ast.Bit | ast.Name) -> str:
+        if isinstance(node, ast.Bit):
+            return f"{node.name.id}[{node.addr}]"
+        else:
+            return f"{node.id}"
+
+    @interp.impl(native.CZPauliChannel)
+    def emit_czp(
+        self,
+        emit: EmitQASM2Gate,
+        frame: EmitQASM2Frame,
+        stmt: native.CZPauliChannel,
+    ):
+        paired: bool = stmt.paired
+        px_ctrl: float = stmt.px_ctrl
+        py_ctrl: float = stmt.py_ctrl
+        pz_ctrl: float = stmt.pz_ctrl
+        px_qarg: float = stmt.pz_qarg
+        py_qarg: float = stmt.py_qarg
+        pz_qarg: float = stmt.pz_qarg
+        ctrls: ilist.IList[ast.Bit, Any] = frame.get(stmt.ctrls)
+        qargs: ilist.IList[ast.Bit, Any] = frame.get(stmt.qargs)
+        frame.body.append(
+            ast.Comment(
+                text=f"native.CZPauliChannel(paired={paired}, p_ctrl[{px_ctrl}, {py_ctrl}, {pz_ctrl}], p_qarg[{px_qarg}, {py_qarg}, {pz_qarg}])"
+            )
+        )
+        frame.body.append(
+            ast.Comment(
+                text=f" -: ctrls: {', '.join([self._convert(q) for q in ctrls])}"
+            )
+        )
+        frame.body.append(
+            ast.Comment(
+                text=f" -: qargs: {', '.join([self._convert(q) for q in qargs])}"
+            )
+        )
+        return ()
+
+    @interp.impl(native.AtomLossChannel)
+    def emit_loss(
+        self,
+        emit: EmitQASM2Gate,
+        frame: EmitQASM2Frame,
+        stmt: native.AtomLossChannel,
+    ):
+        prob: float = stmt.prob
+        qargs: ilist.IList[ast.Bit, Any] = frame.get(stmt.qargs)
+        frame.body.append(ast.Comment(text=f"native.Atomloss(p={prob})"))
+        frame.body.append(
+            ast.Comment(
+                text=f" -: qargs: {', '.join([self._convert(q) for q in qargs])}"
+            )
+        )
+        return ()
+
+    @interp.impl(native.PauliChannel)
+    def emit_pauli(
+        self,
+        emit: EmitQASM2Gate,
+        frame: EmitQASM2Frame,
+        stmt: native.PauliChannel,
+    ):
+        px: float = stmt.px
+        py: float = stmt.py
+        pz: float = stmt.pz
+        qargs: ilist.IList[ast.Bit, Any] = frame.get(stmt.qargs)
+        frame.body.append(
+            ast.Comment(text=f"native.Atomloss(px={px}, py={py}, pz={pz})")
+        )
+        frame.body.append(
+            ast.Comment(
+                text=f" -: qargs: {', '.join([self._convert(q) for q in qargs])}"
+            )
+        )
+        return ()