Skip to content

Commit

Permalink
adding dispatch.
Browse files Browse the repository at this point in the history
  • Loading branch information
weinbe58 committed Sep 27, 2023
1 parent fadefd4 commit 953b703
Showing 1 changed file with 13 additions and 22 deletions.
35 changes: 13 additions & 22 deletions src/bloqade/ir/location/transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from beartype.vale import Is
from typing import Annotated
from beartype import beartype
from plum import dispatch
import numpy as np
from bloqade.ir.scalar import cast

Expand Down Expand Up @@ -46,7 +47,7 @@ def scale(self, scale: ScalarType):

return ListOfLocations(location_list)

@beartype
@dispatch
def _add_position(
self, position: Tuple[ScalarType, ScalarType], filling: Optional[bool] = None
):
Expand All @@ -61,8 +62,8 @@ def _add_position(

return ListOfLocations(location_list)

@beartype
def _add_position_list(
@dispatch
def _add_position( # noqa: F811
self,
position: List[Tuple[ScalarType, ScalarType]],
filling: Optional[List[bool]] = None,
Expand All @@ -86,23 +87,23 @@ def _add_position_list(

return ListOfLocations(location_list)

@beartype
def _add_numpy_position(
@dispatch
def _add_position( # noqa: F811
self, position: PositionArray, filling: Optional[BoolArray] = None
):
return self._add_position_list(
return self.add_position(
list(map(tuple, position.tolist())),
filling.tolist() if filling is not None else None,
)

def add_position(
self,
position: Union[
PositionArray,
Tuple[ScalarType, ScalarType],
List[Tuple[ScalarType, ScalarType]],
PositionArray,
],
filling: Optional[Union[bool, list[bool], BoolArray]] = None,
filling: Optional[Union[BoolArray, List[bool], bool]] = None,
) -> "ListOfLocations":
"""add a position or list of positions to existing atom arrangement.
Expand All @@ -113,25 +114,15 @@ def add_position(
position to add
filling (bool | list[bool]
| numpy.array with shape (n, ) | None, optional):
filling of the added position(s). Defaults to None.
filling of the added position(s). Defaults to None. if None, all
positions are filled.
Returns:
ListOfLocations: new atom arrangement with added positions
"""
if isinstance(position, tuple) and isinstance(filling, (bool, type(None))):
return self._add_position(position, filling)
elif isinstance(position, list) and isinstance(filling, (list, type(None))):
return self._add_position_list(position, filling)
elif isinstance(position, np.ndarray) and isinstance(
filling, (np.ndarray, type(None))
):
return self._add_numpy_position(position, filling)
else:
raise TypeError(
f"cannot interpret arguments, got {type(position)} "
f"for position and {type(filling)} for filling"
)
return self._add_position(position, filling)

@beartype
def apply_defect_count(
Expand Down

0 comments on commit 953b703

Please sign in to comment.