Lowering of squin noise to cirq (#343)

Closes #317

---------

Co-authored-by: Phillip Weinberg <[email protected]>

Author: david-pl

src/bloqade/squin/cirq/__init__.py

@@ -9,8 +9,9 @@
 from .. import kernel
 
 # NOTE: just to register methods
-from .emit import op as op, qubit as qubit
+from .emit import op as op, noise as noise, qubit as qubit
 from .lowering import Squin
+from ..noise.rewrite import RewriteNoiseStmts
 from .emit.emit_circuit import EmitCirq
 
 
@@ -235,7 +236,12 @@ def main():
         )
 
     emitter = EmitCirq(qubits=qubits)
-    return emitter.run(mt, args=())
+
+    # Rewrite noise statements
+    mt_ = mt.similar(mt.dialects)
+    RewriteNoiseStmts(mt_.dialects)(mt_)
+
+    return emitter.run(mt_, args=())
 
 
 def dump_circuit(mt: ir.Method, qubits: Sequence[cirq.Qid] | None = None, **kwargs):

src/bloqade/squin/cirq/emit/noise.py

@@ -0,0 +1,49 @@
+import cirq
+from kirin.emit import EmitError
+from kirin.interp import MethodTable, impl
+
+from ... import noise
+from .runtime import (
+    KronRuntime,
+    BasicOpRuntime,
+    OperatorRuntimeABC,
+    PauliStringRuntime,
+)
+from .emit_circuit import EmitCirq, EmitCirqFrame
+
+
+@noise.dialect.register(key="emit.cirq")
+class EmitCirqNoiseMethods(MethodTable):
+
+    @impl(noise.stmts.StochasticUnitaryChannel)
+    def stochastic_unitary_channel(
+        self,
+        emit: EmitCirq,
+        frame: EmitCirqFrame,
+        stmt: noise.stmts.StochasticUnitaryChannel,
+    ):
+        ops = frame.get(stmt.operators)
+        ps = frame.get(stmt.probabilities)
+
+        error_probabilities = {self._op_to_key(op_): p for op_, p in zip(ops, ps)}
+        cirq_op = cirq.asymmetric_depolarize(error_probabilities=error_probabilities)
+        return (BasicOpRuntime(cirq_op),)
+
+    @staticmethod
+    def _op_to_key(operator: OperatorRuntimeABC) -> str:
+        match operator:
+            case KronRuntime():
+                key_lhs = EmitCirqNoiseMethods._op_to_key(operator.lhs)
+                key_rhs = EmitCirqNoiseMethods._op_to_key(operator.rhs)
+                return key_lhs + key_rhs
+
+            case BasicOpRuntime():
+                return str(operator.gate)
+
+            case PauliStringRuntime():
+                return operator.string
+
+            case _:
+                raise EmitError(
+                    f"Unexpected operator runtime in StochasticUnitaryChannel of type {type(operator).__name__} encountered!"
+                )

src/bloqade/squin/cirq/emit/runtime.py

@@ -21,7 +21,10 @@ def apply(
 
     def unsafe_apply(
         self, qubits: Sequence[cirq.Qid], adjoint: bool = False
-    ) -> list[cirq.Operation]: ...
+    ) -> list[cirq.Operation]:
+        raise NotImplementedError(
+            f"Apply method needs to be implemented in {self.__class__.__name__}"
+        )
 
 
 @dataclass
@@ -38,6 +41,11 @@ class BasicOpRuntime(UnsafeOperatorRuntimeABC):
     def num_qubits(self) -> int:
         return self.gate.num_qubits()
 
+    def unsafe_apply(
+        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
+    ) -> list[cirq.Operation]:
+        return [self.gate(*qubits)]
+
 
 @dataclass
 class UnitaryRuntime(BasicOpRuntime):

src/bloqade/squin/noise/_wrapper.py

@@ -1,3 +1,6 @@
+from typing import Literal
+
+from kirin.dialects import ilist
 from kirin.lowering import wraps
 
 from bloqade.squin.op.types import Op
@@ -18,11 +21,15 @@ def depolarize(p: float) -> Op: ...
 
 
 @wraps(stmts.SingleQubitPauliChannel)
-def single_qubit_pauli_channel(params: tuple[float, float, float]) -> Op: ...
+def single_qubit_pauli_channel(
+    params: ilist.IList[float, Literal[3]] | list[float] | tuple[float, float, float],
+) -> Op: ...
 
 
 @wraps(stmts.TwoQubitPauliChannel)
-def two_qubit_pauli_channel(params: tuple[float, ...]) -> Op: ...
+def two_qubit_pauli_channel(
+    params: ilist.IList[float, Literal[15]] | list[float] | tuple[float, ...],
+) -> Op: ...
 
 
 @wraps(stmts.QubitLoss)

src/bloqade/squin/noise/rewrite.py

@@ -58,12 +58,12 @@ def rewrite_single_qubit_pauli_channel(
     def rewrite_two_qubit_pauli_channel(
         self, node: TwoQubitPauliChannel
     ) -> RewriteResult:
-        paulis = (X(), Y(), Z(), Identity(sites=1))
+        paulis = (Identity(sites=1), X(), Y(), Z())
         for op in paulis:
             op.insert_before(node)
 
-        # NOTE: collect list so we can skip the last entry, which will be two identities
-        combinations = list(itertools.product(paulis, repeat=2))[:-1]
+        # NOTE: collect list so we can skip the first entry, which will be two identities
+        combinations = list(itertools.product(paulis, repeat=2))[1:]
         operators: list[ir.SSAValue] = []
         for pauli_1, pauli_2 in combinations:
             op = Kron(pauli_1.result, pauli_2.result)

test/squin/cirq/test_squin_noise_to_cirq.py

@@ -0,0 +1,88 @@
+import cirq
+
+from bloqade import squin
+
+
+def test_pauli_channel(run_sim: bool = False):
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(2)
+        h = squin.op.h()
+        cx = squin.op.cx()
+        squin.qubit.apply(h, q[0])
+        dpl = squin.noise.depolarize(0.1)
+        squin.qubit.apply(dpl, q[0])
+        squin.qubit.apply(cx, q)
+        single_qubit_noise = squin.noise.single_qubit_pauli_channel([0.1, 0.12, 0.13])
+        squin.qubit.apply(single_qubit_noise, q[1])
+        two_qubit_noise = squin.noise.two_qubit_pauli_channel(
+            [
+                0.036,
+                0.007,
+                0.035,
+                0.022,
+                0.063,
+                0.024,
+                0.006,
+                0.033,
+                0.014,
+                0.019,
+                0.023,
+                0.058,
+                0.0,
+                0.0,
+                0.064,
+            ]
+        )
+        squin.qubit.apply(two_qubit_noise, q)
+        squin.qubit.measure(q)
+
+    main.print()
+
+    circuit = squin.cirq.emit_circuit(main)
+
+    print(circuit)
+
+    if run_sim:
+        sim = cirq.Simulator()
+        sim.run(circuit)
+
+
+def test_pauli_error(run_sim: bool = False):
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(2)
+        x = squin.op.x()
+        n = squin.noise.pauli_error(x, 0.1)
+        squin.qubit.apply(n, q[0])
+        squin.qubit.measure(q)
+
+    main.print()
+
+    circuit = squin.cirq.emit_circuit(main)
+
+    print(circuit)
+
+    if run_sim:
+        sim = cirq.Simulator()
+        sim.run(circuit)
+
+
+def test_pperror(run_sim: bool = False):
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(3)
+        ps = squin.op.pauli_string(string="XYZ")
+        n = squin.noise.pp_error(ps, 0.1)
+        squin.qubit.apply(n, q)
+        squin.qubit.measure(q)
+
+    main.print()
+
+    circuit = squin.cirq.emit_circuit(main)
+
+    print(circuit)
+
+    if run_sim:
+        sim = cirq.Simulator()
+        sim.run(circuit)