Implement squin noise rewrite for depolarize2 (#441)

Another thing that came out of #436.

---------

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

Author: david-pl

src/bloqade/squin/noise/rewrite.py

@@ -3,14 +3,15 @@
 from kirin import ir
 from kirin.passes import Pass
 from kirin.rewrite import Walk
-from kirin.dialects import ilist
+from kirin.dialects import py, ilist
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
 from .stmts import (
     PPError,
     QubitLoss,
     Depolarize,
     PauliError,
+    Depolarize2,
     NoiseChannel,
     TwoQubitPauliChannel,
     SingleQubitPauliChannel,
@@ -58,6 +59,18 @@ def rewrite_single_qubit_pauli_channel(
     def rewrite_two_qubit_pauli_channel(
         self, node: TwoQubitPauliChannel
     ) -> RewriteResult:
+        operator_list = self._insert_two_qubit_paulis_before_node(node)
+        stochastic_unitary = StochasticUnitaryChannel(
+            operators=operator_list, probabilities=node.params
+        )
+
+        node.replace_by(stochastic_unitary)
+        return RewriteResult(has_done_something=True)
+
+    @staticmethod
+    def _insert_two_qubit_paulis_before_node(
+        node: TwoQubitPauliChannel | Depolarize2,
+    ) -> ir.ResultValue:
         paulis = (Identity(sites=1), X(), Y(), Z())
         for op in paulis:
             op.insert_before(node)
@@ -71,12 +84,7 @@ def rewrite_two_qubit_pauli_channel(
             operators.append(op.result)
 
         (operator_list := ilist.New(values=operators)).insert_before(node)
-        stochastic_unitary = StochasticUnitaryChannel(
-            operators=operator_list.result, probabilities=node.params
-        )
-
-        node.replace_by(stochastic_unitary)
-        return RewriteResult(has_done_something=True)
+        return operator_list.result
 
     def rewrite_p_p_error(self, node: PPError) -> RewriteResult:
         (operators := ilist.New(values=(node.op,))).insert_before(node)
@@ -95,8 +103,14 @@ def rewrite_depolarize(self, node: Depolarize) -> RewriteResult:
             op.insert_before(node)
             operators.append(op.result)
 
+        # NOTE: need to divide the probability by 3 to get the correct total error rate
+        (three := py.Constant(3)).insert_before(node)
+        (p_over_3 := py.Div(node.p, three.result)).insert_before(node)
+
         (operator_list := ilist.New(values=operators)).insert_before(node)
-        (ps := ilist.New(values=[node.p for _ in range(3)])).insert_before(node)
+        (ps := ilist.New(values=[p_over_3.result for _ in range(3)])).insert_before(
+            node
+        )
 
         stochastic_unitary = StochasticUnitaryChannel(
             operators=operator_list.result, probabilities=ps.result
@@ -105,6 +119,21 @@ def rewrite_depolarize(self, node: Depolarize) -> RewriteResult:
 
         return RewriteResult(has_done_something=True)
 
+    def rewrite_depolarize2(self, node: Depolarize2) -> RewriteResult:
+        operator_list = self._insert_two_qubit_paulis_before_node(node)
+
+        # NOTE: need to divide the probability by 15 to get the correct total error rate
+        (fifteen := py.Constant(15)).insert_before(node)
+        (p_over_15 := py.Div(node.p, fifteen.result)).insert_before(node)
+        (probs := ilist.New(values=[p_over_15.result] * 15)).insert_before(node)
+
+        stochastic_unitary = StochasticUnitaryChannel(
+            operators=operator_list, probabilities=probs.result
+        )
+        node.replace_by(stochastic_unitary)
+
+        return RewriteResult(has_done_something=True)
+
 
 class RewriteNoiseStmts(Pass):
     def unsafe_run(self, mt: ir.Method):

test/pyqrack/squin/test_noise.py

@@ -209,3 +209,22 @@ def main():
     print(ket)
 
     assert math.isclose(abs(ket[2]) ** 2, 1.0, abs_tol=1e-5)
+
+
+def test_depolarize2():
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(2)
+        err = squin.noise.depolarize2(0.1)
+        squin.qubit.apply(err, q[0], q[1])
+
+    main.print()
+
+    main_ = main.similar(main.dialects)
+
+    result = RewriteNoiseStmts(main.dialects)(main_)
+    assert result.has_done_something
+    main_.print()
+
+    sim = StackMemorySimulator(min_qubits=2)
+    sim.run(main)