bloqade-core 0.6.0__py3-none-any.whl

bloqade/geometry/__init__.py

@@ -0,0 +1,7 @@
+from .dialects.filled import _interface as filled
+from .dialects.grid import _interface as grid
+
+__all__ = [
+    "grid",
+    "filled",
+]

bloqade/geometry/dialects/filled/__init__.py

@@ -0,0 +1,9 @@
+from ._dialect import dialect as dialect
+from ._interface import (
+    fill as fill,
+    get_parent as get_parent,
+    vacate as vacate,
+)
+from .concrete import FilledGridMethods as FilledGridMethods
+from .stmts import Fill as Fill, GetParent as GetParent, Vacate as Vacate
+from .types import FilledGrid as FilledGrid, FilledGridType as FilledGridType

bloqade/geometry/dialects/filled/_dialect.py

@@ -0,0 +1,3 @@
+from kirin import ir
+
+dialect = ir.Dialect("geometry.filled")

bloqade/geometry/dialects/filled/_interface.py

@@ -0,0 +1,62 @@
+from typing import Any, TypeVar
+
+from kirin.dialects import ilist
+from kirin.lowering import wraps as _wraps
+
+from bloqade.geometry.dialects import grid
+
+from .stmts import Fill, GetParent, Vacate
+from .types import FilledGrid
+
+Nx = TypeVar("Nx")
+Ny = TypeVar("Ny")
+
+
+@_wraps(Vacate)
+def vacate(
+    zone: grid.Grid[Nx, Ny],
+    vacancies: ilist.IList[tuple[int, int], Any] | list[tuple[int, int]],
+) -> FilledGrid[Nx, Ny]:
+    """Create a FilledGrid by vacating specified positions from a grid.
+
+    Args:
+        zone: The original grid from which positions will be vacated.
+        vacancies: An IList or list of (x_index, y_index) tuples indicating positions to vacate
+
+    Returns:
+        A FilledGrid with the specified vacancies.
+
+    """
+    ...
+
+
+@_wraps(Fill)
+def fill(
+    zone: grid.Grid[Nx, Ny],
+    filled: ilist.IList[tuple[int, int], Any] | list[tuple[int, int]],
+) -> FilledGrid[Nx, Ny]:
+    """Create a FilledGrid by filling specified positions in a grid.
+
+    Args:
+        zone: The original grid in which positions will be filled.
+        filled: An IList or list of (x_index, y_index) tuples indicating positions to fill
+
+    Returns:
+        A FilledGrid with the specified positions filled.
+
+    """
+    ...
+
+
+@_wraps(GetParent)
+def get_parent(filled_grid: FilledGrid[Nx, Ny]) -> grid.Grid[Nx, Ny]:
+    """Retrieve the parent grid of a FilledGrid.
+
+    Args:
+        filled_grid: The FilledGrid whose parent grid is to be retrieved.
+
+    Returns:
+        The parent grid of the provided FilledGrid.
+
+    """
+    ...

bloqade/geometry/dialects/filled/concrete.py

@@ -0,0 +1,43 @@
+from kirin.dialects import ilist
+from kirin.interp import (
+    Frame,
+    Interpreter,
+    MethodTable,
+    impl,
+)
+
+from bloqade.geometry.dialects.grid.types import Grid
+
+from . import stmts
+from ._dialect import dialect
+from .types import FilledGrid
+
+
+@dialect.register
+class FilledGridMethods(MethodTable):
+    @impl(stmts.Vacate)
+    def vacate(self, interp: Interpreter, frame: Frame, stmt: stmts.Vacate):
+        zone = frame.get_casted(stmt.zone, Grid)
+        vacancies_val = frame.get(stmt.vacancies)
+        vacancies = (
+            vacancies_val
+            if isinstance(vacancies_val, ilist.IList)
+            else ilist.IList(vacancies_val)
+        )
+        return (FilledGrid.vacate(zone, vacancies),)
+
+    @impl(stmts.Fill)
+    def fill(self, interp: Interpreter, frame: Frame, stmt: stmts.Fill):
+        zone = frame.get_casted(stmt.zone, Grid)
+        filled_val = frame.get(stmt.filled)
+        filled = (
+            filled_val
+            if isinstance(filled_val, ilist.IList)
+            else ilist.IList(filled_val)
+        )
+        return (FilledGrid.fill(zone, filled),)
+
+    @impl(stmts.GetParent)
+    def get_parent(self, interp: Interpreter, frame: Frame, stmt: stmts.GetParent):
+        filled_grid = frame.get_casted(stmt.filled_grid, FilledGrid)
+        return (filled_grid.parent,)

bloqade/geometry/dialects/filled/stmts.py

@@ -0,0 +1,43 @@
+from kirin import decl, ir, lowering, types
+from kirin.decl import info
+from kirin.dialects import ilist
+
+from bloqade.geometry.dialects import grid
+
+from ._dialect import dialect
+from .types import FilledGridType
+
+NumVacant = types.TypeVar("NumVacant")
+Nx = types.TypeVar("Nx")
+Ny = types.TypeVar("Ny")
+
+
+@decl.statement(dialect=dialect)
+class Vacate(ir.Statement):
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
+
+    zone: ir.SSAValue = info.argument(grid.GridType[Nx, Ny])
+    vacancies: ir.SSAValue = info.argument(
+        ilist.IListType[types.Tuple[types.Int, types.Int], NumVacant]
+    )
+    result: ir.ResultValue = info.result(FilledGridType[Nx, Ny])
+
+
+@decl.statement(dialect=dialect)
+class Fill(ir.Statement):
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
+
+    zone: ir.SSAValue = info.argument(grid.GridType[Nx, Ny])
+    filled: ir.SSAValue = info.argument(
+        ilist.IListType[types.Tuple[types.Int, types.Int], NumVacant]
+    )
+    result: ir.ResultValue = info.result(FilledGridType[Nx, Ny])
+
+
+@decl.statement(dialect=dialect)
+class GetParent(ir.Statement):
+    name = "get_parent"
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
+
+    filled_grid: ir.SSAValue = info.argument(FilledGridType[Nx, Ny])
+    result: ir.ResultValue = info.result(grid.GridType[Nx, Ny])

bloqade/geometry/dialects/filled/types.py

@@ -0,0 +1,170 @@
+from dataclasses import dataclass, field
+from functools import cached_property
+from itertools import product
+from typing import Any, Iterable, Sequence, TypeVar
+
+from kirin import types
+from kirin.dialects import ilist
+from kirin.serialization.base.deserializer import Deserializer
+from kirin.serialization.base.serializer import Serializer
+from kirin.serialization.core.serializationunit import SerializationUnit
+
+from bloqade.geometry.dialects import grid
+
+from ._dialect import dialect
+
+NumX = TypeVar("NumX")
+NumY = TypeVar("NumY")
+
+
+@dataclass(eq=False)
+@dialect.register
+class FilledGrid(grid.Grid[NumX, NumY]):
+    name = "filled_grid"
+
+    x_spacing: tuple[float, ...] = field(init=False)
+    y_spacing: tuple[float, ...] = field(init=False)
+    x_init: float | None = field(init=False)
+    y_init: float | None = field(init=False)
+
+    parent: grid.Grid[NumX, NumY]
+    vacancies: frozenset[tuple[int, int]]
+
+    def __post_init__(self):
+        self.x_spacing = self.parent.x_spacing
+        self.y_spacing = self.parent.y_spacing
+        self.x_init = self.parent.x_init
+        self.y_init = self.parent.y_init
+
+        self.type = types.Generic(
+            FilledGrid,
+            types.Literal(len(self.x_spacing) + 1),
+            types.Literal(len(self.y_spacing) + 1),
+        )
+
+    def __hash__(self):
+        return hash((self.parent, self.vacancies))
+
+    def __eq__(self, other: Any) -> bool:
+        return (
+            isinstance(other, FilledGrid)
+            and self.parent == other.parent
+            and self.vacancies == other.vacancies
+        )
+
+    def is_equal(self, other: Any) -> bool:
+        return self == other
+
+    @cached_property
+    def positions(self) -> ilist.IList[tuple[float, float], Any]:
+        positions = tuple(
+            (x, y)
+            for (ix, x), (iy, y) in product(
+                enumerate(self.x_positions), enumerate(self.y_positions)
+            )
+            if (ix, iy) not in self.vacancies
+        )
+
+        return ilist.IList(positions)
+
+    @classmethod
+    def fill(
+        cls, grid_obj: grid.Grid[NumX, NumY], filled: Sequence[tuple[int, int]]
+    ) -> "FilledGrid[NumX, NumY]":
+        num_x, num_y = grid_obj.shape
+
+        if isinstance(grid_obj, FilledGrid):
+            vacancies = grid_obj.vacancies
+            parent = grid_obj.parent
+        else:
+            vacancies = frozenset(product(range(num_x), range(num_y)))
+            parent = grid_obj
+
+        vacancies = vacancies - frozenset(filled)
+
+        return cls(parent=parent, vacancies=vacancies)
+
+    @classmethod
+    def vacate(
+        cls, grid_obj: grid.Grid[NumX, NumY], vacancies: Iterable[tuple[int, int]]
+    ) -> "FilledGrid[NumX, NumY]":
+
+        if isinstance(grid_obj, FilledGrid):
+            input_vacancies = grid_obj.vacancies
+            parent = grid_obj.parent
+        else:
+            input_vacancies = frozenset()
+            parent = grid_obj
+
+        input_vacancies = input_vacancies.union(vacancies)
+
+        return cls(parent=parent, vacancies=input_vacancies)
+
+    def get_view(  # type: ignore
+        self, x_indices: ilist.IList[int, Any], y_indices: ilist.IList[int, Any]
+    ):
+        remapping_x = {ix: i for i, ix in enumerate(x_indices)}
+        remapping_y = {iy: i for i, iy in enumerate(y_indices)}
+        return FilledGrid(
+            parent=self.parent.get_view(x_indices, y_indices),
+            vacancies=frozenset(
+                (remapping_x[x], remapping_y[y])
+                for x, y in self.vacancies
+                if x in remapping_x and y in remapping_y
+            ),
+        )
+
+    def shift(self, x_shift: float, y_shift: float):
+        return FilledGrid(
+            parent=self.parent.shift(x_shift, y_shift),
+            vacancies=self.vacancies,
+        )
+
+    def scale(self, x_scale: float, y_scale: float):
+        return FilledGrid(
+            parent=self.parent.scale(x_scale, y_scale),
+            vacancies=self.vacancies,
+        )
+
+    def repeat(self, x_times: int, y_times: int, x_gap: float, y_gap: float):
+        new_parent = self.parent.repeat(x_times, y_times, x_gap, y_gap)
+        x_dim, y_dim = self.shape
+        vacancies = frozenset(
+            (x + x_dim * i, y + y_dim * j)
+            for i, j, (x, y) in product(range(x_times), range(y_times), self.vacancies)
+        )
+        return FilledGrid.vacate(new_parent, vacancies)
+
+    def row_x_pos(self, row_index: int):
+        x_vacancies = {x for x, y in self.vacancies if y == row_index}
+        return ilist.IList(
+            [x for i, x in enumerate(self.parent.x_positions) if i not in x_vacancies]
+        )
+
+    def col_y_pos(self, column_index: int):
+        y_vacancies = {y for x, y in self.vacancies if x == column_index}
+        return ilist.IList(
+            [y for i, y in enumerate(self.y_positions) if i not in y_vacancies]
+        )
+
+    def serialize(self, serializer: "Serializer") -> "SerializationUnit":
+        return SerializationUnit(
+            kind="filled_grid",
+            module_name=dialect.name,
+            class_name=FilledGrid.__name__,
+            data={
+                "parent": serializer.serialize_attribute(self.parent),
+                "vacancies": serializer.serialize(self.vacancies),
+            },
+        )
+
+    @classmethod
+    def deserialize(
+        cls, serUnit: "SerializationUnit", deserializer: "Deserializer"
+    ) -> "FilledGrid":
+        parent = deserializer.deserialize(serUnit.data["parent"])
+        vacancies = deserializer.deserialize(serUnit.data["vacancies"])
+        return FilledGrid(parent=parent, vacancies=vacancies)
+
+
+FilledGridType = types.Generic(FilledGrid, types.TypeVar("NumX"), types.TypeVar("NumY"))

bloqade/geometry/dialects/grid/__init__.py

@@ -0,0 +1,39 @@
+from ._dialect import dialect as dialect
+from ._interface import (
+    col_ypos as col_ypos,
+    from_positions as from_positions,
+    get as get,
+    get_xpos as get_xpos,
+    get_ypos as get_ypos,
+    new as new,
+    positions as positions,
+    repeat as repeat,
+    row_xpos as row_xpos,
+    scale as scale,
+    shape as shape,
+    shift as shift,
+    sub_grid as sub_grid,
+    x_bounds as x_bounds,
+    y_bounds as y_bounds,
+)
+from ._typeinfer import TypeInferMethods as TypeInferMethods
+from .concrete import GridInterpreter as GridInterpreter
+from .stmts import (
+    FromPositions as FromPositions,
+    FromRanges as FromRanges,
+    Get as Get,
+    GetSubGrid as GetSubGrid,
+    GetXBounds as GetXBounds,
+    GetXPos as GetXPos,
+    GetYBounds as GetYBounds,
+    GetYPos as GetYPos,
+    New as New,
+    Positions as Positions,
+    Repeat as Repeat,
+    Scale as Scale,
+    Shape as Shape,
+    Shift as Shift,
+    ShiftSubgridX as ShiftSubgridX,
+    ShiftSubgridY as ShiftSubgridY,
+)
+from .types import Grid as Grid, GridType as GridType

bloqade/geometry/dialects/grid/_dialect.py

@@ -0,0 +1,3 @@
+from kirin.ir import Dialect
+
+dialect = Dialect(name="grid")

bloqade/geometry/dialects/grid/_interface.py

@@ -0,0 +1,309 @@
+import typing
+
+from kirin.dialects import ilist
+from kirin.lowering import wraps as _wraps
+
+from .stmts import (
+    ColYPos,
+    FromPositions,
+    Get,
+    GetSubGrid,
+    GetXBounds,
+    GetXPos,
+    GetYBounds,
+    GetYPos,
+    New,
+    Positions,
+    Repeat,
+    RowXPos,
+    Scale,
+    Shape,
+    Shift,
+    ShiftSubgridX,
+    ShiftSubgridY,
+)
+from .types import Grid
+
+
+@_wraps(New)
+def new(
+    x_spacing: ilist.IList[float, typing.Any] | list[float],
+    y_spacing: ilist.IList[float, typing.Any] | list[float],
+    x_init: float,
+    y_init: float,
+) -> Grid[typing.Any, typing.Any]:
+    """
+    Create a new grid with the given spacing and initial position.
+
+    Args:
+        x_spacing (IList[float] | list[float]): The spacing in the x direction.
+        y_spacing (IList[float] | list[float]): The spacing in the y direction.
+        x_init (float): The initial position in the x direction.
+        y_init (float): The initial position in the y direction.
+
+    Returns:
+        Grid: A new grid object.
+    """
+    ...
+
+
+Nx = typing.TypeVar("Nx")
+Ny = typing.TypeVar("Ny")
+
+
+@typing.overload
+def from_positions(
+    x_positions: ilist.IList[float, Nx], y_positions: ilist.IList[float, Ny]
+) -> Grid[Nx, Ny]: ...
+@typing.overload
+def from_positions(
+    x_positions: ilist.IList[float, Nx], y_positions: list[float]
+) -> Grid[Nx, typing.Any]: ...
+@typing.overload
+def from_positions(
+    x_positions: list[float], y_positions: ilist.IList[float, Ny]
+) -> Grid[typing.Any, Ny]: ...
+@typing.overload
+def from_positions(
+    x_positions: list[float], y_positions: list[float]
+) -> Grid[typing.Any, typing.Any]: ...
+@_wraps(FromPositions)
+def from_positions(x_positions, y_positions):
+    """Construct a grid from the given x and y positions.
+
+    Args:
+        x_positions (IList[float] | list[float]): A list or ilist of floats representing the x-coordinates of grid points.
+        y_positions (IList[float] | list[float]): A list or ilist of floats representing the y-coordinates of grid points.
+
+    Returns:
+        Grid: a grid object
+    """
+
+
+@_wraps(Get)
+def get(grid: Grid, idx: tuple[int, int]) -> tuple[float, float]:
+    """Get the coordinate (x, y) of a grid at the given index.
+
+    Args:
+        grid (Grid): a grid object
+        idx (tuple[int, int]): a tuple of (x, y) indices
+    Returns:
+        tuple[float, float]: a tuple of (x, y) positions
+        tuple[None, None]: if the grid has no initial x or y position
+    """
+    ...
+
+
+@_wraps(GetXPos)
+def get_xpos(grid: Grid[Nx, typing.Any]) -> ilist.IList[float, Nx]:
+    """Get the x positions of a grid.
+
+    Args:
+        grid: a grid object
+    Returns:
+        ilist.IList[float, typing.Any]: a list of x positions
+    """
+    ...
+
+
+@_wraps(GetYPos)
+def get_ypos(grid: Grid[typing.Any, Ny]) -> ilist.IList[float, Ny]:
+    """Get the y positions of a grid.
+
+    Args:
+        grid: a grid object
+    Returns:
+        ilist.IList[float, typing.Any]: a list of y positions
+    """
+    ...
+
+
+@typing.overload
+def sub_grid(
+    grid: Grid, x_indices: ilist.IList[int, Nx], y_indices: ilist.IList[int, Ny]
+) -> Grid[Nx, Ny]: ...
+@typing.overload
+def sub_grid(
+    grid: Grid, x_indices: ilist.IList[int, Nx], y_indices: list[int]
+) -> Grid[Nx, typing.Any]: ...
+@typing.overload
+def sub_grid(
+    grid: Grid, x_indices: list[int], y_indices: ilist.IList[int, Ny]
+) -> Grid[typing.Any, Ny]: ...
+@typing.overload
+def sub_grid(
+    grid: Grid, x_indices: list[int], y_indices: list[int]
+) -> Grid[typing.Any, typing.Any]: ...
+@_wraps(GetSubGrid)
+def sub_grid(grid, x_indices, y_indices):
+    """Get a subgrid from the given grid.
+
+    Args:
+        grid (Grid): a grid object
+        x_indices: a list/ilist of x indices
+        y_indices: a list/ilist of y indices
+    Returns:
+        Grid: a subgrid object
+    """
+    ...
+
+
+@_wraps(GetXBounds)
+def x_bounds(grid: Grid[typing.Any, typing.Any]) -> tuple[float, float]:
+    """Get the x bounds of a grid.
+
+    Args:
+        grid (Grid): a grid object
+    Returns:
+        tuple[float, float]: a tuple of (min_x, max_x)
+    """
+    ...
+
+
+@_wraps(GetYBounds)
+def y_bounds(grid: Grid[typing.Any, typing.Any]) -> tuple[float, float]:
+    """Get the y bounds of a grid.
+
+    Args:
+        grid (Grid): a grid object
+    Returns:
+        tuple[float, float]: a tuple of (min_y, max_y)
+        tuple[None, None]: if the grid has no initial y position
+    """
+    ...
+
+
+@_wraps(Positions)
+def positions(
+    grid: Grid[typing.Any, typing.Any],
+) -> ilist.IList[tuple[float, float], typing.Any]:
+    """Get the positions of a grid as a list of (x, y) tuples.
+
+    Args:
+        grid (Grid): a grid object
+
+    Returns:
+        ilist.IList[tuple[float, float], typing.Any]: a list of (x, y) tuples representing the positions of the grid points
+
+    """
+    ...
+
+
+@_wraps(Repeat)
+def repeat(
+    grid: Grid, x_times: int, y_times: int, x_spacing: float, y_spacing: float
+) -> Grid:
+    """Repeat a grid in the x and y directions.
+
+    Args:
+        grid (Grid): a grid object
+        x_times (int): number of times to repeat in the x direction
+        y_times (int): number of times to repeat in the y direction
+        x_spacing (float): spacing in the x direction
+        y_spacing (float): spacing in the y direction
+    Returns:
+        Grid: a new grid object with the repeated pattern
+    """
+    ...
+
+
+@_wraps(Scale)
+def scale(grid: Grid[Nx, Ny], x_scale: float, y_scale: float) -> Grid[Nx, Ny]:
+    """Scale a grid in the x and y directions.
+
+    Args:
+        grid (Grid): a grid object
+        x_scale (float): scaling factor in the x direction
+        y_scale (float): scaling factor in the y direction
+    Returns:
+        Grid: a new grid object that has been scaled
+    """
+    ...
+
+
+@_wraps(Shift)
+def shift(grid: Grid[Nx, Ny], x_shift: float, y_shift: float) -> Grid[Nx, Ny]:
+    """Shift a grid in the x and y directions.
+
+    Args:
+        grid (Grid): a grid object
+        x_shift (float): shift in the x direction
+        y_shift (float): shift in the y direction
+    Returns:
+        Grid: a new grid object that has been shifted
+    """
+    ...
+
+
+@_wraps(ShiftSubgridX)
+def shift_subgrid_x(
+    grid: Grid[Nx, Ny], x_indices: ilist.IList[int, typing.Any], x_shift: float
+) -> Grid[Nx, Ny]:
+    """Shift a sub grid of grid in the x directions.
+
+    Args:
+        grid (Grid): a grid object
+        x_indices (ilist.IList[int, typing.Any]): a list/ilist of x indices to shift
+        x_shift (float): shift in the x direction
+    Returns:
+        Grid: a new grid object that has been shifted
+    """
+    ...
+
+
+@_wraps(ShiftSubgridY)
+def shift_subgrid_y(
+    grid: Grid[Nx, Ny], y_indices: ilist.IList[int, typing.Any], y_shift: float
+) -> Grid[Nx, Ny]:
+    """Shift a sub grid of grid in the y directions.
+
+    Args:
+        grid (Grid): a grid object
+        y_indices (ilist.IList[int, typing.Any]): a list/ilist of y indices to shift
+        y_shift (float): shift in the y direction
+    Returns:
+        Grid: a new grid object that has been shifted
+    """
+    ...
+
+
+@_wraps(Shape)
+def shape(grid: Grid) -> tuple[int, int]:
+    """Get the shape of a grid.
+
+    Args:
+        grid (Grid): a grid object
+    Returns:
+        tuple[int, int]: a tuple of (num_x, num_y)
+    """
+    ...
+
+
+@_wraps(RowXPos)
+def row_xpos(
+    grid: Grid[typing.Any, typing.Any], row_index: int
+) -> ilist.IList[float, typing.Any]:
+    """Get the x positions of a specific row in the grid.
+
+    Args:
+        grid (Grid): a grid object
+        row_index (int): the index of the row to get x positions for
+    Returns:
+        ilist.IList[float, typing.Any]: a list of x positions for the specified row
+    """
+    ...
+
+
+@_wraps(ColYPos)
+def col_ypos(
+    grid: Grid[typing.Any, typing.Any], column_index: int
+) -> ilist.IList[float, typing.Any]:
+    """Get the y positions of a specific column in the grid.
+
+    Args:
+        grid (Grid): a grid object
+        column_index (int | None): the index of the column to get y positions for, or None for all columns
+    Returns:
+        ilist.IList[float, typing.Any]: a list of y positions for the specified column
+    """
+    ...