diff --git a/tensorcircuit/templates/hamiltonians.py b/tensorcircuit/templates/hamiltonians.py new file mode 100644 index 00000000..0382b11d --- /dev/null +++ b/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 * Σ_{} (X_i X_j + Y_i Y_j + Z_i Z_j) + where the sum is over all unique nearest-neighbor pairs . + + :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 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 diff --git a/tests/test_hamiltonians.py b/tests/test_hamiltonians.py new file mode 100644 index 00000000..28c1ba33 --- /dev/null +++ b/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)