fix according to the review

Author: Stellogic

examples/vqe2d_lattice.py

@@ -0,0 +1,136 @@
+import time
+import optax
+import tensorcircuit as tc
+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
+        nq = circuit._nqubits
+        for i in range(0, nq - 1, 2):
+            j = (i + 1) % nq
+            circuit.X(i)
+            circuit.H(i)
+            circuit.cnot(i, j)
+            circuit.X(j)
+        return circuit
+
+    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.
+        """
+        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.
+        """
+        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()
+        param, opt_state, loss = train_step(param, opt_state)
+        time1 = time.time()
+        if i % 10 == 0:
+            print(
+                f"Step {i:4d}: Loss = {loss:.6f} \t (Time per step: {time1 - time0:.4f}s)"
+            )
+
+    print("Optimization finished.")
+    # Example result on A800 GPU: ~-25.3
+    print(f"Final Loss: {loss:.6f}")
+
+
+if __name__ == "__main__":
+    run_vqe()

tensorcircuit/templates/circuit_utils.py

@@ -15,6 +15,7 @@
     Union,
     TYPE_CHECKING,
     cast,
+    Set,
 )
 
 logger = logging.getLogger(__name__)
@@ -1446,3 +1447,54 @@ def remove_sites(self, identifiers: List[SiteIdentifier]) -> None:
         logger.info(
             f"{len(ids_to_remove)} sites removed. Lattice now has {self.num_sites} sites."
         )
+
+
+def get_compatible_layers(bonds: List[Tuple[int, int]]) -> List[List[Tuple[int, int]]]:
+    """
+    Partitions a list of pairs (bonds) into compatible layers for parallel
+    gate application using a greedy edge-coloring algorithm.
+
+    This function takes a list of pairs, representing connections like
+    nearest-neighbor (NN) or next-nearest-neighbor (NNN) bonds, and
+    partitions them into the minimum number of sets ("layers") where no two
+    pairs in a set share an index. This is a general utility for scheduling
+    non-overlapping operations.
+
+    :Example:
+
+    >>> from tensorcircuit.templates.lattice import SquareLattice
+    >>> sq_lattice = SquareLattice(size=(2, 2), pbc=False)
+    >>> nn_bonds = sq_lattice.get_neighbor_pairs(k=1, unique=True)
+
+    >>> gate_layers = get_compatible_layers(nn_bonds)
+    >>> print(gate_layers)
+    [[[0, 1], [2, 3]], [[0, 2], [1, 3]]]
+
+    :param bonds: A list of tuples, where each tuple represents a bond (i, j)
+        of site indices to be scheduled.
+    :type bonds: List[Tuple[int, int]]
+    :return: A list of layers. Each layer is a list of tuples, where each
+        tuple represents a bond. All bonds within a layer are non-overlapping.
+    :rtype: List[List[Tuple[int, int]]]
+    """
+    uncolored_edges: Set[Tuple[int, int]] = {(min(bond), max(bond)) for bond in bonds}
+
+    layers: List[List[Tuple[int, int]]] = []
+
+    while uncolored_edges:
+        current_layer: List[Tuple[int, int]] = []
+        qubits_in_this_layer: Set[int] = set()
+
+        edges_to_process = sorted(list(uncolored_edges))
+
+        for edge in edges_to_process:
+            i, j = edge
+            if i not in qubits_in_this_layer and j not in qubits_in_this_layer:
+                current_layer.append(edge)
+                qubits_in_this_layer.add(i)
+                qubits_in_this_layer.add(j)
+
+        uncolored_edges -= set(current_layer)
+        layers.append(sorted(current_layer))
+
+    return layers