fix according to the review

Author: Stellogic

examples/lattice_neighbor_time_compare.py

@@ -0,0 +1,118 @@
+"""
+Benchmark: Compare neighbor-building time between KDTree and distance-matrix
+methods in CustomizeLattice for varying lattice sizes.
+"""
+
+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 _timeit(fn, repeats: int) -> float:
+    """Return average wall time (seconds) over repeats for calling fn()."""
+    times: List[float] = []
+    for _ in range(repeats):
+        t0 = time.perf_counter()
+        fn()
+        times.append(time.perf_counter() - t0)
+    return float(np.mean(times))
+
+
+def _gen_coords(n: int, d: int, seed: int) -> np.ndarray:
+    rng = np.random.default_rng(seed)
+    return rng.random((n, d), dtype=float)
+
+
+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)."""
+    coords = _gen_coords(n, d, seed)
+    ids = list(range(n))
+    lat = CustomizeLattice(dimensionality=d, identifiers=ids, coordinates=coords)
+
+    # KDTree path
+    t_kdtree = _timeit(
+        lambda: lat._build_neighbors(max_k=max_k, use_kdtree=True), repeats
+    )
+
+    # Distance-matrix path (fully differentiable)
+    t_matrix = _timeit(
+        lambda: lat._build_neighbors(max_k=max_k, use_kdtree=False), repeats
+    )
+
+    return t_kdtree, t_matrix
+
+
+def parse_sizes(s: str) -> List[int]:
+    return [int(x) for x in s.split(",") if x.strip()]
+
+
+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"Benchmark CustomizeLattice neighbor-building | dims={args.dims} max_k={args.max_k} repeats={args.repeats}"
+    )
+    print("=" * 74)
+    print(f"{'N':>8} | {'KDTree(s)':>12} | {'DistMatrix(s)':>14} | {'Speedup':>7}")
+    print("-" * 74)
+
+    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))
+
+    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__":
+    raise SystemExit(main())

examples/lennard_jones_optimization.py

@@ -67,8 +67,6 @@ def calculate_potential(log_a, epsilon=0.5, sigma=1.0):
     history["a"].append(K.exp(log_a))
     history["energy"].append(energy)
 
-    # (Removed previously added blanket NaN guard per reviewer request to keep example minimal.)
-
     updates, opt_state = optimizer.update(grad, opt_state)
     log_a = optax.apply_updates(log_a, updates)
 
@@ -120,5 +118,3 @@ def calculate_potential(log_a, epsilon=0.5, sigma=1.0):
     plt.legend()
     plt.grid(True)
     plt.show()
-else:
-    print("\nOptimization failed. Final energy is NaN.")

tensorcircuit/backends/abstract_backend.py

@@ -631,8 +631,8 @@ def meshgrid(self: Any, *args: Any, **kwargs: Any) -> Any:
         """
          Return coordinate matrices from coordinate vectors.
 
-         :param args: coordinate vectors
-         :type args: Any
+        :param args: coordinate vectors
+        :type args: Any
         :param kwargs: keyword arguments for meshgrid, typically includes 'indexing'
            which can be 'ij' (matrix indexing) or 'xy' (Cartesian indexing).
            - 'ij': matrix indexing, first dimension corresponds to rows (default)
@@ -647,9 +647,9 @@ def meshgrid(self: Any, *args: Any, **kwargs: Any) -> Any:
                   [0, 1]]
               y = [[0, 0],
                   [2, 2]]
-         :type kwargs: Any
-         :return: list of coordinate matrices
-         :rtype: Any
+        :type kwargs: Any
+        :return: list of coordinate matrices
+        :rtype: Any
         """
         raise NotImplementedError(
             "Backend '{}' has not implemented `meshgrid`.".format(self.name)

tensorcircuit/backends/numpy_backend.py

@@ -137,9 +137,6 @@ def kron(self, a: Tensor, b: Tensor) -> Tensor:
         return np.kron(a, b)
 
     def meshgrid(self, *args: Any, **kwargs: Any) -> Any:
-        """
-        Backend-agnostic meshgrid function.
-        """
         return np.meshgrid(*args, **kwargs)
 
     def dtype(self, a: Tensor) -> str: