Add squin.noise.Depolarize2 statement and rewrite to stim (#368)

Also add rewrite to stim depolarize2

---------

Co-authored-by: John Long <[email protected]>

Author: ehua7365

src/bloqade/squin/noise/__init__.py

@@ -4,6 +4,7 @@
     pp_error as pp_error,
     depolarize as depolarize,
     qubit_loss as qubit_loss,
+    depolarize2 as depolarize2,
     pauli_error as pauli_error,
     two_qubit_pauli_channel as two_qubit_pauli_channel,
     single_qubit_pauli_channel as single_qubit_pauli_channel,

src/bloqade/squin/noise/_wrapper.py

@@ -20,6 +20,10 @@ def pp_error(op: Op, p: float) -> Op: ...
 def depolarize(p: float) -> Op: ...
 
 
+@wraps(stmts.Depolarize2)
+def depolarize2(p: float) -> Op: ...
+
+
 @wraps(stmts.SingleQubitPauliChannel)
 def single_qubit_pauli_channel(
     params: ilist.IList[float, Literal[3]] | list[float] | tuple[float, float, float],

src/bloqade/squin/noise/stmts.py

@@ -33,7 +33,16 @@ class PPError(NoiseChannel):
 @statement(dialect=dialect)
 class Depolarize(NoiseChannel):
     """
-    Apply depolarize error to qubit
+    Apply depolarize error to single qubit
+    """
+
+    p: ir.SSAValue = info.argument(types.Float)
+
+
+@statement(dialect=dialect)
+class Depolarize2(NoiseChannel):
+    """
+    Apply correlated depolarize error to two qubit
     """
 
     p: ir.SSAValue = info.argument(types.Float)

src/bloqade/stim/rewrite/squin_noise.py

@@ -140,6 +140,23 @@ def rewrite_TwoQubitPauliChannel(
         )
         return stim_stmt
 
+    def rewrite_Depolarize2(
+        self,
+        stmt: qubit.Apply | qubit.Broadcast | wire.Broadcast | wire.Apply,
+        qubit_idx_ssas: Tuple[SSAValue],
+    ) -> Statement:
+        """Rewrite squin.noise.Depolarize2 to stim.Depolarize2."""
+
+        squin_channel = stmt.operator.owner
+        assert isinstance(squin_channel, squin_noise.stmts.Depolarize2)
+
+        p = get_const_value(float, squin_channel.p)
+        p_stmt = py.Constant(p)
+        p_stmt.insert_before(stmt)
+
+        stim_stmt = stim_noise.Depolarize2(targets=qubit_idx_ssas, p=p_stmt.result)
+        return stim_stmt
+
     def rewrite_Depolarize(
         self,
         stmt: qubit.Apply | qubit.Broadcast | wire.Broadcast | wire.Apply,

test/stim/passes/stim_reference_programs/noise/broadcast_depolarize2.stim

@@ -0,0 +1 @@
+DEPOLARIZE2(0.01500000) 0 1

test/stim/passes/test_squin_noise_to_stim.py

@@ -150,6 +150,20 @@ def test():
     assert codegen(test) == expected_stim_program
 
 
+def test_broadcast_depolarize2():
+
+    @kernel
+    def test():
+        q = qubit.new(2)
+        channel = noise.depolarize2(p=0.015)
+        qubit.broadcast(channel, q)
+        return
+
+    run_address_and_stim_passes(test)
+    expected_stim_program = load_reference_program("broadcast_depolarize2.stim")
+    assert codegen(test) == expected_stim_program
+
+
 def test_apply_depolarize1():
 
     @kernel