Feature/hamiltonian generation

tensorcircuit/templates/hamiltonians.py

@@ -0,0 +1,139 @@
+import typing
+from typing import cast
+from scipy.sparse import coo_matrix
+import numpy as np
+import tensorcircuit as tc
+from .lattice import AbstractLattice
+
+
+def generate_heisenberg_hamiltonian(
+    lattice: AbstractLattice, j_coupling: float = 1.0
+) -> coo_matrix:
+    """
+    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 constant for the Heisenberg interaction. Defaults to 1.0.
+    :type j_coupling: float, optional
+    :return: The Hamiltonian represented as a SciPy COO sparse matrix.
+    :rtype: coo_matrix
+    """
+    num_sites = lattice.num_sites
+    if num_sites == 0:
+        return coo_matrix((0, 0))
+
+    neighbor_pairs = lattice.get_neighbor_pairs(k=1, unique=True)
+    if not neighbor_pairs:
+        return coo_matrix((2**num_sites, 2**num_sites))
+
+    pauli_map = {"X": 1, "Y": 2, "Z": 3}
+
+    ls: typing.List[typing.List[int]] = []
+    weights: typing.List[float] = []
+
+    for i, j in neighbor_pairs:
+        xx_string = [0] * num_sites
+        xx_string[i] = pauli_map["X"]
+        xx_string[j] = pauli_map["X"]
+        ls.append(xx_string)
+        weights.append(j_coupling)
+
+        yy_string = [0] * num_sites
+        yy_string[i] = pauli_map["Y"]
+        yy_string[j] = pauli_map["Y"]
+        ls.append(yy_string)
+        weights.append(j_coupling)
+
+        zz_string = [0] * num_sites
+        zz_string[i] = pauli_map["Z"]
+        zz_string[j] = pauli_map["Z"]
+        ls.append(zz_string)
+        weights.append(j_coupling)
+
+    hamiltonian_matrix = tc.quantum.PauliStringSum2COO(ls, weight=weights, numpy=True)
+
+    return cast(coo_matrix, hamiltonian_matrix)
+
+
+def generate_rydberg_hamiltonian(
+    lattice: AbstractLattice, omega: float, delta: float, c6: float
+) -> coo_matrix:
+    """
+    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 represented as a SciPy COO sparse matrix.
+    :rtype: coo_matrix
+    """
+    num_sites = lattice.num_sites
+    if num_sites == 0:
+        return coo_matrix((0, 0))
+
+    pauli_map = {"X": 1, "Y": 2, "Z": 3}
+    ls: typing.List[typing.List[int]] = []
+    weights: typing.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)
+
+            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(float(z_coefficients[i]))
+
+    hamiltonian_matrix = tc.quantum.PauliStringSum2COO(ls, weight=weights, numpy=True)
+
+    return cast(coo_matrix, hamiltonian_matrix)

tests/test_hamiltonians.py

@@ -0,0 +1,132 @@
+import pytest
+import numpy as np
+
+from tensorcircuit.templates.lattice import (
+    ChainLattice,
+    SquareLattice,
+    CustomizeLattice,
+)
+from tensorcircuit.templates.hamiltonians import (
+    generate_heisenberg_hamiltonian,
+    generate_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 generate_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 = generate_heisenberg_hamiltonian(empty_lattice)
+        assert h.shape == (0, 0)
+        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 = generate_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 = generate_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(h_generated.toarray(), 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 = generate_heisenberg_hamiltonian(lattice, j_coupling=1.0)
+
+        assert h.shape == (16, 16)
+        assert h.nnz > 0
+        h_dense = h.toarray()
+        assert np.allclose(h_dense, h_dense.conj().T)
+
+
+class TestRydbergHamiltonian:
+    """
+    Test suite for the generate_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 = generate_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(h_generated.toarray(), 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 = generate_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)
+        assert np.allclose(h_generated.toarray(), 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 = generate_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.")