Skip to content

Commit c73644a

Browse files
committed
feat(templates): Add unified and extensible lattice module
- Introduces an AbstractLattice base class to define the core API for all lattice structures. - Implements TILattice for translationally invariant systems (e.g., Square, Honeycomb, Kagome) with efficient, template-based neighbor finding for periodic boundary conditions. - Implements CustomizeLattice for arbitrary geometries, using KDTree for fast neighbor searches and supporting dynamic modifications like adding/removing sites. - Provides a comprehensive test suite covering functionality, robustness, and physical correctness for all lattice types. - Includes a benchmark script demonstrating the performance advantage of the chosen neighbor-finding algorithms. Resolves #12
1 parent 48c3726 commit c73644a

File tree

6 files changed

+3239
-0
lines changed

6 files changed

+3239
-0
lines changed

.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,3 +31,4 @@ docs/source/locale/zh/LC_MESSAGES/textbook.po
3131
docs/source/locale/zh/LC_MESSAGES/whitepapertoc_cn.po
3232
docs/source/locale/zh/LC_MESSAGES/textbooktoc.po
3333
test.qasm
34+
venv/

CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# Change Log
22

33
## Unreleased
4+
- Add `Lattice` module (`tensorcircuit.templates.lattice`) for creating and manipulating various lattice geometries, including `SquareLattice`, `HoneycombLattice`, and `CustomizeLattice`.
45

56
## v1.2.1
67

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
"""
2+
An example script to benchmark neighbor-finding algorithms in CustomizeLattice.
3+
4+
This script demonstrates the performance difference between the KDTree-based
5+
neighbor search and a baseline all-to-all distance matrix method.
6+
As shown by the results, the KDTree approach offers a significant speedup,
7+
especially when calculating for a large number of neighbor shells (large max_k).
8+
9+
To run this script from the project's root directory:
10+
python examples/templates/lattice_neighbor_benchmark.py
11+
"""
12+
13+
import timeit
14+
from typing import Any, Dict, List
15+
16+
17+
def run_benchmark() -> None:
18+
"""
19+
Executes the benchmark test and prints the results in a formatted table.
20+
"""
21+
# --- Benchmark Parameters ---
22+
# A list of lattice sizes (N = number of sites) to test
23+
site_counts: List[int] = [10, 50, 100, 200, 500, 1000, 1500, 2000]
24+
25+
# Use a large k to better showcase the performance of KDTree in
26+
# finding multiple neighbor shells, as suggested by the maintainer.
27+
max_k: int = 2000
28+
29+
# Reduce the number of runs to keep the total benchmark time reasonable,
30+
# especially with a large max_k.
31+
number_of_runs: int = 3
32+
# --------------------------
33+
34+
results: List[Dict[str, Any]] = []
35+
36+
print("=" * 75)
37+
print("Starting neighbor finding benchmark for CustomizeLattice...")
38+
print(f"Parameters: max_k={max_k}, number_of_runs={number_of_runs}")
39+
print("=" * 75)
40+
print(
41+
f"{'Sites (N)':>10} | {'KDTree Time (s)':>18} | {'Baseline Time (s)':>20} | {'Speedup':>10}"
42+
)
43+
print("-" * 75)
44+
45+
for n_sites in site_counts:
46+
# Prepare the setup code for timeit.
47+
# This code generates a random lattice and is executed before timing begins.
48+
# We use a fixed seed to ensure the coordinates are the same for both tests.
49+
setup_code = f"""
50+
import numpy as np
51+
from tensorcircuit.templates.lattice import CustomizeLattice
52+
53+
np.random.seed(42)
54+
coords = np.random.rand({n_sites}, 2)
55+
ids = list(range({n_sites}))
56+
lat = CustomizeLattice(dimensionality=2, identifiers=ids, coordinates=coords)
57+
"""
58+
# Define the Python statements to be timed.
59+
stmt_kdtree = f"lat._build_neighbors(max_k={max_k})"
60+
stmt_baseline = f"lat._build_neighbors_by_distance_matrix(max_k={max_k})"
61+
62+
try:
63+
# Execute the timing. timeit returns the total time for all runs.
64+
time_kdtree = (
65+
timeit.timeit(stmt=stmt_kdtree, setup=setup_code, number=number_of_runs)
66+
/ number_of_runs
67+
)
68+
time_baseline = (
69+
timeit.timeit(
70+
stmt=stmt_baseline, setup=setup_code, number=number_of_runs
71+
)
72+
/ number_of_runs
73+
)
74+
75+
# Calculate and store results, handling potential division by zero.
76+
speedup = time_baseline / time_kdtree if time_kdtree > 0 else float("inf")
77+
results.append(
78+
{
79+
"n_sites": n_sites,
80+
"time_kdtree": time_kdtree,
81+
"time_baseline": time_baseline,
82+
"speedup": speedup,
83+
}
84+
)
85+
print(
86+
f"{n_sites:>10} | {time_kdtree:>18.6f} | {time_baseline:>20.6f} | {speedup:>9.2f}x"
87+
)
88+
89+
except Exception as e:
90+
print(f"An error occurred at N={n_sites}: {e}")
91+
break
92+
93+
print("-" * 75)
94+
print("Benchmark complete.")
95+
96+
97+
if __name__ == "__main__":
98+
run_benchmark()

tensorcircuit/templates/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,5 +5,6 @@
55
from . import graphs
66
from . import measurements
77
from . import conversions
8+
from . import lattice
89

910
costfunctions = measurements

0 commit comments

Comments
 (0)