From 953b70360440f586b710eff15d359b5eeb31fdeb Mon Sep 17 00:00:00 2001
From: Phillip Weinberg <weinbe58@gmail.com>
Date: Wed, 27 Sep 2023 16:30:59 -0400
Subject: [PATCH] adding dispatch.

---
 src/bloqade/ir/location/transform.py | 35 +++++++++++-----------------
 1 file changed, 13 insertions(+), 22 deletions(-)

diff --git a/src/bloqade/ir/location/transform.py b/src/bloqade/ir/location/transform.py
index 41e7a2e4f..5916d0c31 100644
--- a/src/bloqade/ir/location/transform.py
+++ b/src/bloqade/ir/location/transform.py
@@ -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
 
@@ -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
     ):
@@ -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,
@@ -86,11 +87,11 @@ 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,
         )
@@ -98,11 +99,11 @@ def _add_numpy_position(
     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.
 
@@ -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(