fix according to the review

Author: Stellogic

examples/lattice_neighbor_benchmark.py

@@ -1,95 +1,111 @@
 """
-An example script to benchmark neighbor-finding algorithms in CustomizeLattice.
-
-This script demonstrates the performance difference between the KDTree-based
-neighbor search and a baseline all-to-all distance matrix method.
-As shown by the results, the KDTree approach offers a significant speedup,
-especially when calculating for a large number of neighbor shells (large max_k).
+Benchmark: Compare neighbor-building time between KDTree and distance-matrix
+methods in CustomizeLattice for varying lattice sizes.
 """
 
-import timeit
-from typing import Any, Dict, List
+import argparse
+import csv
+import time
+from typing import Iterable, List, Tuple, Optional
+import logging
+
+import numpy as np
+
+# Silence verbose infos from the library during benchmarks
+
+logging.basicConfig(level=logging.WARNING)
+
+from tensorcircuit.templates.lattice import CustomizeLattice
+
+
+def run_once(
+    n: int, d: int, max_k: int, repeats: int, seed: int
+) -> Tuple[float, float]:
+    """Run one size point and return (time_kdtree, time_matrix)."""
+    rng = np.random.default_rng(seed)
+    ids = list(range(n))
+
+    # Collect times for each repeat with different random coordinates
+    kdtree_times: List[float] = []
+    matrix_times: List[float] = []
+
+    for _ in range(repeats):
+        # Generate different coordinates for each repeat
+        coords = rng.random((n, d), dtype=float)
+        lat = CustomizeLattice(dimensionality=d, identifiers=ids, coordinates=coords)
+
+        # KDTree path - single measurement
+        t0 = time.perf_counter()
+        lat._build_neighbors(max_k=max_k, use_kdtree=True)
+        kdtree_times.append(time.perf_counter() - t0)
 
+        # Distance-matrix path - single measurement
+        t0 = time.perf_counter()
+        lat._build_neighbors(max_k=max_k, use_kdtree=False)
+        matrix_times.append(time.perf_counter() - t0)
 
-def run_benchmark() -> None:
-    """
-    Executes the benchmark test and prints the results in a formatted table.
-    """
-    # --- Benchmark Parameters ---
-    # A list of lattice sizes (N = number of sites) to test
-    site_counts: List[int] = [10, 50, 100, 200, 500, 1000, 1500, 2000]
+    return float(np.mean(kdtree_times)), float(np.mean(matrix_times))
 
-    # Use a large k to better showcase the performance of KDTree in
-    # finding multiple neighbor shells, as suggested by the maintainer.
-    max_k: int = 2000
 
-    # Reduce the number of runs to keep the total benchmark time reasonable,
-    # especially with a large max_k.
-    number_of_runs: int = 3
-    # --------------------------
+def parse_sizes(s: str) -> List[int]:
+    return [int(x) for x in s.split(",") if x.strip()]
 
-    results: List[Dict[str, Any]] = []
 
-    print("=" * 75)
-    print("Starting neighbor finding benchmark for CustomizeLattice...")
-    print(f"Parameters: max_k={max_k}, number_of_runs={number_of_runs}")
-    print("=" * 75)
+def format_row(n: int, t_kdtree: float, t_matrix: float) -> str:
+    speedup = (t_matrix / t_kdtree) if t_kdtree > 0 else float("inf")
+    return f"{n:>8} | {t_kdtree:>12.6f} | {t_matrix:>14.6f} | {speedup:>7.2f}x"
+
+
+def main(argv: Optional[Iterable[str]] = None) -> int:
+    p = argparse.ArgumentParser(description="Neighbor-building time comparison")
+    p.add_argument(
+        "--sizes",
+        type=parse_sizes,
+        default=[128, 256, 512, 1024, 2048],
+        help="Comma-separated site counts to benchmark (default: 128,256,512,1024,2048)",
+    )
+    p.add_argument(
+        "--dims", type=int, default=2, help="Lattice dimensionality (default: 2)"
+    )
+    p.add_argument(
+        "--max-k", type=int, default=6, help="Max neighbor shells k (default: 6)"
+    )
+    p.add_argument(
+        "--repeats", type=int, default=5, help="Repeats per measurement (default: 5)"
+    )
+    p.add_argument("--seed", type=int, default=42, help="PRNG seed (default: 42)")
+    p.add_argument("--csv", type=str, default="", help="Optional CSV output path")
+    args = p.parse_args(list(argv) if argv is not None else None)
+
+    print("=" * 74)
     print(
-        f"{'Sites (N)':>10} | {'KDTree Time (s)':>18} | {'Baseline Time (s)':>20} | {'Speedup':>10}"
+        f"Benchmark CustomizeLattice neighbor-building | dims={args.dims} max_k={args.max_k} repeats={args.repeats}"
     )
-    print("-" * 75)
+    print("=" * 74)
+    print(f"{'N':>8} | {'KDTree(s)':>12} | {'DistMatrix(s)':>14} | {'Speedup':>7}")
+    print("-" * 74)
 
-    for n_sites in site_counts:
-        # Prepare the setup code for timeit.
-        # This code generates a random lattice and is executed before timing begins.
-        # We use a fixed seed to ensure the coordinates are the same for both tests.
-        setup_code = f"""
-import numpy as np
-from tensorcircuit.templates.lattice import CustomizeLattice
+    rows: List[Tuple[int, float, float]] = []
+    for n in args.sizes:
+        t_kdtree, t_matrix = run_once(n, args.dims, args.max_k, args.repeats, args.seed)
+        rows.append((n, t_kdtree, t_matrix))
+        print(format_row(n, t_kdtree, t_matrix))
 
-np.random.seed(42)
-coords = np.random.rand({n_sites}, 2)
-ids = list(range({n_sites}))
-lat = CustomizeLattice(dimensionality=2, identifiers=ids, coordinates=coords)
-"""
-        # Define the Python statements to be timed.
-        stmt_kdtree = f"lat._build_neighbors(max_k={max_k})"
-        stmt_baseline = f"lat._build_neighbors_by_distance_matrix(max_k={max_k})"
-
-        try:
-            # Execute the timing. timeit returns the total time for all runs.
-            time_kdtree = (
-                timeit.timeit(stmt=stmt_kdtree, setup=setup_code, number=number_of_runs)
-                / number_of_runs
-            )
-            time_baseline = (
-                timeit.timeit(
-                    stmt=stmt_baseline, setup=setup_code, number=number_of_runs
-                )
-                / number_of_runs
-            )
-
-            # Calculate and store results, handling potential division by zero.
-            speedup = time_baseline / time_kdtree if time_kdtree > 0 else float("inf")
-            results.append(
-                {
-                    "n_sites": n_sites,
-                    "time_kdtree": time_kdtree,
-                    "time_baseline": time_baseline,
-                    "speedup": speedup,
-                }
-            )
-            print(
-                f"{n_sites:>10} | {time_kdtree:>18.6f} | {time_baseline:>20.6f} | {speedup:>9.2f}x"
-            )
-
-        except Exception as e:
-            print(f"An error occurred at N={n_sites}: {e}")
-            break
-
-    print("-" * 75)
-    print("Benchmark complete.")
+    if args.csv:
+        with open(args.csv, "w", newline="", encoding="utf-8") as f:
+            w = csv.writer(f)
+            w.writerow(["N", "time_kdtree_s", "time_distance_matrix_s", "speedup"])
+            for n, t_kdtree, t_matrix in rows:
+                speedup = (t_matrix / t_kdtree) if t_kdtree > 0 else float("inf")
+                w.writerow([n, f"{t_kdtree:.6f}", f"{t_matrix:.6f}", f"{speedup:.2f}"])
+
+        print("-" * 74)
+        print(f"Saved CSV to: {args.csv}")
+
+    print("-" * 74)
+    print("Done.")
+    return 0
 
 
 if __name__ == "__main__":
-    run_benchmark()
+    raise SystemExit(main())

examples/lattice_neighbor_time_compare.py

@@ -5,8 +5,10 @@
 to optimize crystal structure. It finds the equilibrium lattice constant that minimizes
 the total Lennard-Jones potential energy of a 2D square lattice.
 
-The optimization showcases the key Task 3 capability: making lattice parameters
-differentiable for variational material design.
+This example showcases a key capability of the differentiable lattice system: 
+making geometric parameters (like lattice constants) fully differentiable and 
+optimizable using automatic differentiation. This enables variational material design
+where crystal structures can be optimized to minimize physical energy functions.
 """
 
 import optax
@@ -54,7 +56,7 @@ def calculate_potential(log_a, epsilon=0.5, sigma=1.0):
 
 optimizer = optax.adam(learning_rate=0.01)
 
-log_a = K.convert_to_tensor(K.log(K.convert_to_tensor(1.1)))
+log_a = K.convert_to_tensor(K.log(K.convert_to_tensor(2.0)))
 
 opt_state = optimizer.init(log_a)
 

examples/lennard_jones_optimization.py

@@ -629,7 +629,7 @@ def all(self: Any, a: Tensor, axis: Optional[Sequence[int]] = None) -> Tensor:
 
     def meshgrid(self: Any, *args: Any, **kwargs: Any) -> Any:
         """
-         Return coordinate matrices from coordinate vectors.
+        Return coordinate matrices from coordinate vectors.
 
         :param args: coordinate vectors
         :type args: Any

tensorcircuit/backends/abstract_backend.py

@@ -18,11 +18,11 @@
     Set,
 )
 
-logger = logging.getLogger(__name__)
 import itertools
 import math
 import numpy as np
 from scipy.spatial import KDTree
+
 from .. import backend
 
 
@@ -42,6 +42,8 @@
     import matplotlib.axes
     from mpl_toolkits.mplot3d import Axes3D
 
+logger = logging.getLogger(__name__)
+
 Tensor = Any
 SiteIndex = int
 SiteIdentifier = Hashable
@@ -835,7 +837,7 @@ def _build_neighbors(self, max_k: int = 2, **kwargs: Any) -> None:
         :param max_k: The maximum order of neighbors to compute (e.g., k=1 for
             nearest neighbors, k=2 for next-nearest, etc.). Defaults to 2.
         :type max_k: int, optional
-        :param \**kwargs: Additional keyword arguments. May include:
+        :param kwargs: Additional keyword arguments. May include:
             - ``tol`` (float): The numerical tolerance used to determine if two
               distances are equal when identifying shells. Defaults to 1e-6.
         """
@@ -1486,7 +1488,7 @@ def _build_neighbors_kdtree(self, max_k: int, tol: float) -> None:
 
         # For small lattices or cases with potential duplicate coordinates,
         # fall back to distance matrix method for robustness
-        if self.num_sites < 1000:
+        if self.num_sites < 200:
             logger.info(
                 "Small lattice detected, falling back to distance matrix method for robustness"
             )