diff --git a/pyproject.toml b/pyproject.toml index 8886e76c..ec421f78 100644 --- a/pyproject.toml +++ b/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", +] diff --git a/src/bloqade/pyqrack/device.py b/src/bloqade/pyqrack/device.py index ba5fe274..2b58e3af 100644 --- a/src/bloqade/pyqrack/device.py +++ b/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]> None: ... + + +@wraps(stmts.CorrelatedQubitLoss) +def correlated_qubit_loss( + p: float, qubits: ilist.IList[ilist.IList[Qubit, N], Any] +) -> None: ... diff --git a/src/bloqade/squin/noise/stmts.py b/src/bloqade/squin/noise/stmts.py index eddb7942..7ab0042e 100644 --- a/src/bloqade/squin/noise/stmts.py +++ b/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] + ) diff --git a/src/bloqade/squin/stdlib/broadcast/__init__.py b/src/bloqade/squin/stdlib/broadcast/__init__.py index bf55ccbe..f0db57b5 100644 --- a/src/bloqade/squin/stdlib/broadcast/__init__.py +++ b/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, ) diff --git a/src/bloqade/squin/stdlib/broadcast/noise.py b/src/bloqade/squin/stdlib/broadcast/noise.py index 10686007..ebb7d8ba 100644 --- a/src/bloqade/squin/stdlib/broadcast/noise.py +++ b/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 diff --git a/src/bloqade/squin/stdlib/simple/__init__.py b/src/bloqade/squin/stdlib/simple/__init__.py index bf55ccbe..f0db57b5 100644 --- a/src/bloqade/squin/stdlib/simple/__init__.py +++ b/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, ) diff --git a/src/bloqade/squin/stdlib/simple/noise.py b/src/bloqade/squin/stdlib/simple/noise.py index 4bf0ccf8..8f1fb341 100644 --- a/src/bloqade/squin/stdlib/simple/noise.py +++ b/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 diff --git a/src/bloqade/stim/__init__.py b/src/bloqade/stim/__init__.py index eb8644d5..32e2a20e 100644 --- a/src/bloqade/stim/__init__.py +++ b/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, ) diff --git a/src/bloqade/stim/_wrappers.py b/src/bloqade/stim/_wrappers.py index 41ef5bb2..008e6635 100644 --- a/src/bloqade/stim/_wrappers.py +++ b/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: ... diff --git a/src/bloqade/stim/dialects/noise/emit.py b/src/bloqade/stim/dialects/noise/emit.py index 73242a2f..8f3ebefa 100644 --- a/src/bloqade/stim/dialects/noise/emit.py +++ b/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, diff --git a/src/bloqade/stim/dialects/noise/stmts.py b/src/bloqade/stim/dialects/noise/stmts.py index bdfdfbbb..7594d127 100644 --- a/src/bloqade/stim/dialects/noise/stmts.py +++ b/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" diff --git a/src/bloqade/stim/rewrite/qubit_to_stim.py b/src/bloqade/stim/rewrite/qubit_to_stim.py index 1c4302c8..f7c47888 100644 --- a/src/bloqade/stim/rewrite/qubit_to_stim.py +++ b/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): diff --git a/src/bloqade/stim/rewrite/squin_noise.py b/src/bloqade/stim/rewrite/squin_noise.py index 8952792a..10c3cef6 100644 --- a/src/bloqade/stim/rewrite/squin_noise.py +++ b/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): diff --git a/src/bloqade/stim/rewrite/util.py b/src/bloqade/stim/rewrite/util.py index 1efab606..e3b8100f 100644 --- a/src/bloqade/stim/rewrite/util.py +++ b/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): diff --git a/src/bloqade/stim/rewrite/wire_to_stim.py b/src/bloqade/stim/rewrite/wire_to_stim.py index 94640ebd..f411ef5e 100644 --- a/src/bloqade/stim/rewrite/wire_to_stim.py +++ b/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): diff --git a/test/pyqrack/squin/test_noise.py b/test/pyqrack/squin/test_noise.py index 12020069..f746812f 100644 --- a/test/pyqrack/squin/test_noise.py +++ b/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(): diff --git a/test/squin/noise/test_stdlib_noise.py b/test/squin/noise/test_stdlib_noise.py index 2dd90345..c65b6eb9 100644 --- a/test/squin/noise/test_stdlib_noise.py +++ b/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 diff --git a/test/stim/dialects/stim/test_stim_circuits.py b/test/stim/dialects/stim/test_stim_circuits.py index 7b7c1779..84caf6ab 100644 --- a/test/stim/dialects/stim/test_stim_circuits.py +++ b/test/stim/dialects/stim/test_stim_circuits.py @@ -1,3 +1,5 @@ +import re + from bloqade import stim from bloqade.stim.emit import EmitStimMain @@ -94,6 +96,19 @@ def test_qubit_loss(): assert interp.get_output() == "\nI_ERROR[loss](0.10000000, 0.20000000) 0 1 2" +def test_correlated_qubit_loss(): + + @stim.main + def test_correlated_qubit_loss(): + stim.correlated_qubit_loss(probs=(0.1,), targets=(0, 3, 1)) + + interp.run(test_correlated_qubit_loss, args=()) + + assert re.match( + r"\nI_ERROR\[correlated_loss:\d+\]\(0\.10000000\) 0 3 1", interp.get_output() + ) + + def test_collapse(): @stim.main def test_measure():