bloqade-circuit 0.1.0__py3-none-any.whl

bloqade/visual/animation/base.py

@@ -0,0 +1,346 @@
+import itertools
+import dataclasses
+from abc import abstractmethod
+from enum import Enum
+from typing import Any, Dict, List, Tuple, Callable, Optional
+
+import matplotlib.patches as mpatches
+import matplotlib.collections as mcollections
+
+from .gate_event import GateEvent
+
+_gate_color_code = {
+    "GlobalCZGate": "tab:red",
+    "TopHatCZGate": "tab:red",
+    "GlobalHyperfineRotation": "tab:blue",
+    "GlobalHyperfineZRotation": "tab:green",
+    "LocalHyperfineZRotation": "tab:green",
+    "LocalHyperfineZRotationPos": "tab:green",
+    "LocalHyperfineRotation": "tab:blue",
+    "LocalHyperfineRotationPos": "tab:blue",
+}
+
+
+class quera_color_code(str, Enum):
+    # TODO: for python 3.11+, replace traits with StrEnum
+    purple = "#6437FF"
+    red = "#C2477F"
+    yellow = "#EADD08"
+
+
+@dataclasses.dataclass
+class GateArtist:
+    mpl_ax: Any
+
+    @abstractmethod
+    def clear_data(self) -> None:
+        pass
+
+    @abstractmethod
+    def get_artists(self) -> Tuple[Any]:
+        pass
+
+
+@dataclasses.dataclass(init=False)
+class GlobalGateArtist(GateArtist):
+    mpl_obj: mpatches.Rectangle
+
+    def __init__(self, mpl_ax: Any, xmin, ymin, width, height, color):
+        super().__init__(mpl_ax)
+        rc = mpatches.Rectangle(
+            [xmin, ymin], width, height, color=color, alpha=0.6, visible=False
+        )
+        mpl_ax.add_patch(rc)
+        self.mpl_obj = rc
+
+    def clear_data(self) -> None:
+        self.mpl_obj.set_width(0)
+        self.mpl_obj.set_height(0)
+        self.mpl_obj.set_xy([0, 0])
+
+    def get_artists(self) -> Tuple[Any]:
+        return (self.mpl_obj,)
+
+    def set_visible(self, visible: bool):
+        self.mpl_obj.set_visible(visible)
+
+
+@dataclasses.dataclass(init=False)
+class RowRegionGateArtist(GateArtist):
+    """
+    A row region gate artist object.
+
+    bound box is [y_origin - width/2, y_origin + width/2]
+    """
+
+    mpl_obj: mpatches.Rectangle
+    mpl_obj_keepout_top: mpatches.Rectangle
+    mpl_obj_keepout_btm: mpatches.Rectangle
+    width: float
+    xmin: float
+
+    def __init__(
+        self, mpl_ax: Any, xmin, width, ymin, ymin_keepout, ymax, ymax_keepout, color
+    ):
+        super().__init__(mpl_ax)
+        self.width = width
+        self.xmin = xmin
+        rc_btm = mpatches.Rectangle(
+            [xmin, ymin_keepout],
+            width,
+            ymin - ymin_keepout,
+            color=color,
+            alpha=0.3,
+            visible=False,
+        )
+        mpl_ax.add_patch(rc_btm)
+        self.mpl_obj_keepout_btm = rc_btm
+
+        rc = mpatches.Rectangle(
+            [xmin, ymin], width, ymax - ymin, color=color, alpha=0.6, visible=False
+        )
+        mpl_ax.add_patch(rc)
+        self.mpl_obj = rc
+
+        rc_top = mpatches.Rectangle(
+            [xmin, ymax],
+            width,
+            ymax_keepout - ymax,
+            color=color,
+            alpha=0.3,
+            visible=False,
+        )
+        mpl_ax.add_patch(rc_top)
+        self.mpl_obj_keepout_top = rc_top
+
+    def clear_data(self) -> None:
+        self.mpl_obj.set_width(0)
+        self.mpl_obj.set_height(0)
+        self.mpl_obj.set_xy([0, 0])
+
+        self.mpl_obj_keepout_top.set_width(0)
+        self.mpl_obj_keepout_top.set_height(0)
+        self.mpl_obj_keepout_top.set_xy([0, 0])
+
+        self.mpl_obj_keepout_btm.set_width(0)
+        self.mpl_obj_keepout_btm.set_height(0)
+        self.mpl_obj_keepout_btm.set_xy([0, 0])
+
+    def get_artists(self) -> Tuple[Any]:
+        return (self.mpl_obj, self.mpl_obj_keepout_top, self.mpl_obj_keepout_btm)
+
+    def update_data(self, ymin, ymax, ymin_keepout, ymax_keepout):
+        self.mpl_obj.set_height(ymax - ymin)
+        self.mpl_obj.set_width(self.width)
+        self.mpl_obj.set_xy([self.xmin, ymin])
+
+        self.mpl_obj_keepout_top.set_height(ymax_keepout - ymax)
+        self.mpl_obj_keepout_top.set_width(self.width)
+        self.mpl_obj_keepout_top.set_xy([self.xmin, ymax])
+
+        self.mpl_obj_keepout_btm.set_height(ymin - ymin_keepout)
+        self.mpl_obj_keepout_btm.set_width(self.width)
+        self.mpl_obj_keepout_btm.set_xy([self.xmin, ymin_keepout])
+
+    def set_visible(self, visible: bool):
+        self.mpl_obj.set_visible(visible)
+        self.mpl_obj_keepout_btm.set_visible(visible)
+        self.mpl_obj_keepout_top.set_visible(visible)
+
+
+@dataclasses.dataclass
+class SpotGateArtist(GateArtist):
+    mpl_obj: mcollections.PathCollection
+
+    def __init__(self, mpl_ax: Any, xy: List[Tuple[float, float]], radius, color):
+        super().__init__(mpl_ax)
+        self.mpl_obj = mpl_ax.scatter(
+            [x for x, y in xy],
+            [y for x, y in xy],
+            s=radius,
+            visible=False,
+            marker="o",
+            facecolors=color,
+            alpha=0.5,
+            zorder=+100,
+        )
+
+    def clear_data(self) -> None:
+        self.mpl_obj.set_offsets([])
+
+    def get_artists(self) -> Tuple[Any]:
+        return (self.mpl_obj,)
+
+    def update_data(self, xy: List[Tuple[float, float]]):
+        self.mpl_obj.set_offsets(xy)
+
+    def set_visible(self, visible: bool):
+        self.mpl_obj.set_visible(visible)
+
+
+@dataclasses.dataclass
+class FieldOfView:
+    xmin: Optional[float] = dataclasses.field(default=None)
+    xmax: Optional[float] = dataclasses.field(default=None)
+    ymin: Optional[float] = dataclasses.field(default=None)
+    ymax: Optional[float] = dataclasses.field(default=None)
+
+    def not_defined(self):
+        return (
+            self.xmin is None
+            or self.xmax is None
+            or self.ymin is None
+            or self.ymax is None
+        )
+
+    @property
+    def width(self):
+        return self.xmax - self.xmin
+
+    @property
+    def height(self):
+        return self.ymax - self.ymin
+
+    def to_json(self):
+        return {
+            "xmin": self.xmin,
+            "xmax": self.xmax,
+            "ymin": self.ymin,
+            "ymax": self.ymax,
+        }
+
+    @classmethod
+    def from_json(cls, json_dict):
+        return FieldOfView(
+            xmin=json_dict["xmin"],
+            xmax=json_dict["xmax"],
+            ymin=json_dict["ymin"],
+            ymax=json_dict["ymax"],
+        )
+
+
+@dataclasses.dataclass(init=False)
+class GatePainter:
+    artist_objs: Dict[str, GateArtist]
+    artist_methods: Dict[str, Callable]
+
+    def __init__(self, mpl_ax, qpu_fov: FieldOfView, scale: float = 1.0):
+        # gates:
+        self.artist_objs = {}
+        self.artist_methods = {
+            "GlobalCZGate": self.process_global_cz_gate,
+            "GlobalHyperfineRotation": self.process_global_hyperfine_rotation,
+            "TopHatCZGate": self.process_tophat_cz_gate,
+            "LocalHyperfineZRotation": self.process_local_hyperfine_z_rotation,
+            "LocalHyperfineZRotationPos": self.process_local_hyperfine_z_rotation_pos,
+            "LocalHyperfineRotation": self.process_local_hyperfine_rotation,
+            "LocalHyperfineRotationPos": self.process_local_hyperfine_rotation_pos,
+        }
+        self.artist_objs["GlobalCZGate"] = GlobalGateArtist(
+            mpl_ax=mpl_ax,
+            xmin=qpu_fov.xmin,
+            width=qpu_fov.width,
+            ymin=qpu_fov.ymin,
+            height=qpu_fov.height,
+            color=_gate_color_code["GlobalCZGate"],
+        )
+        self.artist_objs["GlobalHyperfineRotation"] = GlobalGateArtist(
+            mpl_ax=mpl_ax,
+            xmin=qpu_fov.xmin,
+            width=qpu_fov.width,
+            ymin=qpu_fov.ymin,
+            height=qpu_fov.height,
+            color=_gate_color_code["GlobalHyperfineRotation"],
+        )
+        self.artist_objs["TopHatCZGate"] = RowRegionGateArtist(
+            mpl_ax=mpl_ax,
+            xmin=qpu_fov.xmin,
+            width=qpu_fov.width,
+            ymin=0,
+            ymin_keepout=0,
+            ymax=0,
+            ymax_keepout=0,
+            color=_gate_color_code["TopHatCZGate"],
+        )
+
+        self.artist_objs["LocalHyperfineZRotation"] = SpotGateArtist(
+            mpl_ax=mpl_ax,
+            xy=[],
+            radius=160 * scale,
+            color=_gate_color_code["LocalHyperfineZRotation"],
+        )
+        self.artist_objs["LocalHyperfineZRotationPos"] = SpotGateArtist(
+            mpl_ax=mpl_ax,
+            xy=[],
+            radius=160 * scale,
+            color=_gate_color_code["LocalHyperfineZRotationPos"],
+        )
+
+        self.artist_objs["LocalHyperfineRotation"] = SpotGateArtist(
+            mpl_ax=mpl_ax,
+            xy=[],
+            radius=80,
+            color=_gate_color_code["LocalHyperfineRotation"],
+        )
+        self.artist_objs["LocalHyperfineRotationPos"] = SpotGateArtist(
+            mpl_ax=mpl_ax,
+            xy=[],
+            radius=80,
+            color=_gate_color_code["LocalHyperfineRotationPos"],
+        )
+
+    @staticmethod
+    def process_global_cz_gate(artist_obj, **kwargs) -> None:
+        artist_obj.set_visible(True)
+
+    @staticmethod
+    def process_global_hyperfine_rotation(artist_obj, **kwargs) -> None:
+        artist_obj.set_visible(True)
+
+    @staticmethod
+    def process_tophat_cz_gate(artist_obj, **kwargs) -> None:
+        artist_obj.update_data(
+            ymin=kwargs["y_min"],
+            ymax=kwargs["y_max"],
+            ymin_keepout=kwargs["y_min_keepout"],
+            ymax_keepout=kwargs["y_max_keepout"],
+        )
+        artist_obj.set_visible(True)
+
+    @staticmethod
+    def process_local_hyperfine_z_rotation_pos(artist_obj, **kwargs) -> None:
+        artist_obj.update_data(xy=[(x, kwargs["y_row_loc"]) for x in kwargs["varargs"]])
+        artist_obj.set_visible(True)
+
+    @staticmethod
+    def process_local_hyperfine_rotation_pos(artist_obj, **kwargs) -> None:
+        artist_obj.update_data(xy=[(x, kwargs["y_row_loc"]) for x in kwargs["varargs"]])
+        artist_obj.set_visible(True)
+
+    @staticmethod
+    def process_local_hyperfine_z_rotation(artist_obj, **kwargs) -> None:
+        raise ValueError(
+            "cannot have local HyperfineZRotation, it should be already lowered to Pos variant!"
+        )
+
+    @staticmethod
+    def process_local_hyperfine_rotation(artist_obj, **kwargs) -> None:
+        raise ValueError(
+            "cannot have local HyperfineRotation, it should be already lowered to Pos variant!"
+        )
+
+    def process_gates(self, gates: List[GateEvent]) -> List:
+
+        for gate_artist in self.artist_objs.values():
+            gate_artist.set_visible(False)
+
+        for gate in gates:
+            self.artist_methods[gate.cls_name](
+                self.artist_objs[gate.cls_name], **gate.kwargs
+            )
+
+        return list(
+            itertools.chain.from_iterable(
+                gate_artist.get_artists() for gate_artist in self.artist_objs.values()
+            )
+        )

bloqade/visual/animation/gate_event.py

@@ -0,0 +1,24 @@
+import dataclasses
+from typing import Any, Dict
+
+
+@dataclasses.dataclass(frozen=True)
+class GateEvent:
+    cls_name: str
+    kwargs: Dict[str, Any]
+    duration: float
+
+    def to_json(self) -> Dict[str, Any]:
+        return {
+            "cls_name": self.cls_name,
+            "kwargs": self.kwargs,
+            "duration": self.duration,
+        }
+
+    @classmethod
+    def from_json(cls, json_dict: Dict[str, Any]) -> "GateEvent":
+        return cls(
+            json_dict["cls_name"],
+            json_dict["kwargs"],
+            json_dict["duration"],
+        )

bloqade/visual/animation/runtime/aod.py

@@ -0,0 +1,36 @@
+import dataclasses
+from typing import Any, Dict, Tuple
+
+from .ppoly import PPoly
+
+
+@dataclasses.dataclass(frozen=True)
+class AODMoveEvent:
+    time: float
+    duration: float
+    x: PPoly
+    y: PPoly
+
+    def __post_init__(self):
+        assert isinstance(self.x, PPoly), "x must be a PPoly instance"
+        assert isinstance(self.y, PPoly), "y must be a PPoly instance"
+
+    def sample(self, t: float) -> Tuple[float, float]:
+        return (self.x(t - self.time), self.y(t - self.time))
+
+    def to_json(self) -> Dict[str, Any]:
+        return {
+            "time": self.time,
+            "duration": self.duration,
+            "x": self.x.to_json(),
+            "y": self.y.to_json(),
+        }
+
+    @classmethod
+    def from_json(cls, json_dict: Dict[str, Any]) -> "AODMoveEvent":
+        return cls(
+            json_dict["time"],
+            json_dict["duration"],
+            PPoly.from_json(json_dict["x"]),
+            PPoly.from_json(json_dict["y"]),
+        )

bloqade/visual/animation/runtime/atoms.py

@@ -0,0 +1,55 @@
+import dataclasses
+from typing import List, Tuple
+from functools import cached_property
+
+import numpy as np
+
+from .ppoly import PPoly
+
+
+@dataclasses.dataclass(frozen=True)
+class AtomTrajectory:
+    id: int
+    x: PPoly
+    y: PPoly
+    events: List[Tuple[float, str]]  # (time, event_type)
+
+    def __post_init__(self):
+        assert all(
+            isinstance(event, tuple) and len(event) == 2 for event in self.events
+        ), "All events must be (float, str) pairs"
+        assert isinstance(self.id, int), "id must be an int"
+        assert isinstance(self.x, PPoly), "x must be a PPoly instance"
+        assert isinstance(self.y, PPoly), "y must be a PPoly instance"
+
+    def position(self, t: float) -> Tuple[np.ndarray, np.ndarray]:
+        return (self.x(t), self.y(t))
+
+    @cached_property
+    def has_lost(self) -> bool:
+        return any(event_type == "Lost" for _, event_type in self.events)
+
+    def is_lost(self, time: float) -> bool:
+        if not self.has_lost:
+            return False
+
+        return (
+            self.events[-1][0] <= time
+        )  # TODO: maybe we need approx_eq to be separate package?
+
+    def to_json(self):
+        return {
+            "id": self.id,
+            "x": self.x.to_json(),
+            "y": self.y.to_json(),
+            "events": self.events,
+        }
+
+    @classmethod
+    def from_json(cls, json_dict):
+        return AtomTrajectory(
+            json_dict["id"],
+            PPoly.from_json(json_dict["x"]),
+            PPoly.from_json(json_dict["y"]),
+            list(map(tuple, json_dict["events"])),
+        )

bloqade/visual/animation/runtime/ppoly.py

@@ -0,0 +1,50 @@
+import dataclasses
+from typing import Any, Dict
+from functools import cached_property
+
+import numpy as np
+import scipy.interpolate as interp
+
+from .utils import array_to_json, json_to_array
+
+
+@dataclasses.dataclass
+class PPoly:
+    c: np.ndarray
+    x: np.ndarray
+    extrapolate: bool = dataclasses.field(default=False)
+
+    def __post_init__(self):
+        super().__setattr__("x", np.asarray(self.x))
+        super().__setattr__("c", np.asarray(self.c))
+
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, PPoly):
+            return False
+        return (
+            np.array_equal(self.x, other.x)
+            and np.array_equal(self.c, other.c)
+            and self.extrapolate == other.extrapolate
+        )
+
+    @cached_property
+    def get_ppoly(self) -> interp.PPoly:
+        return interp.PPoly(self.c, self.x, self.extrapolate)
+
+    def __call__(self, x: float) -> float:
+        return self.get_ppoly(x)
+
+    def to_json(self) -> Dict[str, Any]:
+        return {
+            "c": array_to_json(self.c),
+            "x": array_to_json(self.x),
+            "extrapolate": self.extrapolate,
+        }
+
+    @classmethod
+    def from_json(cls, json_dict: Dict[str, Any]) -> "PPoly":
+        return cls(
+            c=json_to_array(json_dict["c"]),
+            x=json_to_array(json_dict["x"]),
+            extrapolate=json_dict["extrapolate"],
+        )

bloqade/visual/animation/runtime/qpustate.py

@@ -0,0 +1,119 @@
+import dataclasses
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Tuple
+
+import numpy as np
+
+from .aod import AODMoveEvent
+from ..base import FieldOfView
+from .atoms import AtomTrajectory
+from ..gate_event import GateEvent
+
+
+@dataclasses.dataclass
+class QPUStateABC(ABC):
+
+    block_durations: List[float] = dataclasses.field(default_factory=list)
+    gate_events: List[Tuple[float, GateEvent]] = dataclasses.field(default_factory=list)
+    qpu_fov: FieldOfView = dataclasses.field(default_factory=FieldOfView)
+
+    def __post_init__(self):
+        assert all(
+            isinstance(gate, GateEvent) for _, gate in self.gate_events
+        ), "All gate events must be GateEvent instances"
+
+    @abstractmethod
+    def get_slm_sites(self) -> np.ndarray: ...
+
+    @abstractmethod
+    def sample_aod_traps(self, time: float) -> List[Tuple[float, float]]: ...
+
+    @abstractmethod
+    def get_atoms_lost_info(self, time: float) -> List[str]: ...
+
+    @abstractmethod
+    def get_atoms_position(
+        self, time: float, include_lost: bool = False
+    ) -> List[Tuple[int, Tuple[float, float]]]: ...
+
+    def get_gate_events_timing(self) -> List[Tuple[float, float]]:
+        # return [t_start, duration]
+        return [(t, gate.duration) for t, gate in self.gate_events]
+
+    def get_gate_events(self, time: float) -> List[Tuple[float, GateEvent]]:
+        out = []
+        if self.gate_events:
+            for t, gate in self.gate_events:
+                if t > time:
+                    break
+                if t <= time < t + gate.duration:
+                    out.append((t, gate))
+        return out
+
+
+@dataclasses.dataclass
+class AnimateQPUState(QPUStateABC):
+    atoms: List[AtomTrajectory] = dataclasses.field(default_factory=list)
+    slm_zone: List[Tuple[float, float]] = dataclasses.field(default_factory=list)
+    aod_moves: List[AODMoveEvent] = dataclasses.field(default_factory=list)
+
+    def __post_init__(self):
+        assert all(
+            isinstance(atom, AtomTrajectory) for atom in self.atoms
+        ), "All atoms must be AtomTrajectory instances"
+        assert all(
+            isinstance(aod, AODMoveEvent) for aod in self.aod_moves
+        ), "All AOD moves must be AODMoveEvent instances"
+        assert all(
+            isinstance(site, tuple) and len(site) == 2 for site in self.slm_zone
+        ), "All SLM sites must be tuples of length 2"
+
+    def get_slm_sites(self):
+        return np.array(self.slm_zone)
+
+    def sample_aod_traps(self, time: float) -> List[Tuple[float, float]]:
+        out = []
+        for aod_trap in self.aod_moves:
+            if aod_trap.time <= time < aod_trap.time + aod_trap.duration:
+                out.append(aod_trap.sample(time))
+
+        return out
+
+    def get_atoms_lost_info(self, time: float) -> List[str]:
+        return [
+            f"Lost: {atom.id} @{atom.events[-1][0]:.3f} (us)" + "\n"
+            for atom in self.atoms
+            if atom.is_lost(time)
+        ]
+
+    def get_atoms_position(
+        self, time: float, include_lost: bool = False
+    ) -> List[Tuple[int, Tuple[float, float]]]:
+        return [
+            (atom.id, atom.position(time))
+            for atom in self.atoms
+            if not atom.is_lost(time) or include_lost
+        ]
+
+    def to_json(self) -> Dict[str, Any]:
+        return {
+            "block_durations": self.block_durations,
+            "gate_events": [(t, gate.to_json()) for t, gate in self.gate_events],
+            "qpu_fov": self.qpu_fov.to_json(),
+            "atoms": [atom.to_json() for atom in self.atoms],
+            "slm_zone": self.slm_zone,
+            "aod_moves": [aod.to_json() for aod in self.aod_moves],
+        }
+
+    @classmethod
+    def from_json(cls, json_dict: Dict[str, Any]) -> "AnimateQPUState":
+        return cls(
+            block_durations=json_dict["block_durations"],
+            gate_events=[
+                (t, GateEvent.from_json(gate)) for t, gate in json_dict["gate_events"]
+            ],
+            qpu_fov=FieldOfView.from_json(json_dict["qpu_fov"]),
+            atoms=list(map(AtomTrajectory.from_json, json_dict["atoms"])),
+            slm_zone=list(map(tuple, json_dict["slm_zone"])),
+            aod_moves=list(map(AODMoveEvent.from_json, json_dict["aod_moves"])),
+        )

bloqade/visual/animation/runtime/utils.py

@@ -0,0 +1,43 @@
+import zlib
+import base64
+from typing import Any, Dict
+
+import numpy as np
+
+
+def bytes_to_str(byte_array: bytes) -> str:
+    compressed_bytes = zlib.compress(byte_array)
+    return base64.b64encode(compressed_bytes).decode("ascii")
+
+
+def str_to_bytes(string: str) -> bytes:
+    compressed_bytes = base64.b64decode(string)
+    return zlib.decompress(compressed_bytes)
+
+
+def validate_dtype(dtype) -> None:
+    if np.dtype(dtype) == np.dtype(object):
+        raise ValueError("Object arrays are not supported")
+    elif dtype.subdtype is not None:
+        validate_dtype(dtype.subdtype[0])
+    elif dtype.fields is not None:
+        for _, (field_dtype, _) in dtype.fields.items():
+            validate_dtype(field_dtype)
+
+
+def array_to_json(arr: np.ndarray) -> Dict[str, Any]:
+    validate_dtype(arr.dtype)
+
+    return {
+        "dtype": (arr.dtype.descr if arr.dtype.fields is not None else arr.dtype.str),
+        "shape": arr.shape,
+        "buffer": bytes_to_str(arr.tobytes()),
+    }
+
+
+def json_to_array(json_dict: dict) -> np.ndarray:
+    _data = json_dict["buffer"]
+    _dtype = json_dict["dtype"]
+    _shape = json_dict["shape"]
+
+    return np.frombuffer(str_to_bytes(_data), dtype=_dtype).reshape(_shape)