fix according to the review 2

Author: Stellogic

examples/vqe2d_lattice.py

@@ -1,3 +1,9 @@
+"""
+This example demonstrates how to use the VQE algorithm to find the ground state
+of a 2D Heisenberg model on a square lattice. It showcases the setup of the lattice,
+the Heisenberg Hamiltonian, a suitable ansatz, and the optimization process.
+"""
+
 import time
 import optax
 import tensorcircuit as tc
@@ -7,16 +13,18 @@
 # Use JAX for high-performance, especially on GPU.
 K = tc.set_backend("jax")
 tc.set_dtype("complex64")
-# On Windows, cotengra's multiprocessing can cause issues.
-tc.set_contractor("cotengra-8192-8192", parallel=False)
+# On Windows, cotengra's multiprocessing can cause issues, use threads instead.
+tc.set_contractor("cotengra-8192-8192", parallel="threads")
 
 
 def run_vqe():
-    n, m, nlayers = 4, 4, 6
+    """Set up and run the VQE optimization for a 2D Heisenberg model."""
+    n, m, nlayers = 4, 4, 2
     lattice = SquareLattice(size=(n, m), pbc=True, precompute_neighbors=1)
     h = heisenberg_hamiltonian(lattice, j_coupling=[1.0, 1.0, 0.8])  # Jx, Jy, Jz
     nn_bonds = lattice.get_neighbor_pairs(k=1, unique=True)
     gate_layers = get_compatible_layers(nn_bonds)
+    n_params = nlayers * len(nn_bonds) * 3
 
     def singlet_init(circuit):
         # A good initial state for Heisenberg ground state search
@@ -36,22 +44,26 @@ def vqe_forward(param):
         """
         c = tc.Circuit(n * m)
         c = singlet_init(c)
+        param_idx = 0
 
-        for i in range(nlayers):
+        for _ in range(nlayers):
             for layer in gate_layers:
                 for j, k in layer:
-                    c.rzz(int(j), int(k), theta=param[i, 0])
+                    c.rzz(int(j), int(k), theta=param[param_idx])
+                    param_idx += 1
             for layer in gate_layers:
                 for j, k in layer:
-                    c.rxx(int(j), int(k), theta=param[i, 1])
+                    c.rxx(int(j), int(k), theta=param[param_idx])
+                    param_idx += 1
             for layer in gate_layers:
                 for j, k in layer:
-                    c.ryy(int(j), int(k), theta=param[i, 2])
+                    c.ryy(int(j), int(k), theta=param[param_idx])
+                    param_idx += 1
 
         return tc.templates.measurements.operator_expectation(c, h)
 
     vgf = K.jit(K.value_and_grad(vqe_forward))
-    param = tc.backend.implicit_randn(stddev=0.02, shape=[nlayers, 3])
+    param = tc.backend.implicit_randn(stddev=0.02, shape=[n_params])
     optimizer = optax.adam(learning_rate=3e-3)
     opt_state = optimizer.init(param)
 

tests/test_lattice.py

@@ -1,6 +1,5 @@
 from unittest.mock import patch
 import logging
-from typing import List, Set, Tuple
 
 # import time
 
@@ -1669,35 +1668,7 @@ def test_distance_matrix_invariants_for_all_lattice_types(self, lattice):
 #         )
 
 
-class MockLattice(AbstractLattice):
-    """A mock lattice class for testing purposes to precisely control neighbors."""
-
-    def __init__(self, neighbor_pairs: List[Tuple[int, int]]):
-        super().__init__(dimensionality=0)
-        # Ensure bonds are stored in a canonical sorted format for consistency
-        self._neighbor_pairs = [tuple(sorted(p)) for p in neighbor_pairs]
-
-    def get_neighbor_pairs(
-        self, k: int = 1, unique: bool = True
-    ) -> List[Tuple[int, int]]:
-        # The mock lattice only knows about k=1 neighbors
-        if k == 1:
-            return self._neighbor_pairs
-        return []
-
-    def _build_lattice(self, *args, **kwargs) -> None:
-        pass
-
-    def _build_neighbors(self, max_k: int = 1, **kwargs) -> None:
-        pass
-
-    def _compute_distance_matrix(self) -> np.ndarray:
-        return np.array([])
-
-
-def _validate_layers(
-    bonds: List[Tuple[int, int]], layers: List[List[Tuple[int, int]]]
-) -> None:
+def _validate_layers(bonds, layers) -> None:
     """
     A helper function to scientifically validate the output of get_compatible_layers.
     """
@@ -1711,7 +1682,7 @@ def _validate_layers(
     "exactly match the input bonds."
 
     for i, layer in enumerate(layers):
-        qubits_in_layer: Set[int] = set()
+        qubits_in_layer: set[int] = set()
         for edge in layer:
             q1, q2 = edge
             assert (