feat: plot sequence diagram

docs/quickstart.rst

@@ -34,3 +34,10 @@
 
 .. plot:: ../examples/flexible_length.py
     :include-source:
+
+
+画图展示编排结果
+----------------
+
+.. plot:: ../examples/plot.py
+    :include-source:

examples/plot.py

@@ -0,0 +1,25 @@
+import matplotlib.pyplot as plt
+
+from bosing import Barrier, Play, ShiftPhase, Stack
+
+xy = [Play(f"xy{i}", "hann", 1.0, 100e-9) for i in range(2)]
+z = [
+    Stack(Play(f"z{i}", "hann", 1.0, 100e-9), ShiftPhase(f"xy{i}", 1.0))
+    for i in range(2)
+]
+m = Stack(*(Play(f"m{i}", "hann", 1.0, 100e-9, plateau=200e-9) for i in range(2)))
+b = Barrier()
+
+schedule = Stack(
+    xy[0],
+    xy[1],
+    b,
+    z[1],
+    b,
+    xy[1],
+    b,
+    m,
+)
+
+schedule.plot()
+plt.show()

examples/schedule.py

@@ -66,6 +66,9 @@
         Barrier(duration=15e-9),
     )
 
+    schedule.plot(channels=["xy0", "u0", "m0"], max_depth=6)
+    plt.show()
+
     result = generate_waveforms(channels, shapes, schedule, time_tolerance=1e-13)
 
     t = np.arange(length) / 2e9

pyproject.toml

@@ -41,6 +41,7 @@ dev-dependencies = [
   "furo>=2024.8.6",
   "sphinx>=7.4.7",
   "rich>=13.9.2",
+  "pyside6>=6.8.0.1",
 ]
 cache-keys = [
   { file = "src/**/*.rs" },

python/bosing/__init__.py

@@ -6,5 +6,58 @@
     of :math:`0.5` means a phase shift of :math:`\pi` radians.
 """
 
-from ._bosing import *  # noqa: F403
-from ._bosing import __all__ as __all__
+from ._bosing import (
+    Absolute,
+    AbsoluteEntry,
+    Alignment,
+    Barrier,
+    Channel,
+    Direction,
+    Element,
+    Grid,
+    GridEntry,
+    GridLength,
+    GridLengthUnit,
+    Hann,
+    Interp,
+    OscState,
+    Play,
+    Repeat,
+    SetFreq,
+    SetPhase,
+    Shape,
+    ShiftFreq,
+    ShiftPhase,
+    Stack,
+    SwapPhase,
+    generate_waveforms,
+    generate_waveforms_with_states,
+)
+
+__all__ = [
+    "Absolute",
+    "AbsoluteEntry",
+    "Alignment",
+    "Barrier",
+    "Channel",
+    "Direction",
+    "Element",
+    "Grid",
+    "GridEntry",
+    "GridLength",
+    "GridLengthUnit",
+    "Hann",
+    "Interp",
+    "OscState",
+    "Play",
+    "Repeat",
+    "SetFreq",
+    "SetPhase",
+    "Shape",
+    "ShiftFreq",
+    "ShiftPhase",
+    "Stack",
+    "SwapPhase",
+    "generate_waveforms",
+    "generate_waveforms_with_states",
+]

python/bosing/_bosing.pyi

@@ -1,9 +1,10 @@
 # ruff: noqa: PLR0913
 from collections.abc import Iterable, Mapping, Sequence
-from typing import Any, ClassVar, Literal, final
+from typing import Any, ClassVar, Literal, final, type_check_only
 
 import numpy as np
 import numpy.typing as npt
+from matplotlib.axes import Axes
 from typing_extensions import Self, TypeAlias
 
 __all__ = [
@@ -32,6 +33,7 @@ __all__ = [
     "OscState",
     "generate_waveforms",
     "generate_waveforms_with_states",
+    "ItemKind",
 ]
 
 _RichReprResult: TypeAlias = list[Any]
@@ -123,6 +125,13 @@ class Element:
     @property
     def min_duration(self) -> float: ...
     def measure(self) -> float: ...
+    def plot(
+        self,
+        ax: Axes | None = ...,
+        *,
+        channels: Sequence[str] | None = ...,
+        max_depth: int = ...,
+    ) -> Axes: ...
 
 @final
 class Play(Element):
@@ -462,6 +471,34 @@ class OscState:
     def with_time_shift(self, time: float) -> Self: ...
     def __rich_repr__(self) -> _RichReprResult: ...  # undocumented
 
+@type_check_only
+@final
+class PlotItem:
+    @property
+    def channels(self) -> list[str]: ...
+    @property
+    def start(self) -> float: ...
+    @property
+    def span(self) -> float: ...
+    @property
+    def depth(self) -> int: ...
+    @property
+    def kind(self) -> ItemKind: ...
+
+@final
+class ItemKind:
+    Play: ClassVar[ItemKind]
+    ShiftPhase: ClassVar[ItemKind]
+    SetPhase: ClassVar[ItemKind]
+    ShiftFreq: ClassVar[ItemKind]
+    SetFreq: ClassVar[ItemKind]
+    SwapPhase: ClassVar[ItemKind]
+    Barrier: ClassVar[ItemKind]
+    Repeat: ClassVar[ItemKind]
+    Stack: ClassVar[ItemKind]
+    Absolute: ClassVar[ItemKind]
+    Grid: ClassVar[ItemKind]
+
 def generate_waveforms(
     channels: Mapping[str, Channel],
     shapes: Mapping[str, Shape],

python/bosing/_plot.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.patches import PathPatch
+from matplotlib.path import Path
+from matplotlib.ticker import EngFormatter
+
+from bosing._bosing import ItemKind
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator, Sequence
+
+    from matplotlib.axes import Axes
+    from typing_extensions import TypeAlias
+
+    from bosing._bosing import PlotItem
+
+    _RECTS: TypeAlias = defaultdict[ItemKind, list[tuple[float, float, float]]]
+    _MARKERS: TypeAlias = defaultdict[ItemKind, tuple[list[float], list[float]]]
+
+
+COLORS = {
+    ItemKind.Play: "blue",
+    ItemKind.ShiftPhase: "green",
+    ItemKind.SetPhase: "red",
+    ItemKind.ShiftFreq: "green",
+    ItemKind.SetFreq: "red",
+    ItemKind.SwapPhase: "purple",
+    ItemKind.Barrier: "gray",
+    ItemKind.Repeat: "yellow",
+    ItemKind.Stack: "orange",
+    ItemKind.Absolute: "cyan",
+    ItemKind.Grid: "black",
+}
+
+MARKERS = {
+    ItemKind.ShiftPhase: "$\\circlearrowleft$",
+    ItemKind.SetPhase: "$\\circlearrowleft$",
+    ItemKind.ShiftFreq: "$\\Uparrow$",
+    ItemKind.SetFreq: "$\\Uparrow$",
+    ItemKind.SwapPhase: "$\\leftrightarrow$",
+}
+
+LABELS = {
+    ItemKind.Play: "Play",
+    ItemKind.ShiftPhase: "Shift Phase",
+    ItemKind.SetPhase: "Set Phase",
+    ItemKind.ShiftFreq: "Shift Frequency",
+    ItemKind.SetFreq: "Set Frequency",
+    ItemKind.SwapPhase: "Swap Phase",
+    ItemKind.Barrier: "Barrier",
+    ItemKind.Repeat: "Repeat",
+    ItemKind.Stack: "Stack",
+    ItemKind.Absolute: "Absolute",
+    ItemKind.Grid: "Grid",
+}
+
+
+def manage_channel_stack(ch_stack: list[list[str]], x: PlotItem) -> None:
+    prev_depth = len(ch_stack) - 1
+    if x.depth > prev_depth:
+        ch_stack.append(x.channels)
+    elif x.depth < prev_depth:
+        _ = ch_stack.pop()
+        ch_stack[-1] = x.channels
+    else:
+        ch_stack[-1] = x.channels
+
+
+def get_plot_channels(
+    ch_stack: list[list[str]], x: PlotItem, channels: Sequence[str]
+) -> Sequence[str]:
+    if x.kind == ItemKind.Barrier and len(x.channels) == 0:
+        for chs in reversed(ch_stack):
+            if len(chs) > 0:
+                return chs
+        return channels
+    return x.channels
+
+
+def process_blocks(
+    blocks: Iterator[PlotItem],
+    channels: Sequence[str],
+    max_depth: int,
+    channels_ystart: dict[str, int],
+) -> tuple[_RECTS, _MARKERS]:
+    ch_stack: list[list[str]] = []
+    rects: _RECTS = defaultdict(list)
+    markers: _MARKERS = defaultdict(lambda: ([], []))
+
+    for x in blocks:
+        manage_channel_stack(ch_stack, x)
+        if x.depth >= max_depth:
+            continue
+        for c in get_plot_channels(ch_stack, x, channels):
+            if c in channels_ystart:
+                y = channels_ystart[c] + x.depth
+                if x.kind in MARKERS:
+                    mx, my = markers[x.kind]
+                    mx.append(x.start)
+                    my.append(y)
+                else:
+                    rects[x.kind].append((x.start, y, x.span))
+    return rects, markers
+
+
+def plot(
+    ax: Axes | None, blocks: Iterator[PlotItem], channels: Sequence[str], max_depth: int
+) -> Axes:
+    if ax is None:
+        ax = plt.gca()
+
+    channels_ystart = {c: i * (max_depth + 1) for i, c in enumerate(channels)}
+    rects, markers = process_blocks(blocks, channels, max_depth, channels_ystart)
+
+    for k, r in rects.items():
+        # numrects x [x, y, width]
+        r_arr = np.array(r)
+        # numrects x numsides x 2
+        xy = np.empty((r_arr.shape[0], 4, 2))
+        xy[:, :2, 0] = r_arr[:, np.newaxis, 0]
+        xy[:, 2:, 0] = r_arr[:, np.newaxis, 0] + r_arr[:, np.newaxis, 2]
+        xy[:, [0, 3], 1] = r_arr[:, np.newaxis, 1]
+        xy[:, [1, 2], 1] = r_arr[:, np.newaxis, 1] + 1
+        path = Path.make_compound_path_from_polys(xy)
+        patch = PathPatch(path)
+        patch.set_facecolor(COLORS[k])
+        patch.set_label(LABELS[k])
+        _ = ax.add_patch(patch)
+
+    for k, (mx, my) in markers.items():
+        _ = ax.plot(  # pyright: ignore[reportUnknownMemberType]
+            mx,
+            my,
+            linestyle="",
+            marker=MARKERS[k],
+            color=COLORS[k],
+            label=LABELS[k],
+            markersize=12,
+        )
+
+    _ = ax.set_yticks(list(channels_ystart.values()), channels_ystart.keys())  # pyright: ignore[reportUnknownMemberType]
+    ax.xaxis.set_major_formatter(EngFormatter(places=3))
+    _ = ax.set_xlabel("Time")  # pyright: ignore[reportUnknownMemberType]
+    _ = ax.set_ylabel("Channels")  # pyright: ignore[reportUnknownMemberType]
+    _ = ax.legend()  # pyright: ignore[reportUnknownMemberType]
+    ax.autoscale()
+
+    return ax