Merge branch 'feature/hamiltonian-generation' of github.com:Stellogic/tensorcircuit-ng into ngpr22

Author: refraction-ray

tensorcircuit/templates/hamiltonians.py

@@ -0,0 +1,153 @@
+from typing import Any, List, Tuple, Union
+import numpy as np
+from ..cons import dtypestr, backend
+from ..quantum import PauliStringSum2COO
+from .lattice import AbstractLattice
+
+
+def _create_empty_sparse_matrix(shape: Tuple[int, int]) -> Any:
+    """
+    Helper function to create a backend-agnostic empty sparse matrix.
+    """
+    indices = backend.convert_to_tensor(backend.zeros((0, 2), dtype="int32"))
+    values = backend.convert_to_tensor(backend.zeros((0,), dtype=dtypestr))  # type: ignore
+    return backend.coo_sparse_matrix(indices=indices, values=values, shape=shape)  # type: ignore
+
+
+def heisenberg_hamiltonian(
+    lattice: AbstractLattice,
+    j_coupling: Union[float, List[float], Tuple[float, ...]] = 1.0,
+) -> Any:
+    """
+    Generates the sparse matrix of the Heisenberg Hamiltonian for a given lattice.
+
+    The Heisenberg Hamiltonian is defined as:
+    H = J * Σ_{<i,j>} (X_i X_j + Y_i Y_j + Z_i Z_j)
+    where the sum is over all unique nearest-neighbor pairs <i,j>.
+
+    :param lattice: An instance of a class derived from AbstractLattice,
+        which provides the geometric information of the system.
+    :type lattice: AbstractLattice
+    :param j_coupling: The coupling constants. Can be a single float for an
+        isotropic model (Jx=Jy=Jz) or a list/tuple of 3 floats for an
+        anisotropic model (Jx, Jy, Jz). Defaults to 1.0.
+    :type j_coupling: Union[float, List[float], Tuple[float, ...]], optional
+    :return: The Hamiltonian as a backend-agnostic sparse matrix.
+    :rtype: Any
+    """
+    num_sites = lattice.num_sites
+    neighbor_pairs = lattice.get_neighbor_pairs(k=1, unique=True)
+
+    if isinstance(j_coupling, (float, int)):
+        js = [float(j_coupling)] * 3
+    else:
+        if len(j_coupling) != 3:
+            raise ValueError("j_coupling must be a float or a list/tuple of 3 floats.")
+        js = [float(j) for j in j_coupling]
+
+    if not neighbor_pairs:
+        return _create_empty_sparse_matrix(shape=(2**num_sites, 2**num_sites))
+    if num_sites == 0:
+        raise ValueError("Cannot generate a Hamiltonian for a lattice with zero sites.")
+
+    pauli_map = {"X": 1, "Y": 2, "Z": 3}
+
+    ls: List[List[int]] = []
+    weights: List[float] = []
+
+    pauli_terms = ["X", "Y", "Z"]
+    for i, j in neighbor_pairs:
+        for idx, pauli_char in enumerate(pauli_terms):
+            if abs(js[idx]) > 1e-9:
+                string = [0] * num_sites
+                string[i] = pauli_map[pauli_char]
+                string[j] = pauli_map[pauli_char]
+                ls.append(string)
+                weights.append(js[idx])
+
+    hamiltonian_matrix = PauliStringSum2COO(ls, weight=weights, numpy=False)
+
+    return hamiltonian_matrix
+
+
+def rydberg_hamiltonian(
+    lattice: AbstractLattice, omega: float, delta: float, c6: float
+) -> Any:
+    """
+    Generates the sparse matrix of the Rydberg atom array Hamiltonian.
+
+    The Hamiltonian is defined as:
+    H = Σ_i (Ω/2)X_i - Σ_i δ(1 - Z_i)/2 + Σ_{i<j} V_ij (1-Z_i)/2 (1-Z_j)/2
+      = Σ_i (Ω/2)X_i + Σ_i (δ/2)Z_i + Σ_{i<j} (V_ij/4)(Z_iZ_j - Z_i - Z_j)
+    where V_ij = C6 / |r_i - r_j|^6.
+
+    Note: Constant energy offset terms (proportional to the identity operator)
+    are ignored in this implementation.
+
+    :param lattice: An instance of a class derived from AbstractLattice,
+        which provides site coordinates and the distance matrix.
+    :type lattice: AbstractLattice
+    :param omega: The Rabi frequency (Ω) of the driving laser field.
+    :type omega: float
+    :param delta: The laser detuning (δ).
+    :type delta: float
+    :param c6: The Van der Waals interaction coefficient (C6).
+    :type c6: float
+    :return: The Hamiltonian as a backend-agnostic sparse matrix.
+    :rtype: Any
+    """
+    num_sites = lattice.num_sites
+    if num_sites == 0:
+        raise ValueError("Cannot generate a Hamiltonian for a lattice with zero sites.")
+
+    pauli_map = {"X": 1, "Y": 2, "Z": 3}
+    ls: List[List[int]] = []
+    weights: List[float] = []
+
+    for i in range(num_sites):
+        x_string = [0] * num_sites
+        x_string[i] = pauli_map["X"]
+        ls.append(x_string)
+        weights.append(omega / 2.0)
+
+    z_coefficients = np.zeros(num_sites)
+
+    for i in range(num_sites):
+        z_coefficients[i] += delta / 2.0
+
+    dist_matrix = lattice.distance_matrix
+
+    for i in range(num_sites):
+        for j in range(i + 1, num_sites):
+            distance = dist_matrix[i, j]
+
+            if distance < 1e-9:
+                continue
+
+            interaction_strength = c6 / (distance**6)
+            coefficient = interaction_strength / 4.0
+
+            zz_string = [0] * num_sites
+            zz_string[i] = pauli_map["Z"]
+            zz_string[j] = pauli_map["Z"]
+            ls.append(zz_string)
+            weights.append(coefficient)
+
+            # The interaction term V_ij * n_i * n_j, when expanded using
+            # n_i = (1-Z_i)/2, becomes (V_ij/4)*(I - Z_i - Z_j + Z_i*Z_j).
+            # This contributes a positive term (+V_ij/4) to the ZZ interaction,
+            # but negative terms (-V_ij/4) to the single-site Z_i and Z_j operators.
+
+            z_coefficients[i] -= coefficient
+            z_coefficients[j] -= coefficient
+
+    for i in range(num_sites):
+        if abs(z_coefficients[i]) > 1e-9:
+            z_string = [0] * num_sites
+            z_string[i] = pauli_map["Z"]
+            ls.append(z_string)
+            weights.append(z_coefficients[i])  # type: ignore
+
+    hamiltonian_matrix = PauliStringSum2COO(ls, weight=weights, numpy=False)
+
+    return hamiltonian_matrix

tests/test_hamiltonians.py

@@ -0,0 +1,155 @@
+import pytest
+import numpy as np
+import tensorcircuit as tc
+
+
+from tensorcircuit.templates.lattice import (
+    ChainLattice,
+    SquareLattice,
+    CustomizeLattice,
+)
+from tensorcircuit.templates.hamiltonians import (
+    heisenberg_hamiltonian,
+    rydberg_hamiltonian,
+)
+
+PAULI_X = np.array([[0, 1], [1, 0]], dtype=complex)
+PAULI_Y = np.array([[0, -1j], [1j, 0]], dtype=complex)
+PAULI_Z = np.array([[1, 0], [0, -1]], dtype=complex)
+PAULI_I = np.eye(2, dtype=complex)
+
+
+class TestHeisenbergHamiltonian:
+    """
+    Test suite for the heisenberg_hamiltonian function.
+    """
+
+    def test_empty_lattice(self):
+        """
+        Test that an empty lattice produces a 0x0 matrix.
+        """
+        empty_lattice = CustomizeLattice(
+            dimensionality=2, identifiers=[], coordinates=[]
+        )
+        h = heisenberg_hamiltonian(empty_lattice)
+        assert h.shape == (1, 1)
+        assert h.nnz == 0
+
+    def test_single_site(self):
+        """
+        Test that a single-site lattice (no bonds) produces a 2x2 zero matrix.
+        """
+        single_site_lattice = ChainLattice(size=(1,), pbc=False)
+        h = heisenberg_hamiltonian(single_site_lattice)
+        assert h.shape == (2, 2)
+        assert h.nnz == 0
+
+    def test_two_sites_chain(self):
+        """
+        Test a two-site chain against a manually calculated Hamiltonian.
+        This is the most critical test for scientific correctness.
+        """
+        lattice = ChainLattice(size=(2,), pbc=False)
+        j_coupling = -1.5  # Test with a non-trivial coupling constant
+        h_generated = heisenberg_hamiltonian(lattice, j_coupling=j_coupling)
+
+        # Manually construct the expected Hamiltonian: H = J * (X_0X_1 + Y_0Y_1 + Z_0Z_1)
+        xx = np.kron(PAULI_X, PAULI_X)
+        yy = np.kron(PAULI_Y, PAULI_Y)
+        zz = np.kron(PAULI_Z, PAULI_Z)
+        h_expected = j_coupling * (xx + yy + zz)
+
+        assert h_generated.shape == (4, 4)
+        assert np.allclose(tc.backend.to_dense(h_generated), h_expected)
+
+    def test_square_lattice_properties(self):
+        """
+        Test properties of a larger lattice (2x2 square) without full matrix comparison.
+        """
+        lattice = SquareLattice(size=(2, 2), pbc=True)  # 4 sites, 8 bonds with PBC
+        h = heisenberg_hamiltonian(lattice, j_coupling=1.0)
+
+        assert h.shape == (16, 16)
+        assert h.nnz > 0
+        h_dense = tc.backend.to_dense(h)
+        assert np.allclose(h_dense, h_dense.conj().T)
+
+
+class TestRydbergHamiltonian:
+    """
+    Test suite for the rydberg_hamiltonian function.
+    """
+
+    def test_single_site_rydberg(self):
+        """
+        Test a single atom, which should only have driving and detuning terms.
+        """
+        lattice = ChainLattice(size=(1,), pbc=False)
+        omega, delta, c6 = 2.0, 0.5, 100.0
+        h_generated = rydberg_hamiltonian(lattice, omega, delta, c6)
+
+        h_expected = (omega / 2.0) * PAULI_X + (delta / 2.0) * PAULI_Z
+
+        assert h_generated.shape == (2, 2)
+        assert np.allclose(tc.backend.to_dense(h_generated), h_expected)
+
+    def test_two_sites_rydberg(self):
+        """
+        Test a two-site chain for Rydberg Hamiltonian, including interaction.
+        """
+        lattice = ChainLattice(size=(2,), pbc=False, lattice_constant=1.5)
+        omega, delta, c6 = 1.0, -0.5, 10.0
+        h_generated = rydberg_hamiltonian(lattice, omega, delta, c6)
+
+        v_ij = c6 / (1.5**6)
+
+        h1 = (omega / 2.0) * (np.kron(PAULI_X, PAULI_I) + np.kron(PAULI_I, PAULI_X))
+        z0_coeff = delta / 2.0 - v_ij / 4.0
+        z1_coeff = delta / 2.0 - v_ij / 4.0
+        h2 = z0_coeff * np.kron(PAULI_Z, PAULI_I) + z1_coeff * np.kron(PAULI_I, PAULI_Z)
+        h3 = (v_ij / 4.0) * np.kron(PAULI_Z, PAULI_Z)
+
+        h_expected = h1 + h2 + h3
+
+        assert h_generated.shape == (4, 4)
+        h_generated_dense = tc.backend.to_dense(h_generated)
+
+        assert np.allclose(h_generated_dense, h_expected)
+
+    def test_zero_distance_robustness(self):
+        """
+        Test that the function does not crash when two atoms have zero distance.
+        """
+        lattice = CustomizeLattice(
+            dimensionality=2,
+            identifiers=[0, 1],
+            coordinates=[[0.0, 0.0], [0.0, 0.0]],
+        )
+
+        try:
+            h = rydberg_hamiltonian(lattice, omega=1.0, delta=1.0, c6=1.0)
+            # The X terms contribute 8 non-zero elements.
+            # The Z terms (Z0+Z1) have diagonal elements that cancel out,
+            # resulting in only 2 non-zero elements. Total nnz = 8 + 2 = 10.
+            assert h.nnz == 10
+        except ZeroDivisionError:
+            pytest.fail("The function failed to handle zero distance between sites.")
+
+    def test_anisotropic_heisenberg(self):
+        """
+        Test the anisotropic Heisenberg model with different Jx, Jy, Jz.
+        """
+        lattice = ChainLattice(size=(2,), pbc=False)
+        j_coupling = [-1.0, 0.5, 2.0]  # Jx, Jy, Jz
+        h_generated = heisenberg_hamiltonian(lattice, j_coupling=j_coupling)
+
+        # Manually construct the expected Hamiltonian
+        jx, jy, jz = j_coupling
+        xx = np.kron(PAULI_X, PAULI_X)
+        yy = np.kron(PAULI_Y, PAULI_Y)
+        zz = np.kron(PAULI_Z, PAULI_Z)
+        h_expected = jx * xx + jy * yy + jz * zz
+
+        h_generated_dense = tc.backend.to_dense(h_generated)
+        assert h_generated_dense.shape == (4, 4)
+        assert np.allclose(h_generated_dense, h_expected)