fix according to the review

Author: Stellogic

examples/vqe2d_lattice.py

@@ -4,48 +4,22 @@
 from tensorcircuit.templates.lattice import SquareLattice, get_compatible_layers
 from tensorcircuit.templates.hamiltonians import heisenberg_hamiltonian
 
-# ===================
-# Backend and Hardware Configuration
-# ===================
 # Use JAX for high-performance, especially on GPU.
-# For CPU-only environments, TensorFlow can also be efficient.
 K = tc.set_backend("jax")
 tc.set_dtype("complex64")
-# Use a more powerful contractor for better performance on larger graphs.
 # On Windows, cotengra's multiprocessing can cause issues.
-# We disable it here to ensure stability.
 tc.set_contractor("cotengra-8192-8192", parallel=False)
 
 
 def run_vqe():
-    # ===================
-    # Lattice and Hamiltonian Definition
-    # ===================
-    # Define the 2D lattice dimensions and the number of VQE layers.
     n, m, nlayers = 4, 4, 6
-
-    # 1. Create a SquareLattice instance.
-    # This object holds all geometric information, such as site coordinates and neighbors.
     lattice = SquareLattice(size=(n, m), pbc=True, precompute_neighbors=1)
-
-    # 2. Generate the Heisenberg Hamiltonian using the new interface.
-    # This function directly takes the lattice object and coupling constants.
-    # It's cleaner than the old method that required a separate graph object.
     h = heisenberg_hamiltonian(lattice, j_coupling=[1.0, 1.0, 0.8])  # Jx, Jy, Jz
-
-    # 3. Get nearest-neighbor bonds and partition them into compatible layers.
-    # This is the core of the gate scheduling logic. `get_compatible_layers`
-    # ensures that gates within each layer can be applied in parallel without overlap.
     nn_bonds = lattice.get_neighbor_pairs(k=1, unique=True)
     gate_layers = get_compatible_layers(nn_bonds)
 
-    # ===================
-    # VQE Ansatz and Forward Pass
-    # ===================
-
-    def singlet_init(
-        circuit,
-    ):  # A good initial state for Heisenberg ground state search
+    def singlet_init(circuit):
+        # A good initial state for Heisenberg ground state search
         nq = circuit._nqubits
         for i in range(0, nq - 1, 2):
             j = (i + 1) % nq
@@ -58,65 +32,37 @@ def singlet_init(
     def vqe_forward(param):
         """
         Defines the VQE ansatz and computes the energy expectation.
-
-        The ansatz structure is:
-        - Initial state preparation (singlet pairs).
-        - nlayers of parameterized blocks.
-        - Each block consists of RZZ, RXX, and RYY entangling layers.
-        - Gates within each entangling layer are applied according to the pre-computed
-          `gate_layers` for maximum parallelism. All gates of the same type in a
-          VQE layer share the same parameter.
+        The ansatz consists of nlayers of RZZ, RXX, and RYY entangling layers.
         """
         c = tc.Circuit(n * m)
         c = singlet_init(c)
 
         for i in range(nlayers):
-            # RZZ layer
             for layer in gate_layers:
                 for j, k in layer:
                     c.rzz(int(j), int(k), theta=param[i, 0])
-
-            # RXX layer
             for layer in gate_layers:
                 for j, k in layer:
                     c.rxx(int(j), int(k), theta=param[i, 1])
-
-            # RYY layer
             for layer in gate_layers:
                 for j, k in layer:
                     c.ryy(int(j), int(k), theta=param[i, 2])
 
-        # The Hamiltonian is a sparse matrix, so we use the corresponding expectation method.
         return tc.templates.measurements.operator_expectation(c, h)
 
-    # ===================
-    # Training and Optimization (JAX-based for performance)
-    # ===================
-    # Value_and_grad for single (non-batched) training instance.
     vgf = K.jit(K.value_and_grad(vqe_forward))
-
-    # Parameters for a single training instance.
-    # Shape: (nlayers, 3) -> 3 for RZZ, RXX, RYY angles per layer.
     param = tc.backend.implicit_randn(stddev=0.02, shape=[nlayers, 3])
-
-    # Use the Adam optimizer from Optax.
     optimizer = optax.adam(learning_rate=3e-3)
     opt_state = optimizer.init(param)
 
     @K.jit
     def train_step(param, opt_state):
-        """
-        A single training step, JIT-compiled for maximum speed.
-        This follows the standard Optax optimization paradigm.
-        """
+        """A single training step, JIT-compiled for maximum speed."""
         loss_val, grads = vgf(param)
         updates, opt_state = optimizer.update(grads, opt_state, param)
         param = optax.apply_updates(param, updates)
         return param, opt_state, loss_val
 
-    # ===================
-    # Main Training Loop
-    # ===================
     print("Starting VQE optimization...")
     for i in range(1000):
         time0 = time.time()
@@ -128,7 +74,6 @@ def train_step(param, opt_state):
             )
 
     print("Optimization finished.")
-    # Example result on A800 GPU: ~-25.3
     print(f"Final Loss: {loss:.6f}")