bloqade-circuit 0.6.6__py3-none-any.whl → 0.6.8__py3-none-any.whl

bloqade/pyqrack/device.py

@@ -4,6 +4,7 @@ from dataclasses import field, dataclass
 import numpy as np
 from kirin import ir
 from kirin.passes import fold
+from kirin.dialects.ilist import IList
 
 from bloqade.squin import noise as squin_noise
 from pyqrack.pauli import Pauli
@@ -18,6 +19,7 @@ from bloqade.pyqrack.base import (
     _default_pyqrack_args,
 )
 from bloqade.pyqrack.task import PyQrackSimulatorTask
+from pyqrack.qrack_simulator import QrackSimulator
 from bloqade.squin.noise.rewrite import RewriteNoiseStmts
 from bloqade.analysis.address.lattice import AnyAddress
 from bloqade.analysis.address.analysis import AddressAnalysis
@@ -26,6 +28,55 @@ RetType = TypeVar("RetType")
 Params = ParamSpec("Params")
 
 
+def _pyqrack_reduced_density_matrix(
+    inds: tuple[int, ...], sim_reg: QrackSimulator, tol: float = 1e-12
+) -> "np.linalg._linalg.EighResult":
+    """
+    Extract the reduced density matrix representing the state of a list
+    of qubits from a PyQRack simulator register.
+
+    Inputs:
+        inds: A list of integers labeling the qubit registers to extract the reduced density matrix for
+        sim_reg: The PyQRack simulator register to extract the reduced density matrix from
+        tol: The tolerance for density matrix eigenvalues to be considered non-zero.
+    Outputs:
+        An eigh result containing the eigenvalues and eigenvectors of the reduced density matrix.
+    """
+    # Identify the rest of the qubits in the register
+    N = sim_reg.num_qubits()
+    other = tuple(set(range(N)).difference(inds))
+
+    if len(set(inds)) != len(inds):
+        raise ValueError("Qubits must be unique.")
+
+    if max(inds) > N - 1:
+        raise ValueError(
+            f"Qubit indices {inds} exceed the number of qubits in the register {N}."
+        )
+
+    reordering = inds + other
+    # Fix pyqrack edannes to be consistent with Cirq.
+    reordering = tuple(N - 1 - x for x in reordering)
+    # Extract the statevector from the PyQRack qubits
+    statevector = np.array(sim_reg.out_ket())
+    # Reshape into a (2,2,2, ..., 2) tensor
+    vec_f = np.reshape(statevector, (2,) * N)
+    # Reorder the indexes to obey the order of the qubits
+    vec_p = np.transpose(vec_f, reordering)
+    # Rehape into a 2^N by 2^M matrix to compute the singular value decomposition
+    vec_svd = np.reshape(vec_p, (2 ** len(inds), 2 ** len(other)))
+    # The singular values and vectors are the eigenspace of the reduced density matrix
+    s, v, d = np.linalg.svd(vec_svd, full_matrices=False)
+
+    # Remove the negligable singular values
+    nonzero_inds = np.where(np.abs(v) > tol)[0]
+    s = s[:, nonzero_inds]
+    v = v[nonzero_inds] ** 2
+    # Forge into the correct result type
+    result = np.linalg._linalg.EighResult(eigenvalues=v, eigenvectors=s)
+    return result
+
+
 @dataclass
 class PyQrackSimulatorBase(AbstractSimulatorDevice[PyQrackSimulatorTask]):
     """PyQrack simulation device base class."""
@@ -50,7 +101,6 @@ class PyQrackSimulatorBase(AbstractSimulatorDevice[PyQrackSimulatorTask]):
         kwargs: dict[str, Any],
         memory: MemoryType,
     ) -> PyQrackSimulatorTask[Params, RetType, MemoryType]:
-
         if squin_noise in mt.dialects:
             # NOTE: rewrite noise statements
             mt_ = mt.similar(mt.dialects)
@@ -112,6 +162,51 @@ class PyQrackSimulatorBase(AbstractSimulatorDevice[PyQrackSimulatorTask]):
 
         return sim_reg.pauli_expectation(qubit_ids, pauli)
 
+    @staticmethod
+    def quantum_state(
+        qubits: list[PyQrackQubit] | IList[PyQrackQubit, Any], tol: float = 1e-12
+    ) -> "np.linalg._linalg.EighResult":
+        """
+        Extract the reduced density matrix representing the state of a list
+        of qubits from a PyQRack simulator register.
+
+        Inputs:
+            qubits: A list of PyQRack qubits to extract the reduced density matrix for
+            tol: The tolerance for density matrix eigenvalues to be considered non-zero.
+        Outputs:
+            An eigh result containing the eigenvalues and eigenvectors of the reduced density matrix.
+        """
+        if len(qubits) == 0:
+            return np.linalg._linalg.EighResult(
+                eigenvalues=np.array([]), eigenvectors=np.array([]).reshape(0, 0)
+            )
+        sim_reg = qubits[0].sim_reg
+
+        if not all([x.sim_reg is sim_reg for x in qubits]):
+            raise ValueError("All qubits must be from the same simulator register.")
+        inds: tuple[int, ...] = tuple(qubit.addr for qubit in qubits)
+
+        return _pyqrack_reduced_density_matrix(inds, sim_reg, tol)
+
+    @classmethod
+    def reduced_density_matrix(
+        cls, qubits: list[PyQrackQubit] | IList[PyQrackQubit, Any], tol: float = 1e-12
+    ) -> np.ndarray:
+        """
+        Extract the reduced density matrix representing the state of a list
+        of qubits from a PyQRack simulator register.
+
+        Inputs:
+            qubits: A list of PyQRack qubits to extract the reduced density matrix for
+            tol: The tolerance for density matrix eigenvalues to be considered non-zero.
+        Outputs:
+            A dense 2^n x 2^n numpy array representing the reduced density matrix.
+        """
+        rdm = cls.quantum_state(qubits, tol)
+        return np.einsum(
+            "ax,x,bx", rdm.eigenvectors, rdm.eigenvalues, rdm.eigenvectors.conj()
+        )
+
 
 @dataclass
 class StackMemorySimulator(PyQrackSimulatorBase):

bloqade/pyqrack/task.py

@@ -2,6 +2,7 @@ from typing import TypeVar, ParamSpec, cast
 from dataclasses import dataclass
 
 from bloqade.task import AbstractSimulatorTask
+from bloqade.pyqrack.reg import QubitState, PyQrackQubit
 from bloqade.pyqrack.base import (
     MemoryABC,
     PyQrackInterpreter,
@@ -36,3 +37,17 @@ class PyQrackSimulatorTask(AbstractSimulatorTask[Param, RetType, MemoryType]):
         """Returns the state vector of the simulator."""
         self.run()
         return self.state.sim_reg.out_ket()
+
+    def qubits(self) -> list[PyQrackQubit]:
+        """Returns the qubits in the simulator."""
+        try:
+            N = self.state.sim_reg.num_qubits()
+            return [
+                PyQrackQubit(
+                    addr=i, sim_reg=self.state.sim_reg, state=QubitState.Active
+                )
+                for i in range(N)
+            ]
+        except AttributeError:
+            Warning("Task has not been run, there are no qubits!")
+            return []

bloqade/rewrite/rules/split_ifs.py

@@ -46,9 +46,13 @@ class SplitIfStmts(RewriteRule):
         if not isinstance(node, scf.IfElse):
             return RewriteResult()
 
+        # NOTE: only empty else bodies are allowed in valid QASM2
+        if not self._has_empty_else(node):
+            return RewriteResult()
+
         *stmts, yield_or_return = node.then_body.stmts()
 
-        if len(stmts) == 1:
+        if len(stmts) <= 1:
             return RewriteResult()
 
         is_yield = isinstance(yield_or_return, scf.Yield)
@@ -71,3 +75,16 @@ class SplitIfStmts(RewriteRule):
         node.delete()
 
         return RewriteResult(has_done_something=True)
+
+    def _has_empty_else(self, node: scf.IfElse) -> bool:
+        else_stmts = list(node.else_body.stmts())
+        if len(else_stmts) > 1:
+            return False
+
+        if len(else_stmts) == 0:
+            return True
+
+        if not isinstance(else_stmts[0], scf.Yield):
+            return False
+
+        return len(else_stmts[0].values) == 0

bloqade/squin/analysis/nsites/impls.py

@@ -81,6 +81,13 @@ class SquinOp(interp.MethodTable):
         op_sites = frame.get(stmt.op)
         return (op_sites,)
 
+    @interp.impl(op.stmts.PauliString)
+    def pauli_string(
+        self, interp: NSitesAnalysis, frame: interp.Frame, stmt: op.stmts.PauliString
+    ):
+        s = stmt.string
+        return (NumberSites(sites=len(s)),)
+
 
 @scf.dialect.register(key="op.nsites")
 class ScfSquinOp(ScfTypeInfer):

bloqade/squin/cirq/emit/op.py

@@ -9,6 +9,7 @@ from .runtime import (
     SnRuntime,
     SpRuntime,
     U3Runtime,
+    RotRuntime,
     KronRuntime,
     MultRuntime,
     ScaleRuntime,
@@ -123,3 +124,21 @@ class EmitCirqOpMethods(MethodTable):
         self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.PauliString
     ):
         return (PauliStringRuntime(stmt.string),)
+
+    @impl(op.stmts.Rot)
+    def rot(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.Rot):
+        axis_op: HermitianRuntime = frame.get(stmt.axis)
+        angle = frame.get(stmt.angle)
+
+        axis_name = str(axis_op.gate).lower()
+        return (RotRuntime(axis=axis_name, angle=angle),)
+
+    @impl(op.stmts.SqrtX)
+    def sqrt_x(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.SqrtX):
+        cirq_op = cirq.XPowGate(exponent=0.5)
+        return (UnitaryRuntime(cirq_op),)
+
+    @impl(op.stmts.SqrtY)
+    def sqrt_y(self, emit: EmitCirq, frame: EmitCirqFrame, stmt: op.stmts.SqrtY):
+        cirq_op = cirq.YPowGate(exponent=0.5)
+        return (UnitaryRuntime(cirq_op),)

bloqade/squin/cirq/emit/runtime.py

@@ -240,3 +240,18 @@ class PauliStringRuntime(OperatorRuntimeABC):
             qbit: pauli_label for (qbit, pauli_label) in zip(qubits, self.string)
         }
         return [cirq.PauliString(pauli_mapping)]
+
+
+@dataclass
+class RotRuntime(OperatorRuntimeABC):
+    axis: str
+    angle: float
+
+    def num_qubits(self) -> int:
+        return 1
+
+    def unsafe_apply(
+        self, qubits: Sequence[cirq.Qid], adjoint: bool = False
+    ) -> list[cirq.Operation]:
+        rot = getattr(cirq, "R" + self.axis.lower())(rads=self.angle)
+        return [rot(*qubits)]

bloqade/squin/noise/rewrite.py

@@ -3,7 +3,7 @@ import itertools
 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 (
@@ -11,6 +11,7 @@ from .stmts import (
     QubitLoss,
     Depolarize,
     PauliError,
+    Depolarize2,
     NoiseChannel,
     TwoQubitPauliChannel,
     SingleQubitPauliChannel,
@@ -58,6 +59,18 @@ class _RewriteNoiseStmts(RewriteRule):
     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 @@ class _RewriteNoiseStmts(RewriteRule):
             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 @@ class _RewriteNoiseStmts(RewriteRule):
             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 @@ class _RewriteNoiseStmts(RewriteRule):
 
         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):

{bloqade_circuit-0.6.6.dist-info → bloqade_circuit-0.6.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bloqade-circuit
-Version: 0.6.6
+Version: 0.6.8
 Summary: The software development toolkit for neutral atom arrays.
 Author-email: Roger-luo <rluo@quera.com>, kaihsin <khwu@quera.com>, weinbe58 <pweinberg@quera.com>, johnzl-777 <jlong@quera.com>
 License-File: LICENSE

{bloqade_circuit-0.6.6.dist-info → bloqade_circuit-0.6.8.dist-info}/RECORD

@@ -23,10 +23,10 @@ bloqade/cirq_utils/noise/model.py,sha256=06Y_BLChOA-PhhAJcWLSgLVAAJoNjOrAujL1YCw
 bloqade/cirq_utils/noise/transform.py,sha256=tvDt4WMLM8dKPME51y0_peSZk2-jKmjq0urOxm0lWuQ,2309
 bloqade/pyqrack/__init__.py,sha256=lonTS-luJkTVujCCtgdZRC12V7FQdoFcozAI-byXwN0,810
 bloqade/pyqrack/base.py,sha256=g0GRlEgyJ_P8z-lR8RK2CAuRTj6KPfglKX0iwrgg4DM,4408
-bloqade/pyqrack/device.py,sha256=40vduanEgA26GAW3buHoRpyqPA0xUt2tONY3w5JeH5s,7524
+bloqade/pyqrack/device.py,sha256=-zlr1lSzOvcj5l28nnevy1oMYct79DTOdLYyGaT2Yco,11567
 bloqade/pyqrack/reg.py,sha256=uTL07CT1R0xUsInLmwU9YuuNdV6lV0lCs1zhdUz1qIs,1660
 bloqade/pyqrack/target.py,sha256=c78VtLWAiDNp_0sXwvVzhaEoeFsr1fUVsupxWuo6p3s,3661
-bloqade/pyqrack/task.py,sha256=ydVso4wZQ1iIzgSBG0lDxClPTLcsZfkbnZk_1qFV95o,991
+bloqade/pyqrack/task.py,sha256=hzMueE03C-MR_-ti_5Z0lqFZcHDVCxmQHFaX9tPNm2o,1522
 bloqade/pyqrack/noise/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bloqade/pyqrack/noise/native.py,sha256=BmRlRzqCzvSgfSsDlIjuB8L0gx4uaPqmPyfJ5-Ole6M,2751
 bloqade/pyqrack/qasm2/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -124,7 +124,7 @@ bloqade/rewrite/passes/canonicalize_ilist.py,sha256=LmX09aPjagT6NzU-M87V8kZaCGkY
 bloqade/rewrite/rules/__init__.py,sha256=3e1Z5T3INqNtP6OU0Vivu_SdMOR_2KDixeA0Yjoddgg,82
 bloqade/rewrite/rules/flatten_ilist.py,sha256=QoIxMaBXSlatpWzi5s_MAPnV3bV3GeoWc31RBw0WQ3s,1465
 bloqade/rewrite/rules/inline_getitem_ilist.py,sha256=uIXQRCsr3_GPMciDT4ghI-ezhQmkDcGcC6pguABPUVw,875
-bloqade/rewrite/rules/split_ifs.py,sha256=Nm4lpEUHZcnCeewIld0tt7UuGO69LiBGl7Uybuwissw,2119
+bloqade/rewrite/rules/split_ifs.py,sha256=KhwvUx-oBrBO2F1j-J5kwNbdnrnEZcZtJflzyQm-UOI,2613
 bloqade/squin/__init__.py,sha256=b7ZD69ql9GriIPxN6JhWxANJVzuIJh_r-gC24wsW1mM,621
 bloqade/squin/_typeinfer.py,sha256=bilWfC6whTMwewFCqDgB6vDHZsgXPr3azNOYqqnvtB4,780
 bloqade/squin/gate.py,sha256=tCnjfrSVsXHX7VxkEulZ2SQS5ydtmON8QlcGibM6c2I,4028
@@ -139,19 +139,19 @@ bloqade/squin/analysis/address_impl.py,sha256=eMEGlkw88ozj8ZGHcfQ0qW2K_dNxYxnhwz
 bloqade/squin/analysis/schedule.py,sha256=buuC4bFuLuaSDK2BZfkRkh8ZdNicz9HkEv3FAnsDViE,7880
 bloqade/squin/analysis/nsites/__init__.py,sha256=RlQg7ivczXCXG5lMeL3ipYKj2oJKC4THu8orYf5PBYs,263
 bloqade/squin/analysis/nsites/analysis.py,sha256=rIe1RU1MZRItcE2aB8DYahLrv73HfD3IHCX3E_EGQ1c,1773
-bloqade/squin/analysis/nsites/impls.py,sha256=bVqO4E2hDzfYWwSG0pfj2j8oT1m0C3b42LdSHr4HQlo,2874
+bloqade/squin/analysis/nsites/impls.py,sha256=wSNWjNmgwCP35FgmreyLKmRYedxlebWi7LhsEq9jPs4,3097
 bloqade/squin/analysis/nsites/lattice.py,sha256=ruh0808SHtj3ecuT-C3AZTsLY2j3DRhtezGiTZvcuVs,942
 bloqade/squin/cirq/__init__.py,sha256=7OYYboSl5lPwdWOKk4AJgm1s1lBX_WAstVqPfaynSv8,10156
 bloqade/squin/cirq/lowering.py,sha256=F0_skv9lORxUFrhbNaN2ZQpqGnhPyslHt5E92uta5BQ,17959
 bloqade/squin/cirq/emit/emit_circuit.py,sha256=JVFXiaSB7A9kamRwCjLqs03Sel4PVCBT-6kRNRWX-jo,4393
 bloqade/squin/cirq/emit/noise.py,sha256=rESjGC_66s2Y4FwwYda4rY3mYHYjbqLlKE_vnqpZDYI,1534
-bloqade/squin/cirq/emit/op.py,sha256=z54NP5KqMxffXeFGWamEzvunpTNrxmYuluurk4j2-ps,4000
+bloqade/squin/cirq/emit/op.py,sha256=gWMLghj65yQeAbwy2E0YhRLbzM_DkYBlvXKfsgY26fc,4712
 bloqade/squin/cirq/emit/qubit.py,sha256=IegghRU5E1urd0ddV1X2Fh_aEWAPmUPZLVj9AafN2Cg,1930
-bloqade/squin/cirq/emit/runtime.py,sha256=dH7JSMt2mALPhVFjmZETQzvnTUQ3BFY5poe0YZpM5vQ,6819
+bloqade/squin/cirq/emit/runtime.py,sha256=RA_3C_0qkwQCcztalmfOqJe-c4DZgEqWn2ArNjyycZ4,7170
 bloqade/squin/noise/__init__.py,sha256=xST2qojx6ZApRoiKIXOtifDzSpTZgo-67ja309FFvWw,364
 bloqade/squin/noise/_dialect.py,sha256=2IR98J-lXm5Y3srP9g-FD4JC-qTq2seureM6mKKq1xg,63
 bloqade/squin/noise/_wrapper.py,sha256=P8fkr1_2U47PtvqnQqeTI7VzThNIK17cNh2QX6ABh_w,815
-bloqade/squin/noise/rewrite.py,sha256=-IqFfDGnhuaFI-9b6PXjhSuiXFM1C5Qu0ibL5GvZldI,3917
+bloqade/squin/noise/rewrite.py,sha256=8v8AAi2G-ZatZ6dTpqc88hCBMMoWz2TK3nwkLkwk7Fo,5099
 bloqade/squin/noise/stmts.py,sha256=F8AsDp2xsLez9vkSamDW985MUCubz3TMRrE0L745GmI,1875
 bloqade/squin/op/__init__.py,sha256=6JOjPdzc6RKO4299ZFz4Jk-wtVyPlGTkakYewHBueXw,841
 bloqade/squin/op/_dialect.py,sha256=66G1IYqmsqUEaCTyUqn2shSHmGYduiTU8GfDXcoMvw4,55
@@ -230,7 +230,7 @@ bloqade/visual/animation/runtime/atoms.py,sha256=EmjxhujLiHHPS_HtH_B-7TiqeHgvW5u
 bloqade/visual/animation/runtime/ppoly.py,sha256=JB9IP53N1w6adBJEue6J5Nmj818Id9JvrlgrmiQTU1I,1385
 bloqade/visual/animation/runtime/qpustate.py,sha256=rlmxQeJSvaohXrTpXQL5y-NJcpvfW33xPaYM1slv7cc,4270
 bloqade/visual/animation/runtime/utils.py,sha256=ju9IzOWX-vKwfpqUjlUKu3Ssr_UFPFFq-tzH_Nqyo_c,1212
-bloqade_circuit-0.6.6.dist-info/METADATA,sha256=uDJH5m1IwfBujpFvXFYjP89i0aHRbwSMEc6vhx9xwXw,3849
-bloqade_circuit-0.6.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bloqade_circuit-0.6.6.dist-info/licenses/LICENSE,sha256=S5GIJwR6QCixPA9wryYb44ZEek0Nz4rt_zLUqP05UbU,13160
-bloqade_circuit-0.6.6.dist-info/RECORD,,
+bloqade_circuit-0.6.8.dist-info/METADATA,sha256=8pPrKcmeLw4CJWHLi6DJuYS-xhwkGMOKZLtzzvxFAa4,3849
+bloqade_circuit-0.6.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bloqade_circuit-0.6.8.dist-info/licenses/LICENSE,sha256=S5GIJwR6QCixPA9wryYb44ZEek0Nz4rt_zLUqP05UbU,13160
+bloqade_circuit-0.6.8.dist-info/RECORD,,