From 8c666fa63abbac08bb22a579186e3715dddb4d93 Mon Sep 17 00:00:00 2001 From: Stellogic Date: Tue, 29 Jul 2025 13:48:58 +0800 Subject: [PATCH 1/3] feat(templates): Add greedy algorithm for gate layering Introduces the `get_nn_gate_layers` utility function. This function takes a lattice object and partitions its nearest-neighbor pairs into the minimum number of compatible layers for parallel two-qubit gate application. This implementation uses a greedy edge-coloring algorithm to ensure that no two gates within the same layer act on the same qubit. The output is deterministic, with both layers and the edges within them being sorted. This functionality is essential for efficiently scheduling gates in algorithms like Trotterized Hamiltonian evolution and directly addresses Task 2 of the lattice API follow-up plan. --- tensorcircuit/templates/circuit_utils.py | 59 +++++++++++ tests/test_circuit_utils.py | 125 +++++++++++++++++++++++ 2 files changed, 184 insertions(+) create mode 100644 tensorcircuit/templates/circuit_utils.py create mode 100644 tests/test_circuit_utils.py diff --git a/tensorcircuit/templates/circuit_utils.py b/tensorcircuit/templates/circuit_utils.py new file mode 100644 index 00000000..9e5e04dd --- /dev/null +++ b/tensorcircuit/templates/circuit_utils.py @@ -0,0 +1,59 @@ +from typing import List, Set, Tuple + +from .lattice import AbstractLattice + + +def get_nn_gate_layers(lattice: AbstractLattice) -> List[List[Tuple[int, int]]]: + """ + Partitions nearest-neighbor pairs into compatible layers for parallel + gate application using a greedy edge-coloring algorithm. + + In quantum circuits, a single qubit cannot participate in more than one + two-qubit gate simultaneously. This function takes a lattice geometry, + finds its nearest-neighbor graph, and partitions the edges of that graph + (the neighbor pairs) into the minimum number of sets ("layers") where + no two edges in a set share a vertex. + + This is essential for efficiently scheduling gates in algorithms like + Trotterized Hamiltonian evolution. + + :Example: + + >>> import numpy as np + >>> from tensorcircuit.templates.lattice import SquareLattice + >>> sq_lattice = SquareLattice(size=(2, 2), pbc=False) + >>> gate_layers = get_nn_gate_layers(sq_lattice) + >>> print(gate_layers) + [[[0, 1], [2, 3]], [[0, 2], [1, 3]]] + + :param lattice: An initialized `AbstractLattice` object from which to + extract nearest-neighbor connectivity. + :type lattice: AbstractLattice + :return: A list of layers. Each layer is a list of tuples, where each + tuple represents a nearest-neighbor pair (i, j) of site indices. + All pairs within a layer are non-overlapping. + :rtype: List[List[Tuple[int, int]]] + """ + uncolored_edges: Set[Tuple[int, int]] = set( + lattice.get_neighbor_pairs(k=1, unique=True) + ) + + layers: List[List[Tuple[int, int]]] = [] + + while uncolored_edges: + current_layer: List[Tuple[int, int]] = [] + qubits_in_this_layer: Set[int] = set() + edges_to_remove: Set[Tuple[int, int]] = set() + + for edge in sorted(list(uncolored_edges)): + 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) + edges_to_remove.add(edge) + + layers.append(sorted(current_layer)) + uncolored_edges -= edges_to_remove + + return layers diff --git a/tests/test_circuit_utils.py b/tests/test_circuit_utils.py new file mode 100644 index 00000000..10c738b0 --- /dev/null +++ b/tests/test_circuit_utils.py @@ -0,0 +1,125 @@ +from typing import List, Set, Tuple + +import pytest +import numpy as np + +from tensorcircuit.templates.circuit_utils import get_nn_gate_layers +from tensorcircuit.templates.lattice import ( + AbstractLattice, + ChainLattice, + HoneycombLattice, + SquareLattice, +) + + +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) + self._neighbor_pairs = neighbor_pairs + + def get_neighbor_pairs( + self, k: int = 1, unique: bool = True + ) -> List[Tuple[int, int]]: + return self._neighbor_pairs + + 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( + lattice: AbstractLattice, layers: List[List[Tuple[int, int]]] +) -> None: + """ + A helper function to scientifically validate the output of get_nn_gate_layers. + """ + expected_edges = set(lattice.get_neighbor_pairs(k=1, unique=True)) + actual_edges = set(tuple(sorted(edge)) for layer in layers for edge in layer) + + assert ( + expected_edges == actual_edges + ), "Completeness check failed: The set of all edges in the layers must " + "exactly match the lattice's unique nearest-neighbor pairs." + + for i, layer in enumerate(layers): + qubits_in_layer: Set[int] = set() + for edge in layer: + q1, q2 = edge + assert ( + q1 not in qubits_in_layer + ), f"Compatibility check failed: Qubit {q1} is reused in layer {i}." + qubits_in_layer.add(q1) + assert ( + q2 not in qubits_in_layer + ), f"Compatibility check failed: Qubit {q2} is reused in layer {i}." + qubits_in_layer.add(q2) + + +@pytest.mark.parametrize( + "lattice_instance", + [ + SquareLattice(size=(3, 2), pbc=False), + SquareLattice(size=(3, 3), pbc=True), + HoneycombLattice(size=(2, 2), pbc=False), + ], + ids=[ + "SquareLattice_3x2_OBC", + "SquareLattice_3x3_PBC", + "HoneycombLattice_2x2_OBC", + ], +) +def test_various_lattices_layering(lattice_instance: AbstractLattice): + """Tests gate layering for various standard lattice types.""" + layers = get_nn_gate_layers(lattice_instance) + assert len(layers) > 0, "Layers should not be empty for non-trivial lattices." + _validate_layers(lattice_instance, layers) + + +def test_1d_chain_pbc(): + """Test layering on a 1D chain with periodic boundaries (a cycle graph).""" + lattice_even = ChainLattice(size=(6,), pbc=True) + layers_even = get_nn_gate_layers(lattice_even) + _validate_layers(lattice_even, layers_even) + + lattice_odd = ChainLattice(size=(5,), pbc=True) + layers_odd = get_nn_gate_layers(lattice_odd) + assert len(layers_odd) == 3, "A 5-site cycle graph should be 3-colorable." + _validate_layers(lattice_odd, layers_odd) + + +def test_custom_star_graph(): + """Test layering on a custom lattice forming a star graph.""" + star_edges = [(0, 1), (0, 2), (0, 3)] + lattice = MockLattice(star_edges) + layers = get_nn_gate_layers(lattice) + assert len(layers) == 3, "A star graph S_4 requires 3 layers." + _validate_layers(lattice, layers) + + +def test_edge_cases(): + """Test various edge cases: empty, single-site, and no-edge lattices.""" + empty_lattice = MockLattice([]) + layers = get_nn_gate_layers(empty_lattice) + assert layers == [], "Layers should be empty for an empty lattice." + + single_site_lattice = MockLattice([]) + layers = get_nn_gate_layers(single_site_lattice) + assert layers == [], "Layers should be empty for a single-site lattice." + + disconnected_lattice = MockLattice([]) + layers = get_nn_gate_layers(disconnected_lattice) + assert layers == [], "Layers should be empty for a lattice with no neighbors." + + single_edge_lattice = MockLattice([(0, 1)]) + layers = get_nn_gate_layers(single_edge_lattice) + # The tuple inside the list might be (0, 1) or (1, 0) after sorting. + # We check for the sorted version to be deterministic. + assert layers == [[(0, 1)]] + _validate_layers(single_edge_lattice, layers) From ebdebc95bd86023c5f3674e5c47d0ae434daee01 Mon Sep 17 00:00:00 2001 From: Stellogic Date: Wed, 30 Jul 2025 15:39:59 +0800 Subject: [PATCH 2/3] fix according to the review --- examples/vqe2d_lattice.py | 81 +++++++++++++++ tensorcircuit/templates/circuit_utils.py | 59 ----------- tensorcircuit/templates/lattice.py | 52 ++++++++++ tests/test_circuit_utils.py | 125 ----------------------- tests/test_lattice.py | 113 ++++++++++++++++++++ 5 files changed, 246 insertions(+), 184 deletions(-) create mode 100644 examples/vqe2d_lattice.py delete mode 100644 tensorcircuit/templates/circuit_utils.py delete mode 100644 tests/test_circuit_utils.py diff --git a/examples/vqe2d_lattice.py b/examples/vqe2d_lattice.py new file mode 100644 index 00000000..9a935e66 --- /dev/null +++ b/examples/vqe2d_lattice.py @@ -0,0 +1,81 @@ +import time +import optax +import tensorcircuit as tc +from tensorcircuit.templates.lattice import SquareLattice, get_compatible_layers +from tensorcircuit.templates.hamiltonians import heisenberg_hamiltonian + +# 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) + + +def run_vqe(): + n, m, nlayers = 4, 4, 6 + 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) + + 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 consists of nlayers of RZZ, RXX, and RYY entangling layers. + """ + c = tc.Circuit(n * m) + c = singlet_init(c) + + for i in range(nlayers): + for layer in gate_layers: + for j, k in layer: + c.rzz(int(j), int(k), theta=param[i, 0]) + for layer in gate_layers: + for j, k in layer: + c.rxx(int(j), int(k), theta=param[i, 1]) + for layer in gate_layers: + for j, k in layer: + c.ryy(int(j), int(k), theta=param[i, 2]) + + 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]) + 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.""" + 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 + + 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.") + print(f"Final Loss: {loss:.6f}") + + +if __name__ == "__main__": + run_vqe() diff --git a/tensorcircuit/templates/circuit_utils.py b/tensorcircuit/templates/circuit_utils.py deleted file mode 100644 index 9e5e04dd..00000000 --- a/tensorcircuit/templates/circuit_utils.py +++ /dev/null @@ -1,59 +0,0 @@ -from typing import List, Set, Tuple - -from .lattice import AbstractLattice - - -def get_nn_gate_layers(lattice: AbstractLattice) -> List[List[Tuple[int, int]]]: - """ - Partitions nearest-neighbor pairs into compatible layers for parallel - gate application using a greedy edge-coloring algorithm. - - In quantum circuits, a single qubit cannot participate in more than one - two-qubit gate simultaneously. This function takes a lattice geometry, - finds its nearest-neighbor graph, and partitions the edges of that graph - (the neighbor pairs) into the minimum number of sets ("layers") where - no two edges in a set share a vertex. - - This is essential for efficiently scheduling gates in algorithms like - Trotterized Hamiltonian evolution. - - :Example: - - >>> import numpy as np - >>> from tensorcircuit.templates.lattice import SquareLattice - >>> sq_lattice = SquareLattice(size=(2, 2), pbc=False) - >>> gate_layers = get_nn_gate_layers(sq_lattice) - >>> print(gate_layers) - [[[0, 1], [2, 3]], [[0, 2], [1, 3]]] - - :param lattice: An initialized `AbstractLattice` object from which to - extract nearest-neighbor connectivity. - :type lattice: AbstractLattice - :return: A list of layers. Each layer is a list of tuples, where each - tuple represents a nearest-neighbor pair (i, j) of site indices. - All pairs within a layer are non-overlapping. - :rtype: List[List[Tuple[int, int]]] - """ - uncolored_edges: Set[Tuple[int, int]] = set( - lattice.get_neighbor_pairs(k=1, unique=True) - ) - - layers: List[List[Tuple[int, int]]] = [] - - while uncolored_edges: - current_layer: List[Tuple[int, int]] = [] - qubits_in_this_layer: Set[int] = set() - edges_to_remove: Set[Tuple[int, int]] = set() - - for edge in sorted(list(uncolored_edges)): - 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) - edges_to_remove.add(edge) - - layers.append(sorted(current_layer)) - uncolored_edges -= edges_to_remove - - return layers diff --git a/tensorcircuit/templates/lattice.py b/tensorcircuit/templates/lattice.py index d4e3e541..52f152c9 100644 --- a/tensorcircuit/templates/lattice.py +++ b/tensorcircuit/templates/lattice.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 diff --git a/tests/test_circuit_utils.py b/tests/test_circuit_utils.py deleted file mode 100644 index 10c738b0..00000000 --- a/tests/test_circuit_utils.py +++ /dev/null @@ -1,125 +0,0 @@ -from typing import List, Set, Tuple - -import pytest -import numpy as np - -from tensorcircuit.templates.circuit_utils import get_nn_gate_layers -from tensorcircuit.templates.lattice import ( - AbstractLattice, - ChainLattice, - HoneycombLattice, - SquareLattice, -) - - -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) - self._neighbor_pairs = neighbor_pairs - - def get_neighbor_pairs( - self, k: int = 1, unique: bool = True - ) -> List[Tuple[int, int]]: - return self._neighbor_pairs - - 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( - lattice: AbstractLattice, layers: List[List[Tuple[int, int]]] -) -> None: - """ - A helper function to scientifically validate the output of get_nn_gate_layers. - """ - expected_edges = set(lattice.get_neighbor_pairs(k=1, unique=True)) - actual_edges = set(tuple(sorted(edge)) for layer in layers for edge in layer) - - assert ( - expected_edges == actual_edges - ), "Completeness check failed: The set of all edges in the layers must " - "exactly match the lattice's unique nearest-neighbor pairs." - - for i, layer in enumerate(layers): - qubits_in_layer: Set[int] = set() - for edge in layer: - q1, q2 = edge - assert ( - q1 not in qubits_in_layer - ), f"Compatibility check failed: Qubit {q1} is reused in layer {i}." - qubits_in_layer.add(q1) - assert ( - q2 not in qubits_in_layer - ), f"Compatibility check failed: Qubit {q2} is reused in layer {i}." - qubits_in_layer.add(q2) - - -@pytest.mark.parametrize( - "lattice_instance", - [ - SquareLattice(size=(3, 2), pbc=False), - SquareLattice(size=(3, 3), pbc=True), - HoneycombLattice(size=(2, 2), pbc=False), - ], - ids=[ - "SquareLattice_3x2_OBC", - "SquareLattice_3x3_PBC", - "HoneycombLattice_2x2_OBC", - ], -) -def test_various_lattices_layering(lattice_instance: AbstractLattice): - """Tests gate layering for various standard lattice types.""" - layers = get_nn_gate_layers(lattice_instance) - assert len(layers) > 0, "Layers should not be empty for non-trivial lattices." - _validate_layers(lattice_instance, layers) - - -def test_1d_chain_pbc(): - """Test layering on a 1D chain with periodic boundaries (a cycle graph).""" - lattice_even = ChainLattice(size=(6,), pbc=True) - layers_even = get_nn_gate_layers(lattice_even) - _validate_layers(lattice_even, layers_even) - - lattice_odd = ChainLattice(size=(5,), pbc=True) - layers_odd = get_nn_gate_layers(lattice_odd) - assert len(layers_odd) == 3, "A 5-site cycle graph should be 3-colorable." - _validate_layers(lattice_odd, layers_odd) - - -def test_custom_star_graph(): - """Test layering on a custom lattice forming a star graph.""" - star_edges = [(0, 1), (0, 2), (0, 3)] - lattice = MockLattice(star_edges) - layers = get_nn_gate_layers(lattice) - assert len(layers) == 3, "A star graph S_4 requires 3 layers." - _validate_layers(lattice, layers) - - -def test_edge_cases(): - """Test various edge cases: empty, single-site, and no-edge lattices.""" - empty_lattice = MockLattice([]) - layers = get_nn_gate_layers(empty_lattice) - assert layers == [], "Layers should be empty for an empty lattice." - - single_site_lattice = MockLattice([]) - layers = get_nn_gate_layers(single_site_lattice) - assert layers == [], "Layers should be empty for a single-site lattice." - - disconnected_lattice = MockLattice([]) - layers = get_nn_gate_layers(disconnected_lattice) - assert layers == [], "Layers should be empty for a lattice with no neighbors." - - single_edge_lattice = MockLattice([(0, 1)]) - layers = get_nn_gate_layers(single_edge_lattice) - # The tuple inside the list might be (0, 1) or (1, 0) after sorting. - # We check for the sorted version to be deterministic. - assert layers == [[(0, 1)]] - _validate_layers(single_edge_lattice, layers) diff --git a/tests/test_lattice.py b/tests/test_lattice.py index d332e6cd..5c03b7bb 100644 --- a/tests/test_lattice.py +++ b/tests/test_lattice.py @@ -1,5 +1,6 @@ from unittest.mock import patch import logging +from typing import List, Set, Tuple # import time @@ -23,6 +24,8 @@ RectangularLattice, SquareLattice, TriangularLattice, + AbstractLattice, + get_compatible_layers, ) @@ -1664,3 +1667,113 @@ def test_distance_matrix_invariants_for_all_lattice_types(self, lattice): # "The specialized PBC implementation is significantly slower " # "than the general-purpose implementation." # ) + + +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: + """ + A helper function to scientifically validate the output of get_compatible_layers. + """ + # MODIFICATION: This function now takes the original bonds list for comparison. + expected_edges = set(tuple(sorted(b)) for b in bonds) + actual_edges = set(tuple(sorted(edge)) for layer in layers for edge in layer) + + assert ( + expected_edges == actual_edges + ), "Completeness check failed: The set of all edges in the layers must " + "exactly match the input bonds." + + for i, layer in enumerate(layers): + qubits_in_layer: Set[int] = set() + for edge in layer: + q1, q2 = edge + assert ( + q1 not in qubits_in_layer + ), f"Compatibility check failed: Qubit {q1} is reused in layer {i}." + qubits_in_layer.add(q1) + assert ( + q2 not in qubits_in_layer + ), f"Compatibility check failed: Qubit {q2} is reused in layer {i}." + qubits_in_layer.add(q2) + + +@pytest.mark.parametrize( + "lattice_instance", + [ + SquareLattice(size=(3, 2), pbc=False), + SquareLattice(size=(3, 3), pbc=True), + HoneycombLattice(size=(2, 2), pbc=False), + ], + ids=[ + "SquareLattice_3x2_OBC", + "SquareLattice_3x3_PBC", + "HoneycombLattice_2x2_OBC", + ], +) +def test_layering_on_various_lattices(lattice_instance: AbstractLattice): + """Tests gate layering for various standard lattice types.""" + bonds = lattice_instance.get_neighbor_pairs(k=1, unique=True) + layers = get_compatible_layers(bonds) + + assert len(layers) > 0, "Layers should not be empty for non-trivial lattices." + _validate_layers(bonds, layers) + + +def test_layering_on_1d_chain_pbc(): + """Test layering on a 1D chain with periodic boundaries (a cycle graph).""" + lattice_even = ChainLattice(size=(6,), pbc=True) + bonds_even = lattice_even.get_neighbor_pairs(k=1, unique=True) + layers_even = get_compatible_layers(bonds_even) + _validate_layers(bonds_even, layers_even) + + lattice_odd = ChainLattice(size=(5,), pbc=True) + bonds_odd = lattice_odd.get_neighbor_pairs(k=1, unique=True) + layers_odd = get_compatible_layers(bonds_odd) + assert len(layers_odd) == 3, "A 5-site cycle graph should be 3-colorable." + _validate_layers(bonds_odd, layers_odd) + + +def test_layering_on_custom_star_graph(): + """Test layering on a custom lattice forming a star graph.""" + star_edges = [(0, 1), (0, 2), (0, 3)] + layers = get_compatible_layers(star_edges) + assert len(layers) == 3, "A star graph S_4 requires 3 layers." + _validate_layers(star_edges, layers) + + +def test_layering_on_edge_cases(): + """Test various edge cases: empty, single-site, and no-edge lattices.""" + layers_empty = get_compatible_layers([]) + assert layers_empty == [], "Layers should be empty for an empty set of bonds." + + single_edge = [(0, 1)] + layers_single = get_compatible_layers(single_edge) + assert layers_single == [[(0, 1)]] + _validate_layers(single_edge, layers_single) From 947d35799b9159a08cb32590a319f8466bcd0d65 Mon Sep 17 00:00:00 2001 From: Stellogic Date: Fri, 1 Aug 2025 16:01:45 +0800 Subject: [PATCH 3/3] fix according to the review 2 --- examples/vqe2d_lattice.py | 28 ++++++++++++++++++++-------- tests/test_lattice.py | 33 ++------------------------------- 2 files changed, 22 insertions(+), 39 deletions(-) diff --git a/examples/vqe2d_lattice.py b/examples/vqe2d_lattice.py index 9a935e66..13d78876 100644 --- a/examples/vqe2d_lattice.py +++ b/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) diff --git a/tests/test_lattice.py b/tests/test_lattice.py index 5c03b7bb..12354f13 100644 --- a/tests/test_lattice.py +++ b/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 (