refactor(templates): Improve hamiltonian functions based on review

This commit refactors the hamiltonian generation functions to be more robust, flexible, and aligned with the project's backend-agnostic architecture, based on feedback from the recent code review.

Key changes include:

- **Backend Agnosticism:**
  The `PauliStringSum2COO` call now uses `numpy=False`, ensuring the returned sparse matrix is a native tensor of the currently active backend (TensorFlow, JAX, etc.). The return type hints have been updated to `Any` to reflect this.

- **Anisotropic Heisenberg Model:**
  The `heisenberg_hamiltonian` function now accepts a list or tuple for `j_coupling`, allowing for the definition of anisotropic models with different Jx, Jy, and Jz values.

- **Test Suite Enhancement:**
  Tests have been updated to use `tc.backend.to_dense` for backend-agnostic validation. A new test case for the anisotropic Heisenberg model has also been added to ensure its correctness.

- **Code Cleanup:**
  A redundant `float()` cast was removed from the Rydberg hamiltonian logic.

Author: Stellogic

tensorcircuit/templates/hamiltonians.py

@@ -1,14 +1,24 @@
 import typing
-from typing import cast
-from scipy.sparse import coo_matrix
+from typing import Any, List, Tuple, Union
 import numpy as np
+from tensorcircuit.cons import dtypestr, backend
 import tensorcircuit as tc
 from .lattice import AbstractLattice
 
 
-def generate_heisenberg_hamiltonian(
-    lattice: AbstractLattice, j_coupling: float = 1.0
-) -> coo_matrix:
+def _create_empty_sparse_matrix(shape: Tuple[int, int]) -> Any:
+    """
+    Helper function to create a backend-agnostic empty sparse matrix.
+    """
+    indices = tc.backend.convert_to_tensor(backend.zeros((0, 2), dtype="int32"))
+    values = tc.backend.convert_to_tensor(backend.zeros((0,), dtype=dtypestr))  # type: ignore
+    return tc.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.
 
@@ -19,51 +29,49 @@ def generate_heisenberg_hamiltonian(
     :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
+    :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
-    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))
+
+    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 num_sites == 0 or not neighbor_pairs:
+        return _create_empty_sparse_matrix(shape=(2**num_sites, 2**num_sites))
 
     pauli_map = {"X": 1, "Y": 2, "Z": 3}
 
     ls: typing.List[typing.List[int]] = []
     weights: typing.List[float] = []
 
+    pauli_terms = ["X", "Y", "Z"]
     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)
+        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 = tc.quantum.PauliStringSum2COO(ls, weight=weights, numpy=True)
+    hamiltonian_matrix = tc.quantum.PauliStringSum2COO(ls, weight=weights, numpy=False)
 
-    return cast(coo_matrix, hamiltonian_matrix)
+    return hamiltonian_matrix
 
 
-def generate_rydberg_hamiltonian(
+def rydberg_hamiltonian(
     lattice: AbstractLattice, omega: float, delta: float, c6: float
-) -> coo_matrix:
+) -> Any:
     """
     Generates the sparse matrix of the Rydberg atom array Hamiltonian.
 
@@ -84,12 +92,12 @@ def generate_rydberg_hamiltonian(
     :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
+    :return: The Hamiltonian as a backend-agnostic sparse matrix.
+    :rtype: Any
     """
     num_sites = lattice.num_sites
     if num_sites == 0:
-        return coo_matrix((0, 0))
+        return _create_empty_sparse_matrix(shape=(1, 1))
 
     pauli_map = {"X": 1, "Y": 2, "Z": 3}
     ls: typing.List[typing.List[int]] = []
@@ -132,8 +140,8 @@ def generate_rydberg_hamiltonian(
             z_string = [0] * num_sites
             z_string[i] = pauli_map["Z"]
             ls.append(z_string)
-            weights.append(float(z_coefficients[i]))
+            weights.append(z_coefficients[i])  # type: ignore
 
-    hamiltonian_matrix = tc.quantum.PauliStringSum2COO(ls, weight=weights, numpy=True)
+    hamiltonian_matrix = tc.quantum.PauliStringSum2COO(ls, weight=weights, numpy=False)
 
-    return cast(coo_matrix, hamiltonian_matrix)
+    return hamiltonian_matrix

tests/test_hamiltonians.py

@@ -1,14 +1,16 @@
 import pytest
 import numpy as np
+import tensorcircuit as tc
+
 
 from tensorcircuit.templates.lattice import (
     ChainLattice,
     SquareLattice,
     CustomizeLattice,
 )
 from tensorcircuit.templates.hamiltonians import (
-    generate_heisenberg_hamiltonian,
-    generate_rydberg_hamiltonian,
+    heisenberg_hamiltonian,
+    rydberg_hamiltonian,
 )
 
 PAULI_X = np.array([[0, 1], [1, 0]], dtype=complex)
@@ -19,7 +21,7 @@
 
 class TestHeisenbergHamiltonian:
     """
-    Test suite for the generate_heisenberg_hamiltonian function.
+    Test suite for the heisenberg_hamiltonian function.
     """
 
     def test_empty_lattice(self):
@@ -29,16 +31,16 @@ def test_empty_lattice(self):
         empty_lattice = CustomizeLattice(
             dimensionality=2, identifiers=[], coordinates=[]
         )
-        h = generate_heisenberg_hamiltonian(empty_lattice)
-        assert h.shape == (0, 0)
+        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 = generate_heisenberg_hamiltonian(single_site_lattice)
+        h = heisenberg_hamiltonian(single_site_lattice)
         assert h.shape == (2, 2)
         assert h.nnz == 0
 
@@ -49,7 +51,7 @@ def test_two_sites_chain(self):
         """
         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)
+        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)
@@ -58,24 +60,24 @@ def test_two_sites_chain(self):
         h_expected = j_coupling * (xx + yy + zz)
 
         assert h_generated.shape == (4, 4)
-        assert np.allclose(h_generated.toarray(), h_expected)
+        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 = generate_heisenberg_hamiltonian(lattice, j_coupling=1.0)
+        h = heisenberg_hamiltonian(lattice, j_coupling=1.0)
 
         assert h.shape == (16, 16)
         assert h.nnz > 0
-        h_dense = h.toarray()
+        h_dense = tc.backend.to_dense(h)
         assert np.allclose(h_dense, h_dense.conj().T)
 
 
 class TestRydbergHamiltonian:
     """
-    Test suite for the generate_rydberg_hamiltonian function.
+    Test suite for the rydberg_hamiltonian function.
     """
 
     def test_single_site_rydberg(self):
@@ -84,20 +86,20 @@ def test_single_site_rydberg(self):
         """
         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_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(h_generated.toarray(), h_expected)
+        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 = generate_rydberg_hamiltonian(lattice, omega, delta, c6)
+        h_generated = rydberg_hamiltonian(lattice, omega, delta, c6)
 
         v_ij = c6 / (1.5**6)
 
@@ -110,7 +112,9 @@ def test_two_sites_rydberg(self):
         h_expected = h1 + h2 + h3
 
         assert h_generated.shape == (4, 4)
-        assert np.allclose(h_generated.toarray(), h_expected)
+        h_generated_dense = tc.backend.to_dense(h_generated)
+
+        assert np.allclose(h_generated_dense, h_expected)
 
     def test_zero_distance_robustness(self):
         """
@@ -123,10 +127,29 @@ def test_zero_distance_robustness(self):
         )
 
         try:
-            h = generate_rydberg_hamiltonian(lattice, omega=1.0, delta=1.0, c6=1.0)
+            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)