correlated loss support

pyproject.toml

@@ -118,3 +118,6 @@ include = ["src/bloqade/*"]
 
 [tool.pytest.ini_options]
 testpaths = "test/"
+filterwarnings = [
+    "ignore:In cirq 1.6 the default value of `use_repetition_ids`:FutureWarning:cirq.circuits.circuit_operation",
+]

src/bloqade/pyqrack/device.py

@@ -30,7 +30,7 @@ class QuantumState(NamedTuple):
     A representation of a quantum state as a density matrix, where the density matrix is
     rho = sum_i eigenvalues[i] |eigenvectors[:,i]><eigenvectors[:,i]|.
 
-    This reprsentation is efficient for low-rank density matrices by only storing
+    This representation is efficient for low-rank density matrices by only storing
     the non-zero eigenvalues and corresponding eigenvectors of the density matrix.
     For example, a pure state has only one non-zero eigenvalue equal to 1.0.
 

src/bloqade/pyqrack/squin/noise/native.py

@@ -5,6 +5,7 @@
     QubitLoss,
     Depolarize,
     Depolarize2,
+    CorrelatedQubitLoss,
     TwoQubitPauliChannel,
     SingleQubitPauliChannel,
 )
@@ -88,6 +89,17 @@ def qubit_loss(
             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,

src/bloqade/squin/__init__.py

@@ -37,6 +37,7 @@
     sqrt_y_adj as sqrt_y_adj,
     sqrt_z_adj as sqrt_z_adj,
     depolarize2 as depolarize2,
+    correlated_qubit_loss as correlated_qubit_loss,
     two_qubit_pauli_channel as two_qubit_pauli_channel,
     single_qubit_pauli_channel as single_qubit_pauli_channel,
 )

src/bloqade/squin/noise/_interface.py

@@ -37,3 +37,9 @@ def two_qubit_pauli_channel(
 
 @wraps(stmts.QubitLoss)
 def qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+@wraps(stmts.CorrelatedQubitLoss)
+def correlated_qubit_loss(
+    p: float, qubits: ilist.IList[ilist.IList[Qubit, N], Any]
+) -> None: ...

src/bloqade/squin/noise/stmts.py

@@ -96,4 +96,16 @@ class QubitLoss(SingleQubitNoiseChannel):
 
     # NOTE: qubit loss error (not supported by Stim)
     p: ir.SSAValue = info.argument(types.Float)
-    qubits: ir.SSAValue = info.argument(ilist.IListType)
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, types.Any])
+
+
+@statement(dialect=dialect)
+class CorrelatedQubitLoss(NoiseChannel):
+    """
+    Apply a correlated atom loss channel.
+    """
+
+    p: ir.SSAValue = info.argument(types.Float)
+    qubits: ir.SSAValue = info.argument(
+        ilist.IListType[ilist.IListType[QubitType, N], types.Any]
+    )

src/bloqade/squin/stdlib/broadcast/__init__.py

@@ -28,6 +28,7 @@
     depolarize as depolarize,
     qubit_loss as qubit_loss,
     depolarize2 as depolarize2,
+    correlated_qubit_loss as correlated_qubit_loss,
     two_qubit_pauli_channel as two_qubit_pauli_channel,
     single_qubit_pauli_channel as single_qubit_pauli_channel,
 )

src/bloqade/squin/stdlib/broadcast/noise.py

@@ -103,6 +103,32 @@ def qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None:
     noise.qubit_loss(p, qubits)
 
 
+@kernel
+def correlated_qubit_loss(
+    p: float, qubits: ilist.IList[ilist.IList[Qubit, N], Any]
+) -> None:
+    """
+    Apply correlated qubit loss channels to groups of qubits.
+
+    For each group of qubits, applies a correlated loss channel where all qubits
+    within the group are lost together with probability `p`. Loss events are independent
+    between different groups.
+
+    Args:
+        p (float): Loss probability for each group.
+        qubits (IList[IList[Qubit, N], Any]): List of qubit groups. Each sublist
+            represents a group of qubits to which a correlated loss channel is applied.
+
+    Example:
+        >>> q1 = squin.qubit.new(3) # First group: qubits 0, 1, 2
+        >>> q2 = squin.qubit.new(3) # Second group: qubits 3, 4, 5
+        >>> squin.broadcast.correlated_qubit_loss(0.5, [q1, q2])
+        # Each group has 50% chance: either all qubits lost or none lost.
+        # Group 1 and Group 2 outcomes are independent.
+    """
+    noise.correlated_qubit_loss(p, qubits)
+
+
 # NOTE: actual stdlib that doesn't wrap statements starts here
 
 

src/bloqade/squin/stdlib/simple/__init__.py

@@ -28,6 +28,7 @@
     depolarize as depolarize,
     qubit_loss as qubit_loss,
     depolarize2 as depolarize2,
+    correlated_qubit_loss as correlated_qubit_loss,
     two_qubit_pauli_channel as two_qubit_pauli_channel,
     single_qubit_pauli_channel as single_qubit_pauli_channel,
 )

src/bloqade/squin/stdlib/simple/noise.py

@@ -1,4 +1,4 @@
-from typing import Literal, TypeVar
+from typing import Any, Literal, TypeVar
 
 from kirin.dialects import ilist
 
@@ -97,6 +97,20 @@ def qubit_loss(p: float, qubit: Qubit) -> None:
     broadcast.qubit_loss(p, ilist.IList([qubit]))
 
 
+@kernel
+def correlated_qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None:
+    """
+    Apply a correlated qubit loss channel to the given qubits.
+
+    All qubits are lost together with a probability `p`.
+
+    Args:
+        p (float): Probability of the qubits being lost.
+        qubits (IList[Qubit, Any]): The list of qubits to which the correlated noise channel is applied.
+    """
+    broadcast.correlated_qubit_loss(p, ilist.IList([qubits]))
+
+
 # NOTE: actual stdlib that doesn't wrap statements starts here
 
 

src/bloqade/stim/__init__.py

@@ -39,4 +39,5 @@
     pauli_channel1 as pauli_channel1,
     pauli_channel2 as pauli_channel2,
     observable_include as observable_include,
+    correlated_qubit_loss as correlated_qubit_loss,
 )

src/bloqade/stim/_wrappers.py

@@ -194,3 +194,9 @@ def z_error(p: float, targets: tuple[int, ...]) -> None: ...
 
 @wraps(noise.QubitLoss)
 def qubit_loss(probs: tuple[float, ...], targets: tuple[int, ...]) -> None: ...
+
+
+@wraps(noise.CorrelatedQubitLoss)
+def correlated_qubit_loss(
+    probs: tuple[float, ...], targets: tuple[int, ...]
+) -> None: ...

src/bloqade/stim/dialects/noise/emit.py

@@ -81,6 +81,7 @@ def non_stim_error(
         return ()
 
     @impl(stmts.TrivialCorrelatedError)
+    @impl(stmts.CorrelatedQubitLoss)
     def non_stim_corr_error(
         self,
         emit: EmitStimMain,

src/bloqade/stim/dialects/noise/stmts.py

@@ -89,8 +89,8 @@ class NonStimError(ir.Statement):
 class NonStimCorrelatedError(ir.Statement):
     name = "NonStimCorrelatedError"
     traits = frozenset({lowering.FromPythonCall()})
-    nonce: int = (
-        info.attribute()
+    nonce: int = info.attribute(
+        default_factory=lambda: __import__("random").getrandbits(32)
     )  # Must be a unique value, otherwise stim might merge two correlated errors with equal probabilities
     probs: tuple[ir.SSAValue, ...] = info.argument(types.Float)
     targets: tuple[ir.SSAValue, ...] = info.argument(types.Int)
@@ -109,3 +109,8 @@ class TrivialError(NonStimError):
 @statement(dialect=dialect)
 class QubitLoss(NonStimError):
     name = "loss"
+
+
+@statement(dialect=dialect)
+class CorrelatedQubitLoss(NonStimCorrelatedError):
+    name = "correlated_loss"

src/bloqade/stim/rewrite/qubit_to_stim.py

@@ -8,6 +8,7 @@
     SQUIN_STIM_OP_MAPPING,
     rewrite_Control,
     rewrite_QubitLoss,
+    rewrite_CorrelatedQubitLoss,
     insert_qubit_idx_from_address,
 )
 
@@ -38,6 +39,9 @@ def rewrite_Apply_and_Broadcast(
         if isinstance(applied_op, noise.stmts.QubitLoss):
             return rewrite_QubitLoss(stmt)
 
+        if isinstance(applied_op, noise.stmts.CorrelatedQubitLoss):
+            return rewrite_CorrelatedQubitLoss(stmt)
+
         assert isinstance(applied_op, op.stmts.Operator)
 
         if isinstance(applied_op, op.stmts.Control):

src/bloqade/stim/rewrite/squin_noise.py

@@ -31,8 +31,10 @@ def rewrite_Apply_and_Broadcast(
 
         # this is an SSAValue, need it to be the actual operator
         applied_op = stmt.operator.owner
-
-        if isinstance(applied_op, squin_noise.stmts.QubitLoss):
+        if isinstance(
+            applied_op,
+            (squin_noise.stmts.QubitLoss, squin_noise.stmts.CorrelatedQubitLoss),
+        ):
             return RewriteResult()
 
         if isinstance(applied_op, squin_noise.stmts.NoiseChannel):

src/bloqade/stim/rewrite/util.py

@@ -21,6 +21,7 @@
     op.stmts.Identity: gate.Identity,
     op.stmts.Reset: collapse.RZ,
     squin_noise.stmts.QubitLoss: stim_noise.QubitLoss,
+    squin_noise.stmts.CorrelatedQubitLoss: stim_noise.CorrelatedQubitLoss,
 }
 
 # Squin allows creation of control gates where the gate can be any operator,
@@ -201,6 +202,33 @@ def rewrite_QubitLoss(
     return RewriteResult(has_done_something=True)
 
 
+def rewrite_CorrelatedQubitLoss(
+    stmt: qubit.Apply | qubit.Broadcast | wire.Broadcast | wire.Apply,
+) -> RewriteResult:
+    """
+    Rewrite CorrelatedQubitLoss statements to Stim's TrivialCorrelatedError.
+    """
+
+    squin_loss_op = stmt.operator.owner
+    assert isinstance(squin_loss_op, squin_noise.stmts.CorrelatedQubitLoss)
+
+    qubit_idx_ssas = insert_qubit_idx_after_apply(stmt=stmt)
+    if qubit_idx_ssas is None:
+        return RewriteResult()
+
+    stim_loss_stmt = stim_noise.CorrelatedQubitLoss(
+        targets=qubit_idx_ssas,
+        probs=(squin_loss_op.p,),
+    )
+
+    if isinstance(stmt, (wire.Apply, wire.Broadcast)):
+        create_wire_passthrough(stmt)
+
+    stmt.replace_by(stim_loss_stmt)
+
+    return RewriteResult(has_done_something=True)
+
+
 def create_wire_passthrough(stmt: wire.Apply | wire.Broadcast) -> None:
 
     for input_wire, output_wire in zip(stmt.inputs, stmt.results):

src/bloqade/stim/rewrite/wire_to_stim.py

@@ -6,6 +6,7 @@
     SQUIN_STIM_OP_MAPPING,
     rewrite_Control,
     rewrite_QubitLoss,
+    rewrite_CorrelatedQubitLoss,
     insert_qubit_idx_from_wire_ssa,
 )
 
@@ -29,6 +30,9 @@ def rewrite_Apply_and_Broadcast(
         if isinstance(applied_op, noise.stmts.QubitLoss):
             return rewrite_QubitLoss(stmt)
 
+        if isinstance(applied_op, noise.stmts.CorrelatedQubitLoss):
+            return rewrite_CorrelatedQubitLoss(stmt)
+
         assert isinstance(applied_op, op.stmts.Operator)
 
         if isinstance(applied_op, op.stmts.Control):

test/pyqrack/squin/test_noise.py

@@ -16,6 +16,22 @@ def main():
     assert not qubit.is_active()
 
 
+def test_correlated_loss():
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(5)
+        squin.correlated_qubit_loss(0.5, q[0:4])
+        return q
+
+    target = PyQrack(5)
+    for _ in range(10):
+        qubits = target.run(main)
+        qubits_active = [q.is_active() for q in qubits[:4]]
+        assert all(qubits_active) or not any(qubits_active)
+
+        assert qubits[4].is_active()
+
+
 def test_pauli_channel():
     @squin.kernel
     def single_qubit():

test/squin/noise/test_stdlib_noise.py

@@ -1,3 +1,6 @@
+import numpy as np
+import pytest
+
 from bloqade import squin
 from bloqade.pyqrack import PyQrackQubit, StackMemorySimulator
 
@@ -19,6 +22,62 @@ def main():
     assert not qubit.is_active()
 
 
+@pytest.mark.parametrize(
+    "seed, expected_loss_triggered",
+    [
+        (0, False),  # Seed 0: no loss
+        (2, True),  # Seed 2: qubits 0-3 are lost
+    ],
+)
+def test_correlated_loss(seed, expected_loss_triggered):
+
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(5)
+        squin.correlated_qubit_loss(0.5, q[0:4])
+        return q
+
+    rng = np.random.default_rng(seed=seed)
+    sim = StackMemorySimulator(min_qubits=5, rng_state=rng)
+    qubits = sim.run(main)
+
+    for q in qubits:
+        assert isinstance(q, PyQrackQubit)
+
+    for q in qubits[:4]:
+        assert not q.is_active() if expected_loss_triggered else q.is_active()
+
+    assert qubits[4].is_active()
+
+
+@pytest.mark.parametrize(
+    "seed, expected_loss_triggered",
+    [(0, (False, True)), (1, (False, False)), (2, (True, True)), (8, (True, False))],
+)
+def test_correlated_loss_broadcast(seed, expected_loss_triggered):
+
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(6)
+        q1 = q[:3]
+        q2 = q[3:]
+        squin.broadcast.correlated_qubit_loss(0.5, [q1, q2])
+        return q
+
+    rng = np.random.default_rng(seed=seed)
+    sim = StackMemorySimulator(min_qubits=5, rng_state=rng)
+    qubits = sim.run(main)
+
+    for q in qubits:
+        assert isinstance(q, PyQrackQubit)
+
+    for q in qubits[:3]:
+        assert not q.is_active() if expected_loss_triggered[0] else q.is_active()
+
+    for q in qubits[3:]:
+        assert not q.is_active() if expected_loss_triggered[1] else q.is_active()
+
+
 def test_bit_flip():
 
     @squin.kernel