Refactor standard library and updating dependencies. (#60)

* adding new submodule for standard library

* updating lock

* Apply suggestions from code review

Co-authored-by: Copilot <[email protected]>

* adding imports

* adding imports

---------

Co-authored-by: Copilot <[email protected]>

Author: weinbe58

pyproject.toml

@@ -7,12 +7,8 @@ authors = [
 ]
 dependencies = [
     "bloqade-geometry ~=0.3.0",
-    "bloqade-circuit ~=0.6.0",
+    "bloqade-circuit >=0.6.0,<0.8.0-dev",
     "kirin-toolchain >= 0.17.17,<0.18.0-dev",
-    "pint>=0.24.3",
-    "pyarrow>=17.0.0",
-    "pytest>=8.3.2",
-    "numpy>=1.26.4",
 ]
 
 readme = "README.md"
@@ -23,8 +19,6 @@ visualization = [
     "matplotlib>=3.9.4",
 ]
 
-
-
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"

src/bloqade/shuttle/__init__.py

@@ -1,4 +1,3 @@
-from . import arch as arch, passes as passes, stdlib as stdlib
 from .dialects.action import _interface as action
 from .dialects.atom import _interface as atom
 from .dialects.filled import _interface as filled

src/bloqade/shuttle/stdlib/__init__.py

@@ -0,0 +1 @@
+from . import layouts as layouts

src/bloqade/shuttle/stdlib/layouts/__init__.py

@@ -0,0 +1 @@
+from . import single_col_zone as single_col_zone, two_col_zone as two_col_zone

src/bloqade/shuttle/stdlib/layouts/asserts.py

@@ -0,0 +1,7 @@
+from bloqade.shuttle.prelude import tweezer
+
+
+@tweezer
+def assert_sorted(indices):
+    for i in range(1, len(indices)):
+        assert indices[i - 1] < indices[i], "Indices must be sorted in ascending order."

src/bloqade/shuttle/stdlib/layouts/single_col_zone.py

@@ -0,0 +1,133 @@
+from itertools import repeat
+from typing import Any, TypeVar
+
+from bloqade.geometry.dialects import grid
+from kirin.dialects import ilist
+
+from bloqade.shuttle import action, gate, schedule, spec
+from bloqade.shuttle.prelude import move, tweezer
+
+from .asserts import assert_sorted
+
+
+def get_spec(num_x: int, num_y: int, spacing: float = 10.0) -> spec.ArchSpec:
+    """Create a static trap spec with a single zone. compatible with the stdlib
+
+    Args:
+        num_x (int): Number of traps in the x direction.
+        num_y (int): Number of traps in the y direction.
+        spacing (float): Spacing between traps in both directions. Default is 10.0.
+
+    Returns:
+        spec.Spec: A specification object containing the layout with a single zone.
+
+    """
+    x_spacing = tuple(repeat(spacing, num_x - 1))
+    y_spacing = tuple(repeat(spacing, num_y - 1))
+
+    return spec.ArchSpec(
+        layout=spec.Layout(
+            static_traps={"traps": grid.Grid(x_spacing, y_spacing, 0.0, 0.0)},
+            fillable=set(["traps"]),
+        )
+    )
+
+
+NumX = TypeVar("NumX")
+NumY = TypeVar("NumY")
+
+
+@tweezer
+def single_zone_move_cz(
+    zone: grid.Grid[Any, Any],
+    ctrl_x_ids: ilist.IList[int, NumX],
+    ctrl_y_ids: ilist.IList[int, NumY],
+    qarg_x_ids: ilist.IList[int, NumX],
+    qarg_y_ids: ilist.IList[int, NumY],
+    shift_x: float,
+    shift_y: float,
+):
+    assert len(ctrl_x_ids) == len(
+        qarg_x_ids
+    ), "Control and target x indices must have the same length."
+
+    assert len(ctrl_y_ids) == len(
+        qarg_y_ids
+    ), "Control and target y indices must have the same length."
+
+    assert_sorted(ctrl_x_ids)
+    assert_sorted(ctrl_y_ids)
+    assert_sorted(qarg_x_ids)
+    assert_sorted(qarg_y_ids)
+
+    start = grid.sub_grid(zone, ctrl_x_ids, ctrl_y_ids)
+    target_atoms = grid.sub_grid(zone, qarg_x_ids, qarg_y_ids)
+
+    end = grid.shift(target_atoms, shift_x, shift_y)
+    first_pos = grid.shift(start, shift_x, shift_y)
+    second_pos = grid.from_positions(grid.get_xpos(end), grid.get_ypos(first_pos))
+
+    action.set_loc(start)
+    action.turn_on(action.ALL, action.ALL)
+    action.move(first_pos)
+    action.move(second_pos)
+    action.move(end)
+
+
+@move
+def default_move_cz_impl(
+    zone: grid.Grid[Any, Any],
+    x_shift: float,
+    y_shift: float,
+    ctrl_x_ids: ilist.IList[int, NumX],
+    ctrl_y_ids: ilist.IList[int, NumY],
+    qarg_x_ids: ilist.IList[int, NumX],
+    qarg_y_ids: ilist.IList[int, NumY],
+):
+    """Move atoms from the start ids and run cz gate with the atoms at the end ids.
+
+    Args:
+        zone (grid.Grid[Any, Any]): The grid representing the trap layout (zone) in which atoms are moved.
+        x_shift (float): The amount to shift atoms in the x direction during the move.
+        y_shift (float): The amount to shift atoms in the y direction during the move.
+        ctrl_x_ids (ilist.IList[int, NumX]): The x-indices of the starting positions.
+        ctrl_y_ids (ilist.IList[int, NumY]): The y-indices of the starting positions.
+        qarg_x_ids (ilist.IList[int, NumX]): The x-indices of the ending positions.
+        qarg_y_ids (ilist.IList[int, NumY]): The y-indices of the ending positions.
+    """
+    if len(ctrl_x_ids) < 1 or len(qarg_x_ids) < 1:
+        return
+
+    fwd_kernel = schedule.device_fn(
+        single_zone_move_cz,
+        ilist.range(len(ctrl_x_ids)),
+        ilist.range(len(ctrl_y_ids)),
+    )
+    bwd_kernel = schedule.reverse(fwd_kernel)
+
+    fwd_kernel(zone, ctrl_x_ids, ctrl_y_ids, qarg_x_ids, qarg_y_ids, x_shift, y_shift)
+    gate.top_hat_cz(zone)
+    bwd_kernel(zone, qarg_x_ids, qarg_y_ids, ctrl_x_ids, ctrl_y_ids, x_shift, y_shift)
+
+
+DEFAULT_X_SHIFT = 2.0
+DEFAULT_Y_SHIFT = 2.0
+
+
+@move
+def cz_move(
+    ctrl_x_ids: ilist.IList[int, NumX],
+    ctrl_y_ids: ilist.IList[int, NumY],
+    qarg_x_ids: ilist.IList[int, NumX],
+    qarg_y_ids: ilist.IList[int, NumY],
+):
+    zone = spec.get_static_trap(zone_id="traps")
+    default_move_cz_impl(
+        zone,
+        DEFAULT_X_SHIFT,
+        DEFAULT_Y_SHIFT,
+        ctrl_x_ids,
+        ctrl_y_ids,
+        qarg_x_ids,
+        qarg_y_ids,
+    )

src/bloqade/shuttle/stdlib/layouts/two_col_zone.py

@@ -0,0 +1,103 @@
+from itertools import repeat
+from typing import TypeVar
+
+from bloqade.geometry.dialects import grid
+from kirin.dialects import ilist
+
+from bloqade.shuttle import action, schedule, spec
+from bloqade.shuttle.prelude import move, tweezer
+
+from .asserts import assert_sorted
+
+
+def get_spec(
+    num_x: int, num_y: int, spacing: float = 10.0, gate_spacing: float = 2.0
+) -> spec.ArchSpec:
+    """Create a static trap spec with a single zone with pairs of traps oriented
+    horizontally.
+
+    Args:
+        num_x (int): Number of pairs of traps in the x direction.
+        num_y (int): Number of pairs of traps in the y direction.
+        spacing (float): Spacing between traps in both directions. Default is 10.0.
+        gate_spacing (float): Spacing between gates. Default is 2.0.
+
+    Returns:
+        spec.Spec: A specification object containing the layout with a single zone.
+
+    """
+    x_spacing = sum(repeat((gate_spacing, spacing), num_x - 1), ())
+    y_spacing = tuple(repeat(spacing, num_y - 1))
+
+    return spec.ArchSpec(
+        layout=spec.Layout(
+            static_traps={"traps": grid.Grid(x_spacing, y_spacing, 0.0, 0.0)},
+            fillable=set(["traps"]),
+        )
+    )
+
+
+NumX = TypeVar("NumX")
+NumY = TypeVar("NumY")
+
+
+@tweezer
+def rearrange_impl(
+    src_x: ilist.IList[int, NumX],
+    src_y: ilist.IList[int, NumY],
+    dst_x: ilist.IList[int, NumX],
+    dst_y: ilist.IList[int, NumY],
+):
+    assert len(src_x) == len(
+        dst_x
+    ), "Source and destination x indices must have the same length."
+    assert len(src_y) == len(
+        dst_y
+    ), "Source and destination y indices must have the same length."
+
+    assert_sorted(src_x)
+    assert_sorted(src_y)
+    assert_sorted(dst_x)
+    assert_sorted(dst_y)
+
+    zone = spec.get_static_trap(zone_id="traps")
+
+    start = grid.sub_grid(zone, src_x, src_y)
+    end = grid.sub_grid(zone, dst_x, dst_y)
+
+    x_positions = grid.get_xpos(start)
+
+    def parking_x(index: int):
+        return x_positions[index] + 3.0 * (2 * (index % 2) - 1)
+
+    src_parking = grid.from_positions(ilist.map(parking_x, src_x), grid.get_ypos(zone))
+
+    dst_parking = grid.shift(end, 0.0, 3.0)
+    mid_pos = grid.from_positions(
+        grid.get_xpos(src_parking), grid.get_ypos(dst_parking)
+    )
+
+    action.set_loc(start)
+    action.turn_on(action.ALL, action.ALL)
+    action.move(src_parking)
+    action.move(mid_pos)
+    action.move(dst_parking)
+    action.move(end)
+    action.turn_off(action.ALL, action.ALL)
+
+
+@move
+def rearrange(
+    src_x: ilist.IList[int, NumX],
+    src_y: ilist.IList[int, NumY],
+    dst_x: ilist.IList[int, NumX],
+    dst_y: ilist.IList[int, NumY],
+):
+    if len(src_x) < 1 or len(dst_x) < 1:
+        return
+
+    x_tones = ilist.range(len(src_x))
+    y_tones = ilist.range(len(src_y))
+
+    device_fn = schedule.device_fn(rearrange_impl, x_tones, y_tones)
+    device_fn(src_x, src_y, dst_x, dst_y)

src/bloqade/shuttle/stdlib/moves.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Any, TypeVar
 
 from bloqade.geometry.dialects import grid
@@ -6,6 +7,14 @@
 from bloqade.shuttle import action, gate, schedule, spec
 from bloqade.shuttle.prelude import move, tweezer
 
+warnings.warn(
+    (
+        "This module has been moved to `stdlib.layouts.single_col_zone` submodule and"
+        " will be removed in 0.8.0"
+    ),
+    DeprecationWarning,
+)
+
 NumX = TypeVar("NumX")
 NumY = TypeVar("NumY")
 

src/bloqade/shuttle/stdlib/spec.py

@@ -1,9 +1,18 @@
+import warnings
 from itertools import repeat
 
 from bloqade.geometry.dialects import grid
 
 from bloqade.shuttle import spec
 
+warnings.warn(
+    (
+        "This module's contents have been moved to `stdlib.layouts.single_col_zone` submodule and"
+        " will be removed in 0.8.0"
+    ),
+    DeprecationWarning,
+)
+
 
 def single_zone_spec(num_x: int, num_y: int, spacing: float = 10.0) -> spec.ArchSpec:
     """Create a static trap spec with a single zone. compatible with the stdlib