feat(lattice): Make lattice geometries differentiable and backend-agn…

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/lennard_jones_optimization.py

@@ -0,0 +1,121 @@
+"""
+Lennard-Jones Potential Optimization Example
+
+This script demonstrates how to use TensorCircuit's differentiable lattice geometries
+to optimize crystal structure. It finds the equilibrium lattice constant that minimizes
+the total Lennard-Jones potential energy of a 2D square lattice.
+
+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
+import numpy as np
+import matplotlib.pyplot as plt
+import tensorcircuit as tc
+
+
+tc.set_dtype("float64")  # Use tc for universal control
+K = tc.set_backend("jax")
+
+
+def calculate_potential(log_a, epsilon=0.5, sigma=1.0):
+    """
+    Calculate the total Lennard-Jones potential energy for a given logarithm of the lattice constant (log_a).
+    This version creates the lattice inside the function to demonstrate truly differentiable geometry.
+    """
+    lattice_constant = K.exp(log_a)
+
+    # Create lattice with the differentiable parameter
+    size = (4, 4)  # Smaller size for demonstration
+    lattice = tc.templates.lattice.SquareLattice(
+        size, lattice_constant=lattice_constant, pbc=True
+    )
+    d = lattice.distance_matrix
+
+    d_safe = K.where(d > 1e-9, d, K.convert_to_tensor(1e-9))
+
+    term12 = K.power(sigma / d_safe, 12)
+    term6 = K.power(sigma / d_safe, 6)
+    potential_matrix = 4 * epsilon * (term12 - term6)
+
+    num_sites = lattice.num_sites
+    # Zero out self-interactions (diagonal elements)
+    eye_mask = K.eye(num_sites, dtype=potential_matrix.dtype)
+    potential_matrix = potential_matrix * (1 - eye_mask)
+
+    potential_energy = K.sum(potential_matrix) / 2.0
+
+    return potential_energy
+
+
+# Create value and grad function for optimization
+value_and_grad_fun = K.jit(K.value_and_grad(calculate_potential))
+
+optimizer = optax.adam(learning_rate=0.01)
+
+log_a = K.convert_to_tensor(K.log(K.convert_to_tensor(2.0)))
+
+opt_state = optimizer.init(log_a)
+
+history = {"a": [], "energy": []}
+
+print("Starting optimization of lattice constant...")
+for i in range(200):
+    energy, grad = value_and_grad_fun(log_a)
+
+    history["a"].append(K.exp(log_a))
+    history["energy"].append(energy)
+
+    updates, opt_state = optimizer.update(grad, opt_state)
+    log_a = optax.apply_updates(log_a, updates)
+
+    if (i + 1) % 20 == 0:
+        current_a = K.exp(log_a)
+        print(
+            f"Iteration {i+1}/200: Total Energy = {energy:.4f}, Lattice Constant = {current_a:.4f}"
+        )
+
+final_a = K.exp(log_a)
+final_energy = calculate_potential(log_a)
+
+print("\nOptimization finished!")
+print(f"Final optimized lattice constant: {final_a:.6f}")
+print(f"Corresponding minimum total energy: {final_energy:.6f}")
+
+# Vectorized calculation for the potential curve
+a_vals = np.linspace(0.8, 1.5, 200)
+log_a_vals = K.log(K.convert_to_tensor(a_vals))
+
+# Use vmap to create a vectorized version of the potential function
+vmap_potential = K.vmap(lambda la: calculate_potential(la))
+potential_curve = vmap_potential(log_a_vals)
+
+plt.figure(figsize=(10, 6))
+plt.plot(a_vals, potential_curve, label="Lennard-Jones Potential", color="blue")
+plt.scatter(
+    history["a"],
+    history["energy"],
+    color="red",
+    s=20,
+    zorder=5,
+    label="Optimization Steps",
+)
+plt.scatter(
+    final_a,
+    final_energy,
+    color="green",
+    s=100,
+    zorder=6,
+    marker="*",
+    label="Final Optimized Point",
+)
+
+plt.title("Lennard-Jones Potential Optimization")
+plt.xlabel("Lattice Constant (a)")
+plt.ylabel("Total Potential Energy")
+plt.legend()
+plt.grid(True)
+plt.show()