diff --git a/src/braket/ahs/atom_arrangement.py b/src/braket/ahs/atom_arrangement.py index 24d4fc9aa..9329694ba 100644 --- a/src/braket/ahs/atom_arrangement.py +++ b/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,40 +66,41 @@ 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) @@ -96,28 +108,29 @@ def coordinate_list(self, coordinate_index: Number) -> list[Number]: 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") + diff --git a/test/unit_tests/braket/ahs/test_atom_arrangement.py b/test/unit_tests/braket/ahs/test_atom_arrangement.py index 06a926163..f9f5e5ff1 100644 --- a/test/unit_tests/braket/ahs/test_atom_arrangement.py +++ b/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 +