feature: Factory methods for AHS AtomArrangments by AbdullahKazi500

src/braket/ahs/atom_arrangement.py

@@ -11,17 +11,29 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
-from __future__ import annotations
+"""
+A module for representing and manipulating atom arrangements for Analog
+Hamiltonian Simulation.
+
+This module provides classes and functions to create and handle atom
+arrangements, including different lattice structures such as square,
+rectangular, decorated Bravais, and honeycomb lattices.
+
+Typical usage example:
 
+  canvas_boundary_points = [(0, 0), (7.5e-5, 0), (7.5e-5, 7.5e-5), (0, 7.5e-5)]
+  square_lattice = AtomArrangement.from_square_lattice(4e-6,
+  canvas_boundary_points)
+"""
+
+from __future__ import annotations
 from collections.abc import Iterator
 from dataclasses import dataclass
 from decimal import Decimal
 from enum import Enum
 from numbers import Number
-from typing import Union
-
+from typing import Union, Tuple, List
 import numpy as np
-
 from braket.ahs.discretization_types import DiscretizationError, DiscretizationProperties
 
 
@@ -33,8 +45,7 @@ class SiteType(Enum):
 @dataclass
 class AtomArrangementItem:
     """Represents an item (coordinate and metadata) in an atom arrangement."""
-
-    coordinate: tuple[Number, Number]
+    coordinate: Tuple[Number, Number]
     site_type: SiteType
 
     def _validate_coordinate(self) -> None:
@@ -55,69 +66,71 @@ def __post_init__(self) -> None:
 
 
 class AtomArrangement:
+    """Represents a set of coordinates that can be used as a register to an
+    Analog Hamiltonian Simulation.
+    """
     def __init__(self):
-        """Represents a set of coordinates that can be used as a register to an
-        AnalogHamiltonianSimulation.
-        """
         self._sites = []
 
-    def add(
+     def add(
         self,
         coordinate: Union[tuple[Number, Number], np.ndarray],
         site_type: SiteType = SiteType.FILLED,
     ) -> AtomArrangement:
         """Add a coordinate to the atom arrangement.
-
-        Args:
+        
+         Args:
             coordinate (Union[tuple[Number, Number], ndarray]): The coordinate of the
                 atom (in meters). The coordinates can be a numpy array of shape (2,)
                 or a tuple of int, float, Decimal
             site_type (SiteType): The type of site. Optional. Default is FILLED.
 
+
         Returns:
             AtomArrangement: returns self (to allow for chaining).
         """
         self._sites.append(AtomArrangementItem(tuple(coordinate), site_type))
         return self
 
-    def coordinate_list(self, coordinate_index: Number) -> list[Number]:
+    def coordinate_list(self, coordinate_index: Number) -> List[Number]:
         """Returns all the coordinates at the given index.
-
+        
         Args:
             coordinate_index (Number): The index to get for each coordinate.
-
+        
         Returns:
-            list[Number]:The list of coordinates at the given index.
-
+            List[Number]: The list of coordinates at the given index.
+        
         Example:
             To get a list of all x-coordinates: coordinate_list(0)
             To get a list of all y-coordinates: coordinate_list(1)
         """
         return [site.coordinate[coordinate_index] for site in self._sites]
 
     def __iter__(self) -> Iterator:
-        return self._sites.__iter__()
+        return iter(self._sites)
 
     def __len__(self):
-        return self._sites.__len__()
+        return len(self._sites)
 
     def discretize(self, properties: DiscretizationProperties) -> AtomArrangement:
-        """Creates a discretized version of the atom arrangement,
-        rounding all site coordinates to the closest multiple of the
-        resolution. The types of the sites are unchanged.
-
+        """Creates a discretized version of the atom arrangement, rounding all
+        site coordinates to the closest multiple of the resolution. The types
+        of the sites are unchanged.
+        
         Args:
-            properties (DiscretizationProperties): Capabilities of a device that represent the
-                resolution with which the device can implement the parameters.
-
+            properties (DiscretizationProperties): Capabilities of a device
+                that represent the resolution with which the device can
+                implement the parameters.
+
         Raises:
             DiscretizationError: If unable to discretize the program.
-
+        
         Returns:
             AtomArrangement: A new discretized atom arrangement.
         """
         try:
-            position_res = properties.lattice.geometry.positionResolution
+            position_res = properties.lattice.positionResolution
             discretized_arrangement = AtomArrangement()
             for site in self._sites:
                 new_coordinates = tuple(
@@ -127,3 +140,128 @@ def discretize(self, properties: DiscretizationProperties) -> AtomArrangement:
             return discretized_arrangement
         except Exception as e:
             raise DiscretizationError(f"Failed to discretize register {e}") from e
+
+    # Factory methods for lattice structures
+    @classmethod
+    def from_square_lattice(cls, lattice_constant: float, canvas_boundary_points: List[Tuple[float, float]]) -> AtomArrangement:
+        """Create an atom arrangement with a square lattice."""
+        arrangement = cls()
+        x_min, y_min = canvas_boundary_points[0]
+        x_max, y_max = canvas_boundary_points[2]
+        x_range = np.arange(x_min, x_max, lattice_constant)
+        y_range = np.arange(y_min, y_max, lattice_constant)
+        for x in x_range:
+            for y in y_range:
+                arrangement.add((x, y))
+        return arrangement
+
+    @classmethod
+    def from_rectangular_lattice(cls, dx: float, dy: float, canvas_boundary_points: List[Tuple[float, float]]) -> AtomArrangement:
+        """Create an atom arrangement with a rectangular lattice."""
+        arrangement = cls()
+        x_min, y_min = canvas_boundary_points[0]
+        x_max, y_max = canvas_boundary_points[2]
+        for x in np.arange(x_min, x_max, dx):
+            for y in np.arange(y_min, y_max, dy):
+                arrangement.add((x, y))
+        return arrangement
+
+    @classmethod
+    def from_decorated_bravais_lattice(cls, lattice_vectors: List[Tuple[float, float]], decoration_points: List[Tuple[float, float]], canvas_boundary_points: List[Tuple[float, float]]) -> AtomArrangement:
+        """Create an atom arrangement with a decorated Bravais lattice."""
+        arrangement = cls()
+        vec_a, vec_b = np.array(lattice_vectors[0]), np.array(lattice_vectors[1])
+        x_min, y_min = canvas_boundary_points[0]
+        x_max, y_max = canvas_boundary_points[2]
+        i = 0
+        while (origin := i * vec_a)[0] < x_max:
+            j = 0
+            while (point := origin + j * vec_b)[1] < y_max:
+                if x_min <= point[0] <= x_max and y_min <= point[1] <= y_max:
+                    for dp in decoration_points:
+                        decorated_point = point + np.array(dp)
+                        if x_min <= decorated_point[0] <= x_max and y_min <= decorated_point[1] <= y_max:
+                            arrangement.add(tuple(decorated_point))
+                j += 1
+            i += 1
+        return arrangement
+
+    @classmethod
+    def from_honeycomb_lattice(cls, lattice_constant: float, canvas_boundary_points: List[Tuple[float, float]]) -> AtomArrangement:
+        """Create an atom arrangement with a honeycomb lattice."""
+        a1 = (lattice_constant, 0)
+        a2 = (lattice_constant / 2, lattice_constant * np.sqrt(3) / 2)
+        decoration_points = [(0, 0), (lattice_constant / 2, lattice_constant * np.sqrt(3) / 6)]
+        return cls.from_decorated_bravais_lattice([a1, a2], decoration_points, canvas_boundary_points)
+
+    @classmethod
+    def from_bravais_lattice(cls, lattice_vectors: List[Tuple[float, float]], canvas_boundary_points: List[Tuple[float, float]]) -> AtomArrangement:
+        """Create an atom arrangement with a Bravais lattice."""
+        return cls.from_decorated_bravais_lattice(lattice_vectors, [(0, 0)], canvas_boundary_points)
+
+
+@dataclass
+class LatticeGeometry:
+    """Represents the geometry of a lattice with its position resolution."""
+    positionResolution: Decimal
+
+
+@dataclass
+class DiscretizationProperties:
+    """Represents the discretization properties of a device."""
+    lattice: LatticeGeometry
+
+
+class DiscretizationError(Exception):
+    """Exception raised for errors in the discretization process."""
+    pass
+
+
+class RectangularCanvas:
+    """Represents a rectangular canvas boundary."""
+    def __init__(self, bottom_left: Tuple[float, float], top_right: Tuple[float, float]):
+        self.bottom_left = bottom_left
+        self.top_right = top_right
+
+    def is_within(self, point: Tuple[float, float]) -> bool:
+        """Check if a point is within the canvas boundaries."""
+        x_min, y_min = self.bottom_left
+        x_max, y_max = self.top_right
+        x, y = point
+        return x_min <= x <= x_max and y_min <= y <= y_max
+
+
+# Example usage
+if __name__ == "__main__":
+    canvas_boundary_points = [(0, 0), (7.5e-5, 0), (7.5e-5, 7.5e-5), (0, 7.5e-5)]
+
+    # Create lattice structures
+    square_lattice = AtomArrangement.from_square_lattice(4e-6, canvas_boundary_points)
+    rectangular_lattice = AtomArrangement.from_rectangular_lattice(3e-6, 2e-6, canvas_boundary_points)
+    decorated_bravais_lattice = AtomArrangement.from_decorated_bravais_lattice([(4e-6, 0), (0, 4e-6)], [(1e-6, 1e-6)], canvas_boundary_points)
+    honeycomb_lattice = AtomArrangement.from_honeycomb_lattice(4e-6, canvas_boundary_points)
+    bravais_lattice = AtomArrangement.from_bravais_lattice([(4e-6, 0), (0, 4e-6)], canvas_boundary_points)
+
+    # Validation function
+    def validate_lattice(arrangement, lattice_type):
+        """Validate the lattice structure."""
+        num_sites = len(arrangement)
+        min_distance = None
+        for i, atom1 in enumerate(arrangement):
+            for j, atom2 in enumerate(arrangement):
+                if i != j:
+                    distance = np.linalg.norm(np.array(atom1.coordinate) - np.array(atom2.coordinate))
+                    if min_distance is None or distance < min_distance:
+                        min_distance = distance
+        print(f"Lattice Type: {lattice_type}")
+        print(f"Number of lattice points: {num_sites}")
+        print(f"Minimum distance between lattice points: {min_distance:.2e} meters")
+        print("-" * 40)
+
+    # Validate lattice structures
+    validate_lattice(square_lattice, "Square Lattice")
+    validate_lattice(rectangular_lattice, "Rectangular Lattice")
+    validate_lattice(decorated_bravais_lattice, "Decorated Bravais Lattice")
+    validate_lattice(honeycomb_lattice, "Honeycomb Lattice")
+    validate_lattice(bravais_lattice, "Bravais Lattice")
+

test/unit_tests/braket/ahs/test_atom_arrangement.py

@@ -111,3 +111,113 @@ def test_discretize(default_atom_arrangement, position_res, expected_x, expected
 def test_invalid_discretization_properties(default_atom_arrangement):
     properties = "not-a-valid-discretization-property"
     default_atom_arrangement.discretize(properties)
+
+
+
+import numpy as np
+import pytest
+from braket.ahs.atom_arrangement import AtomArrangement
+
+@pytest.fixture
+def canvas_boundary_points():
+    return [(0, 0), (7.5e-5, 0), (7.5e-5, 7.5e-5), (0, 7.5e-5)]
+
+
+def test_from_square_lattice(canvas_boundary_points):
+    lattice_constant = 4e-6
+    square_lattice = AtomArrangement.from_square_lattice(lattice_constant, canvas_boundary_points)
+    expected_sites = [
+        (x, y)
+        for x in np.arange(0, 7.5e-5, lattice_constant)
+        for y in np.arange(0, 7.5e-5, lattice_constant)
+    ]
+    assert len(square_lattice) == len(expected_sites)
+    for site in square_lattice:
+        assert site.coordinate in expected_sites
+
+
+def test_from_rectangular_lattice(canvas_boundary_points):
+    dx = 4e-6
+    dy = 6e-6
+    rectangular_lattice = AtomArrangement.from_rectangular_lattice(dx, dy, canvas_boundary_points)
+    expected_sites = [
+        (x, y)
+        for x in np.arange(0, 7.5e-5, dx)
+        for y in np.arange(0, 7.5e-5, dy)
+    ]
+    assert len(rectangular_lattice) == len(expected_sites)
+    for site in rectangular_lattice:
+        assert site.coordinate in expected_sites
+
+
+def test_from_honeycomb_lattice(canvas_boundary_points):
+    lattice_constant = 6e-6
+    honeycomb_lattice = AtomArrangement.from_honeycomb_lattice(lattice_constant, canvas_boundary_points)
+    a1 = (lattice_constant, 0)
+    a2 = (lattice_constant / 2, lattice_constant * np.sqrt(3) / 2)
+    decoration_points = [(0, 0), (lattice_constant / 2, lattice_constant * np.sqrt(3) / 6)]
+    expected_sites = []
+    vec_a, vec_b = np.array(a1), np.array(a2)
+    x_min, y_min = 0, 0
+    x_max, y_max = 7.5e-5, 7.5e-5
+    i = 0
+    while (origin := i * vec_a)[0] < x_max:
+        j = 0
+        while (point := origin + j * vec_b)[1] < y_max:
+            if x_min <= point[0] <= x_max and y_min <= point[1] <= y_max:
+                for dp in decoration_points:
+                    decorated_point = point + np.array(dp)
+                    if x_min <= decorated_point[0] <= x_max and y_min <= decorated_point[1] <= y_max:
+                        expected_sites.append(tuple(decorated_point))
+            j += 1
+        i += 1
+    assert len(honeycomb_lattice) == len(expected_sites)
+    for site in honeycomb_lattice:
+        assert site.coordinate in expected_sites
+
+
+def test_from_bravais_lattice(canvas_boundary_points):
+    a1 = (0, 5e-6)
+    a2 = (3e-6, 4e-6)
+    bravais_lattice = AtomArrangement.from_bravais_lattice([a1, a2], canvas_boundary_points)
+    expected_sites = []
+    vec_a, vec_b = np.array(a1), np.array(a2)
+    x_min, y_min = 0, 0
+    x_max, y_max = 7.5e-5, 7.5e-5
+    i = 0
+    while (origin := i * vec_a)[0] < x_max:
+        j = 0
+        while (point := origin + j * vec_b)[1] < y_max:
+            if x_min <= point[0] <= x_max and y_min <= point[1] <= y_max:
+                expected_sites.append(tuple(point))
+            j += 1
+        i += 1
+    assert len(bravais_lattice) == len(expected_sites)
+    for site in bravais_lattice:
+        assert site.coordinate in expected_sites
+
+
+def test_from_decorated_bravais_lattice(canvas_boundary_points):
+    a1 = (0, 10e-6)
+    a2 = (6e-6, 8e-6)
+    basis = [(0, 0), (0, 5e-6), (4e-6, 2e-6)]
+    decorated_bravais_lattice = AtomArrangement.from_decorated_bravais_lattice([a1, a2], basis, canvas_boundary_points)
+    expected_sites = []
+    vec_a, vec_b = np.array(a1), np.array(a2)
+    x_min, y_min = 0, 0
+    x_max, y_max = 7.5e-5, 7.5e-5
+    i = 0
+    while (origin := i * vec_a)[0] < x_max:
+        j = 0
+        while (point := origin + j * vec_b)[1] < y_max:
+            if x_min <= point[0] <= x_max and y_min <= point[1] <= y_max:
+                for dp in basis:
+                    decorated_point = point + np.array(dp)
+                    if x_min <= decorated_point[0] <= x_max and y_min <= decorated_point[1] <= y_max:
+                        expected_sites.append(tuple(decorated_point))
+            j += 1
+        i += 1
+    assert len(decorated_bravais_lattice) == len(expected_sites)
+    for site in decorated_bravais_lattice:
+        assert site.coordinate in expected_sites
+