feat: Implement AbstractLattice and initial GeneralLattice

Author: Stellogic

tensorcircuit/templates/lattice.py

@@ -0,0 +1,157 @@
+# -*- coding: utf-8 -*-
+"""
+The lattice module for defining and manipulating lattice geometries.
+"""
+import abc
+import numpy as np
+import matplotlib.pyplot as plt
+from typing import List, Tuple, Dict, Any, Optional, Hashable, Iterator, Union
+
+# --- Type Aliases for Readability ---
+SiteIndex = int
+SiteIdentifier = Hashable
+Coordinates = np.ndarray
+NeighborMap = Dict[SiteIndex, List[SiteIndex]]
+
+
+class AbstractLattice(abc.ABC):
+    """
+    Abstract base class for describing lattice systems.
+
+    This class defines the common interface for all lattice structures,
+    providing access to fundamental properties like site information
+
+    (count, coordinates, identifiers) and neighbor relationships.
+    Subclasses are responsible for implementing the specific logic for
+    generating the lattice points and calculating neighbor connections.
+    """
+
+    def __init__(self, dimensionality: int):
+        """
+        Initializes the base lattice class.
+
+        Args:
+            dimensionality (int): The spatial dimension of the lattice (e.g., 1, 2, 3).
+        """
+        self._dimensionality = dimensionality
+        
+        # --- Internal Data Structures (to be populated by subclasses) ---
+        self._indices: List[SiteIndex] = []
+        self._identifiers: List[SiteIdentifier] = []
+        self._coordinates: List[Coordinates] = []
+        self._ident_to_idx: Dict[SiteIdentifier, SiteIndex] = {}
+        
+        # Neighbor information, structured as {k: {site_index: [neighbor_indices]}}
+        # k=1 for nearest neighbors, k=2 for next-nearest, etc.
+        self._neighbor_maps: Dict[int, NeighborMap] = {}
+
+    @property
+    def num_sites(self) -> int:
+        """Returns the total number of sites (N) in the lattice."""
+        return len(self._indices)
+
+    @property
+    def dimensionality(self) -> int:
+        """Returns the spatial dimension of the lattice."""
+        return self._dimensionality
+
+    def __len__(self) -> int:
+        """Returns the total number of sites, enabling `len(lattice)`."""
+        return self.num_sites
+    
+    # Other common methods like get_site_info, get_neighbors can be added here.
+
+    @abc.abstractmethod
+    def _build_lattice(self, *args, **kwargs):
+        """
+        Abstract method for subclasses to generate the lattice data.
+        
+        This method should populate the following internal lists and dicts:
+        - self._indices
+        - self._identifiers
+        - self._coordinates
+        - self._ident_to_idx
+        """
+        pass
+
+    @abc.abstractmethod
+    def _build_neighbors(self, max_k: int = 1):
+        """
+        Abstract method for subclasses to calculate neighbor relationships.
+
+        This method should populate the `self._neighbor_maps` dictionary.
+
+        Args:
+            max_k (int): The maximum order of neighbors to compute (e.g., max_k=2 for NN and NNN).
+        """
+        pass
+
+    @abc.abstractmethod
+    def show(self, **kwargs):
+        """
+        Abstract method for visualizing the lattice structure.
+        
+        Subclasses should implement the specific plotting logic.
+        """
+        pass
+
+
+class GeneralLattice(AbstractLattice):
+    """
+    A general lattice built from an explicit list of sites and coordinates.
+    """
+    def __init__(self,
+                 dimensionality: int,
+                 identifiers: List[SiteIdentifier],
+                 coordinates: List[Union[List[float], np.ndarray]]):
+        
+        super().__init__(dimensionality)
+        assert len(identifiers) == len(coordinates), "Identifiers and coordinates must have the same length."
+
+        # The lattice is built directly upon initialization.
+        self._build_lattice(identifiers=identifiers, coordinates=coordinates)
+        
+        # Neighbor relationships can be calculated later if needed.
+        print(f"GeneralLattice with {self.num_sites} sites created.")
+
+
+    def _build_lattice(self, identifiers: List[SiteIdentifier], coordinates: List[Union[List[float], np.ndarray]]):
+        """Implements the lattice building for GeneralLattice."""
+        self._identifiers = list(identifiers)
+        self._coordinates = [np.array(c) for c in coordinates]
+        self._indices = list(range(len(identifiers)))
+        self._ident_to_idx = {ident: idx for idx, ident in enumerate(identifiers)}
+
+    def _build_neighbors(self, max_k: int = 1):
+        """Calculates neighbors based on distance (to be implemented)."""
+        print(f"Neighbor calculation for k={max_k} is not implemented yet.")
+        # The logic for _find_neighbors_by_distance will be implemented here in the future.
+        pass
+
+    def show(self, show_indices: bool = True, show_bonds: bool = False, **kwargs):
+        """
+        A simple visualization of the lattice sites.
+        
+        Args:
+            show_indices (bool): If True, display the identifier of each site.
+            show_bonds (bool): If True, display lines connecting neighbors (not implemented yet).
+        """
+        if self.dimensionality != 2:
+            print("Warning: show() is currently only implemented for 2D lattices.")
+            return
+
+        coords = np.array(self._coordinates)
+        
+        plt.figure(figsize=(6, 6))
+        plt.scatter(coords[:, 0], coords[:, 1], s=100, zorder=2)
+
+        if show_indices:
+            for i in range(self.num_sites):
+                plt.text(coords[i, 0] + 0.1, coords[i, 1] + 0.1, str(self._identifiers[i]), fontsize=12)
+        
+        plt.axis('equal')
+        plt.grid(True)
+        plt.title(f"{self.__class__.__name__} ({self.num_sites} sites)")
+        plt.xlabel("x coordinate")
+        plt.ylabel("y coordinate")
+        plt.show()