bloqade-circuit 0.1.0__py3-none-any.whl

bloqade/qbraid/lowering.py

@@ -0,0 +1,324 @@
+from typing import Dict, List, Tuple, Sequence
+from dataclasses import field, dataclass
+
+from kirin import ir, types, passes
+from kirin.dialects import func, ilist
+
+from bloqade import noise, qasm2
+from bloqade.qbraid import schema
+from bloqade.qasm2.dialects import glob, parallel
+
+
+@ir.dialect_group(
+    [func, qasm2.core, qasm2.uop, parallel, glob, qasm2.expr, noise.native, ilist]
+)
+def qbraid_noise(
+    self,
+):
+    fold_pass = passes.Fold(self)
+    typeinfer_pass = passes.TypeInfer(self)
+
+    def run_pass(
+        method: ir.Method,
+        *,
+        fold: bool = True,
+    ):
+        method.verify()
+
+        if fold:
+            fold_pass(method)
+
+        typeinfer_pass(method)
+        method.code.verify_type()
+
+    return run_pass
+
+
+@dataclass
+class Lowering:
+    qubit_list: List[ir.SSAValue] = field(init=False, default_factory=list)
+    qubit_id_map: Dict[int, ir.SSAValue] = field(init=False, default_factory=dict)
+    bit_id_map: Dict[int, ir.SSAValue] = field(init=False, default_factory=dict)
+    block_list: List[ir.Statement] = field(init=False, default_factory=list)
+
+    def lower(
+        self,
+        sym_name: str,
+        noise_model: schema.NoiseModel,
+        return_qreg: bool = False,
+    ) -> ir.Method:
+        """Lower the noise model to a method.
+
+        Args:
+            name (str): The name of the method to generate.
+            return_qreg (bool): Use the quantum register as the return value.
+
+        Returns:
+            Method: The generated kirin method.
+
+        """
+        self.process_noise_model(noise_model, return_qreg)
+        block = ir.Block(stmts=self.block_list)
+        ret_type = qasm2.types.QRegType if return_qreg else qasm2.types.CRegType
+        block.args.append_from(types.MethodType[[], ret_type], name=f"{sym_name}_self")
+        region = ir.Region(block)
+        func_stmt = func.Function(
+            sym_name=sym_name,
+            signature=func.Signature(inputs=(), output=qasm2.types.QRegType),
+            body=region,
+        )
+
+        mt = ir.Method(
+            mod=None,
+            py_func=None,
+            sym_name=sym_name,
+            dialects=qbraid_noise,
+            code=func_stmt,
+            arg_names=[],
+        )
+        qbraid_noise.run_pass(mt)  # type: ignore
+        return mt
+
+    def process_noise_model(self, noise_model: schema.NoiseModel, return_qreg: bool):
+        num_qubits = self.lower_number(noise_model.num_qubits)
+
+        reg = qasm2.core.QRegNew(num_qubits)
+        creg = qasm2.core.CRegNew(num_qubits)
+        self.block_list.append(reg)
+        self.block_list.append(creg)
+
+        for idx_value, qubit in enumerate(noise_model.all_qubits):
+            idx = self.lower_number(idx_value)
+            self.block_list.append(qubit_stmt := qasm2.core.QRegGet(reg.result, idx))
+            self.block_list.append(bit_stmt := qasm2.core.CRegGet(creg.result, idx))
+
+            self.qubit_id_map[qubit] = qubit_stmt.result
+            self.bit_id_map[qubit] = bit_stmt.result
+            self.qubit_list.append(qubit_stmt.result)
+
+        for gate_event in noise_model.gate_events:
+            self.process_gate_event(gate_event)
+
+        if return_qreg:
+            self.block_list.append(func.Return(reg.result))
+        else:
+            self.block_list.append(func.Return(creg.result))
+
+    def process_gate_event(self, node: schema.GateEvent):
+        self.lower_atom_loss(node.error.survival_prob)
+
+        if isinstance(node.operation, schema.CZ):
+            assert isinstance(node.error, schema.CZError), "Only CZError is supported"
+            self.process_cz_pauli_error(node.operation.participants, node.error)
+            self.lower_cz_gates(node.operation)
+        else:
+
+            error = node.error
+            assert isinstance(
+                error, schema.SingleQubitError
+            ), "Only SingleQubitError is supported"
+            self.lower_pauli_errors(error.operator_error)
+
+            operation = node.operation
+            assert isinstance(
+                operation,
+                (
+                    schema.GlobalW,
+                    schema.LocalW,
+                    schema.GlobalRz,
+                    schema.LocalRz,
+                    schema.Measurement,
+                ),
+            ), f"Only W and Rz gates are supported, found {type(operation)}.__name__"
+
+            if isinstance(operation, schema.GlobalW):
+                self.lower_w_gates(
+                    tuple(self.qubit_id_map.keys()), operation.theta, operation.phi
+                )
+            elif isinstance(operation, schema.LocalW):
+                self.lower_w_gates(
+                    operation.participants, operation.theta, operation.phi
+                )
+            elif isinstance(operation, schema.GlobalRz):
+                self.lower_rz_gates(tuple(self.qubit_id_map.keys()), operation.phi)
+            elif isinstance(operation, schema.LocalRz):
+                self.lower_rz_gates(operation.participants, operation.phi)
+            elif isinstance(operation, schema.Measurement):
+                self.lower_measurement(operation)
+
+    def process_cz_pauli_error(
+        self,
+        participants: Tuple[Tuple[int] | Tuple[int, int], ...],
+        node: schema.CZError[schema.PauliErrorModel],
+    ):
+
+        storage_error = node.storage_error
+        single_error = node.single_error
+        entangled_error = node.entangled_error
+
+        assert isinstance(
+            storage_error, schema.PauliErrorModel
+        ), "Only PauliErrorModel is supported"
+        assert isinstance(
+            single_error, schema.PauliErrorModel
+        ), "Only PauliErrorModel is supported"
+        assert isinstance(
+            entangled_error, schema.PauliErrorModel
+        ), "Only PauliErrorModel is supported"
+
+        self.lower_pauli_errors(storage_error)
+
+        single_error_dict = dict(single_error.errors)
+        entangled_error_dict = dict(entangled_error.errors)
+
+        single_layers = {}
+        paired_layers = {}
+        unpaired_layers = {}
+
+        for participant in participants:
+            if len(participant) == 1:
+                p_single = single_error_dict[participant[0]]
+                single_layers.setdefault(p_single, []).append(participant[0])
+            elif len(participant) == 2:
+                p_ctrl = entangled_error_dict[participant[0]]
+                up_ctrl = single_error_dict[participant[0]]
+                p_qarg = entangled_error_dict[participant[1]]
+                up_qarg = single_error_dict[participant[1]]
+                paired_layers.setdefault((p_ctrl, p_qarg), []).append(participant)
+                unpaired_layers.setdefault((up_ctrl, up_qarg), []).append(participant)
+
+        for (px, py, pz), qubits in single_layers.items():
+            self.block_list.append(
+                qargs := ilist.New(values=tuple(self.qubit_id_map[q] for q in qubits))
+            )
+            self.block_list.append(
+                noise.native.PauliChannel(px=px, py=py, pz=pz, qargs=qargs.result)
+            )
+
+        for (p_ctrl, p_qarg), qubits in paired_layers.items():
+            ctrls, qargs = list(zip(*qubits))
+            self.block_list.append(
+                ctrls := ilist.New(values=tuple(self.qubit_id_map[q] for q in ctrls))
+            )
+            self.block_list.append(
+                qargs := ilist.New(values=tuple(self.qubit_id_map[q] for q in qargs))
+            )
+            self.block_list.append(
+                noise.native.CZPauliChannel(
+                    paired=True,
+                    px_ctrl=p_ctrl[0],
+                    py_ctrl=p_ctrl[1],
+                    pz_ctrl=p_ctrl[2],
+                    px_qarg=p_qarg[0],
+                    py_qarg=p_qarg[1],
+                    pz_qarg=p_qarg[2],
+                    ctrls=ctrls.result,
+                    qargs=qargs.result,
+                )
+            )
+
+        for (p_ctrl, p_qarg), qubits in unpaired_layers.items():
+            ctrls, qargs = list(zip(*qubits))
+            self.block_list.append(
+                ctrls := ilist.New(values=tuple(self.qubit_id_map[q] for q in ctrls))
+            )
+            self.block_list.append(
+                qargs := ilist.New(values=tuple(self.qubit_id_map[q] for q in qargs))
+            )
+            self.block_list.append(
+                noise.native.CZPauliChannel(
+                    paired=False,
+                    px_ctrl=p_ctrl[0],
+                    py_ctrl=p_ctrl[1],
+                    pz_ctrl=p_ctrl[2],
+                    px_qarg=p_qarg[0],
+                    py_qarg=p_qarg[1],
+                    pz_qarg=p_qarg[2],
+                    ctrls=ctrls.result,
+                    qargs=qargs.result,
+                )
+            )
+
+    def lower_cz_gates(self, node: schema.CZ):
+        ctrls, qargs = list(zip(*(p for p in node.participants if len(p) == 2)))
+        self.block_list.append(
+            ctrls := ilist.New(values=tuple(self.qubit_id_map[q] for q in ctrls))
+        )
+        self.block_list.append(
+            qargs := ilist.New(values=tuple(self.qubit_id_map[q] for q in qargs))
+        )
+        self.block_list.append(parallel.CZ(ctrls=ctrls.result, qargs=qargs.result))
+
+    def lower_w_gates(self, participants: Sequence[int], theta: float, phi: float):
+        self.block_list.append(
+            qargs := ilist.New(values=tuple(self.qubit_id_map[q] for q in participants))
+        )
+        self.block_list.append(
+            parallel.UGate(
+                theta=self.lower_full_turns(theta),
+                phi=self.lower_full_turns(phi + 0.5),
+                lam=self.lower_full_turns(-(0.5 + phi)),
+                qargs=qargs.result,
+            )
+        )
+
+    def lower_rz_gates(self, participants: Tuple[int, ...], phi: float):
+        self.block_list.append(
+            qargs := ilist.New(values=tuple(self.qubit_id_map[q] for q in participants))
+        )
+        self.block_list.append(
+            parallel.RZ(theta=self.lower_full_turns(phi), qargs=qargs.result)
+        )
+
+    def lower_pauli_errors(self, operator_error: schema.PauliErrorModel):
+        assert isinstance(
+            operator_error, schema.PauliErrorModel
+        ), "Only PauliErrorModel is supported"
+
+        layers = {}
+
+        for qubit_num, pauli_errors in operator_error.errors:
+            layers.setdefault(pauli_errors, []).append(qubit_num)
+
+        for (px, py, pz), qubits in layers.items():
+            self.block_list.append(
+                qargs := ilist.New(values=tuple(self.qubit_id_map[q] for q in qubits))
+            )
+            self.block_list.append(
+                noise.native.PauliChannel(px=px, py=py, pz=pz, qargs=qargs.result)
+            )
+
+    def lower_measurement(self, operation: schema.Measurement):
+        for participant in operation.participants:
+            qubit = self.qubit_id_map[participant]
+            bit = self.bit_id_map[participant]
+            self.block_list.append(qasm2.core.Measure(qarg=qubit, carg=bit))
+
+    def lower_atom_loss(self, survival_probs: Tuple[float, ...]):
+        layers = {}
+
+        for qubit_num, survival_prob in zip(self.qubit_list, survival_probs):
+            layers.setdefault(survival_prob, []).append(qubit_num)
+
+        for survival_prob, qubits in layers.items():
+            self.block_list.append(qargs := ilist.New(values=qubits))
+            self.block_list.append(
+                noise.native.AtomLossChannel(prob=survival_prob, qargs=qargs.result)
+            )
+
+    def lower_number(self, value: float | int) -> ir.SSAValue:
+        if isinstance(value, int):
+            stmt = qasm2.expr.ConstInt(value=value)
+        else:
+            stmt = qasm2.expr.ConstFloat(value=value)
+
+        self.block_list.append(stmt)
+        return stmt.result
+
+    def lower_full_turns(self, value: float) -> ir.SSAValue:
+        const_pi = qasm2.expr.ConstPI()
+        self.block_list.append(const_pi)
+        turns = self.lower_number(2 * value)
+        mul = qasm2.expr.Mul(const_pi.result, turns)
+        self.block_list.append(mul)
+        return mul.result

bloqade/qbraid/schema.py

@@ -0,0 +1,252 @@
+from typing import List, Tuple, Union, Generic, Literal, TypeVar
+
+from pydantic import Field, BaseModel
+
+
+class Operation(BaseModel, frozen=True, extra="forbid"):
+    """Base class for operations."""
+
+    op_type: str = Field(init=False)
+
+
+class CZ(Operation):
+    """A CZ gate operation.
+
+    Fields:
+        op_type (str): The type of operation (Literal["CZ"]).
+        participants (Tuple[Union[Tuple[int], Tuple[int, int]], ...]): The qubit indices that are participating in the CZ gate.
+
+    """
+
+    op_type: Literal["CZ"] = Field(init=False, default="CZ")
+    participants: Tuple[Union[Tuple[int], Tuple[int, int]], ...]
+
+
+class GlobalRz(Operation):
+    """GlobalRz operation.
+
+    Fields:
+        op_type (str): The type of operation (Literal["GlobalRz"]).
+        phi (float): The angle of rotation.
+    """
+
+    op_type: Literal["GlobalRz"] = Field(init=False, default="GlobalRz")
+    phi: float
+
+
+class GlobalW(Operation):
+    """GlobalW operation.
+
+    Fields:
+        op_type (str): The type of operation (Literal["GlobalW"]).
+        theta (float): The angle of rotation.
+        phi (float): The angle of rotation.
+    """
+
+    op_type: Literal["GlobalW"] = Field(init=False, default="GlobalW")
+    theta: float
+    phi: float
+
+
+class LocalRz(Operation):
+    """LocalRz operation.
+
+    Fields:
+        op_type (str): The type of operation (Literal["LocalRz"]).
+        participants (Tuple[int, ...]): The qubit indices that are participating in the local Rz gate.
+        phi (float): The angle of rotation.
+
+    """
+
+    op_type: Literal["LocalRz"] = Field(init=False, default="LocalRz")
+    participants: Tuple[int, ...]
+    phi: float
+
+
+class LocalW(Operation):
+    """LocalW operation.
+
+    Fields:
+        op_type (str): The type of operation (Literal["LocalW"]).
+        participants (Tuple[int, ...]): The qubit indices that are participating in the local W gate.
+        theta (float): The angle of rotation.
+        phi (float): The angle of rotation.
+
+    """
+
+    op_type: Literal["LocalW"] = Field(init=False, default="LocalW")
+    participants: Tuple[int, ...]
+    theta: float
+    phi: float
+
+
+class Measurement(Operation):
+    """Measurement operation.
+
+    Fields:
+        op_type (str): The type of operation (Literal["Measurement"]).
+        measure_tag (str): The tag to use for the measurement.
+        participants (Tuple[int, ...]): The qubit indices that are participating in the measurement.
+
+    """
+
+    op_type: Literal["Measurement"] = Field(init=False, default="Measurement")
+    measure_tag: str = Field(default="m")
+    participants: Tuple[int, ...]
+
+
+OperationType = TypeVar(
+    "OperationType", bound=Union[CZ, GlobalRz, GlobalW, LocalRz, LocalW, Measurement]
+)
+
+
+class ErrorModel(BaseModel, frozen=True, extra="forbid"):
+    """Base class for error models."""
+
+    error_model_type: str = Field(init=False)
+
+
+class PauliErrorModel(ErrorModel):
+    """Pauli error model.
+
+    Fields:
+        error_model_type (str): The type of error model (Literal["PauliNoise"]).
+        errors (Tuple[Tuple[int, Tuple[float, float, float]], ...]): The qubit indices and the error rates for each qubit.
+
+    """
+
+    error_model_type: Literal["PauliNoise"] = Field(default="PauliNoise", init=False)
+    errors: Tuple[Tuple[int, Tuple[float, float, float]], ...] = Field(
+        default_factory=tuple
+    )
+
+
+ErrorModelType = TypeVar("ErrorModelType", bound=PauliErrorModel)
+
+
+class ErrorOperation(BaseModel, Generic[ErrorModelType], frozen=True, extra="forbid"):
+    """Base class for error operations."""
+
+    error_type: str = Field(init=False)
+    survival_prob: Tuple[float, ...]
+
+
+class CZError(ErrorOperation[ErrorModelType]):
+    """CZError operation.
+
+    Fields:
+        survival_prob (Tuple[float, ...]): The survival probabilities for each qubit.
+        error_type (str): The type of error (Literal["CZError"]).
+        storage_error (ErrorModelType): The error model for storage.
+        entangled_error (ErrorModelType): The error model for entangled qubits.
+        single_error (ErrorModelType): The error model for single qubits.
+
+    """
+
+    error_type: Literal["CZError"] = Field(default="CZError", init=False)
+    storage_error: ErrorModelType
+    entangled_error: ErrorModelType
+    single_error: ErrorModelType
+
+
+class SingleQubitError(ErrorOperation[ErrorModelType]):
+    """SingleQubitError operation.
+
+    Fields:
+        survival_prob (Tuple[float, ...]): The survival probabilities for each qubit.
+        error_type (str): The type of error (Literal["SingleQubitError"]).
+        operator_error (ErrorModelType): The error model for the single qubit.
+
+    """
+
+    error_type: Literal["SingleQubitError"] = Field(
+        default="SingleQubitError", init=False
+    )
+    operator_error: ErrorModelType
+
+
+class GateEvent(BaseModel, Generic[ErrorModelType], frozen=True, extra="forbid"):
+    """A gate event.
+
+    Fields:
+        error (Union[SingleQubitError[ErrorModelType], CZError[ErrorModelType]]): The error model for the gate event.
+        operation (OperationType): The operation for the gate event.
+
+    """
+
+    error: Union[SingleQubitError[ErrorModelType], CZError[ErrorModelType]] = Field(
+        union_mode="left_to_right", discriminator="error_type"
+    )
+    operation: OperationType = Field(
+        union_mode="left_to_right", discriminator="op_type"
+    )
+
+    def __pydantic_post_init__(self):
+        assert (isinstance(self.operation, CZ) and isinstance(self.error, CZError)) or (
+            not isinstance(self.operation, CZ)
+            and isinstance(self.error, SingleQubitError)
+        ), "Operation and error must be of the same type"
+
+
+class NoiseModel(BaseModel, Generic[ErrorModelType], extra="forbid"):
+    """Noise model for a circuit.
+
+    Fields:
+        all_qubits (Tuple[int, ...]): The qubit indices for the noise model.
+        gate_events (List[GateEvent[ErrorModelType]]): The gate events for the noise model.
+
+    """
+
+    all_qubits: Tuple[int, ...] = Field(default_factory=tuple)
+    gate_events: List[GateEvent[ErrorModelType]] = Field(default_factory=list)
+
+    @property
+    def num_qubits(self) -> int:
+        """Return the number of qubits in the noise model."""
+        return len(self.all_qubits)
+
+    def __add__(self, other: "NoiseModel") -> "NoiseModel":
+        if not isinstance(other, NoiseModel):
+            raise ValueError(f"Cannot add {type(other)} to Circuit")
+
+        if self.all_qubits != other.all_qubits:
+            raise ValueError("Circuits must have the same number of qubits")
+
+        return NoiseModel(
+            all_qubits=self.all_qubits,
+            gate_events=self.gate_events + other.gate_events,
+        )
+
+    def lower_noise_model(self, sym_name: str, return_qreg: bool = False):
+        """Lower the noise model to a method.
+
+        Args:
+            sym_name (str): The name of the method to generate.
+            return_qreg (bool): Whether to return the quantum register after the method
+                has completed execution. Useful for obtaining the full state vector.
+
+        Returns:
+            Method: The generated kirin method.
+
+        """
+        from bloqade.qbraid.lowering import Lowering
+
+        return Lowering().lower(sym_name, self, return_qreg)
+
+    def decompiled_circuit(self) -> str:
+        """Clean the circuit of noise.
+
+        Returns:
+            str: The decompiled circuit from hardware execution.
+
+        """
+        from bloqade.noise import native
+        from bloqade.qasm2.emit import QASM2
+        from bloqade.qasm2.passes import glob, parallel
+
+        mt = self.lower_noise_model("method")
+
+        native.RemoveNoisePass(mt.dialects)(mt)
+        parallel.ParallelToUOp(mt.dialects)(mt)
+        glob.GlobalToUOP(mt.dialects)(mt)
+        return QASM2(qelib1=True).emit_str(mt)

bloqade/qbraid/simulation_result.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2024, QuEra Computing Inc.
+# All rights reserved.
+
+from io import StringIO
+from typing import Any, Dict, Optional
+from dataclasses import dataclass
+
+import pandas as pd
+from pandas import DataFrame
+
+from bloqade.visual.animation.runtime import qpustate as vis_qpustate
+
+from .schema import NoiseModel
+
+
+@dataclass
+class QuEraSimulationResult:
+    """Results of the QuEra hardware model simulation.
+
+    Fields:
+        flair_visual_version (str): The version of the Flair Visual package used to generate the simulation result.
+        counts (dict[str, int]): The measurement bitstrings of the simulation.
+        logs (DataFrame): Grainular logs events of what happened to each atom during the simulation.
+        atom_animation_state (vis_qpustate.AnimateQPUState): Object used to play back atom trajectories and events during the simulation.
+        noise_model (NoiseModel): The noise model used in the simulation.
+
+    """
+
+    flair_visual_version: str
+    counts: dict[str, int]
+    logs: DataFrame
+    atom_animation_state: vis_qpustate.AnimateQPUState
+    noise_model: NoiseModel
+
+    @classmethod
+    def from_json(cls, json: dict) -> "QuEraSimulationResult":
+        """deserialize the object from a JSON serializable dictionary."""
+        flair_visual_version = json["flair_visual_version"]
+        counts = json["counts"]
+        logs = pd.read_csv(StringIO(json["logs"]), index_col=0)
+        atom_animation_state = vis_qpustate.AnimateQPUState.from_json(
+            json["atom_animation_state"]
+        )
+        noise_model = NoiseModel(**json["noise_model"])
+
+        return cls(
+            flair_visual_version=flair_visual_version,
+            counts=counts,
+            logs=logs,
+            atom_animation_state=atom_animation_state,
+            noise_model=noise_model,
+        )
+
+    def to_json(self) -> Dict[str, Any]:
+        """Turn the object into a JSON serializable dictionary."""
+        return {
+            "flair_visual_version": self.flair_visual_version,
+            "counts": self.counts,
+            "logs": self.logs.to_csv(),
+            "atom_animation_state": self.atom_animation_state.to_json(),
+            "noise_model": self.noise_model.model_dump(mode="json"),
+        }
+
+    def animate(
+        self,
+        dilation_rate: float = 0.05,
+        fps: int = 30,
+        gate_display_dilation: float = 1.0,
+        save_mpeg: bool = False,
+        filename: str = "vqpu_animation",
+        start_block: int = 0,
+        n_blocks: Optional[int] = None,
+    ):
+        """animate the qpu state
+
+        Args:
+            dilation_rate (float): Conversion factor from the qpu time to animation time units. when dilation_rate=1.0, 1 (us) of qpu exec time corresponds to 1 second of animation time.
+            fps (int, optional): frame per second. Defaults to 30.
+            gate_display_dilation (float, optional): relative dilation rate of a gate event. Defaults to 1. When setting higher value, the gate event will be displayed longer.
+            save_mpeg (bool, optional): Save as mpeg. Defaults to False.
+            filename (str, optional): The file name of saved mpeg file. Defaults to "vqpu_animation". When `save_mpeg` is False, this argument is ignored.
+            start_block (int, optional): The start block to animate. Defaults to 0.
+            n_blocks (int, optional): number of blocks to animate. Defaults to None. When None, animate all blocks after `start_block`.
+        Returns:
+            ani: matplotlib animation object
+        """
+        from bloqade.visual.animation.animate import animate_qpu_state
+
+        ani = animate_qpu_state(
+            state=self.atom_animation_state,
+            dilation_rate=dilation_rate,
+            fps=fps,
+            gate_display_dilation=gate_display_dilation,
+            start_block=start_block,
+            n_blocks=n_blocks,
+            save_mpeg=save_mpeg,
+            filename=filename,
+        )
+        return ani