tensorcircuit-nightly 1.3.0.dev20250728__py3-none-any.whl → 1.4.0.dev20251103__py3-none-any.whl

tests/test_fgs.py

@@ -1,318 +0,0 @@
-import numpy as np
-import pytest
-from pytest_lazyfixture import lazy_fixture as lf
-import tensorflow as tf
-
-try:
-    import openfermion as _
-except ModuleNotFoundError:
-    pytestmark = pytest.mark.skip("skip fgs test due to missing openfermion module")
-
-import tensorcircuit as tc
-
-F = tc.fgs.FGSSimulator
-FT = tc.fgs.FGSTestSimulator
-
-
-@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
-def test_cmatrix(backend, highp):
-    import openfermion
-
-    def asymmetric_hopping(J, gamma, i, j, L):
-        m0 = tc.fgs.onehot_matrix(i, j, 2 * L) - tc.fgs.onehot_matrix(
-            j + L, i + L, 2 * L
-        )
-        m = (J + gamma) * m0 + (J - gamma) * tc.backend.adjoint(m0)
-        return m / 2
-
-    def asymmetric_hopping_jw(J, gamma, i, j, L):
-        op = (J + gamma) * openfermion.FermionOperator(f"{str(i)}^ {str(j)}") + np.conj(
-            J - gamma
-        ) * openfermion.FermionOperator(f"{str(j)}^ {str(i)}")
-        sop = openfermion.transforms.jordan_wigner(op)
-        m = openfermion.get_sparse_operator(sop, n_qubits=L).todense()
-        return m
-
-    c = tc.fgs.FGSSimulator(4, [0])
-    c1 = tc.fgs.FGSTestSimulator(4, [0])
-    c.evol_hp(0, 1, 1.2)
-    c1.evol_hp(0, 1, 1.2)
-    c.evol_icp(2, 0.5)
-    c1.evol_icp(2, 0.5)
-    c.evol_cp(0, -0.8)
-    c1.evol_cp(0, -0.8)
-    c.evol_sp(1, 0, 0.5)
-    c1.evol_sp(1, 0, 0.5)
-    c.evol_ghamiltonian(asymmetric_hopping(1, 0.5, 0, 2, 4))
-    c1.evol_ghamiltonian(np.array(asymmetric_hopping_jw(1, 0.5, 0, 2, 4)))
-    c.evol_sp(1, 2, 1.5)
-    c1.evol_sp(1, 2, 1.5)
-    c.evol_ihamiltonian(c.chemical_potential(1.0, 1, 4))
-    c1.evol_ihamiltonian(c1.chemical_potential_jw(1.0, 1, 4))
-    np.testing.assert_allclose(
-        c1.get_cmatrix_majorana(), c.get_cmatrix_majorana(), atol=1e-5
-    )
-    np.testing.assert_allclose(c1.get_cmatrix(), c.get_cmatrix(), atol=1e-5)
-    np.testing.assert_allclose(c1.entropy([0, 2]), c.entropy([0, 2]), atol=1e-5)
-    np.testing.assert_allclose(
-        c1.renyi_entropy(2, [1]), c.renyi_entropy(2, [1]), atol=1e-5
-    )
-
-
-def test_full_trace():
-    c = tc.FGSSimulator(4)
-    with pytest.raises(ValueError):
-        c.entropy([0, 1, 2, 3])
-
-
-@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
-def test_otoc(backend, highp):
-    c = tc.FGSSimulator(4, [0, 1])
-    c1 = tc.fgs.FGSTestSimulator(4, [0, 1])
-    h = c.hopping(0.6, 1, 2, 4)
-    h += c.hopping(-0.8, 3, 2, 4)
-    h += c.chemical_potential(0.3, 2, 4)
-    h1 = c1.hopping_jw(0.6, 1, 2, 4)
-    h1 += c1.hopping_jw(-0.8, 3, 2, 4)
-    h1 += c1.chemical_potential_jw(0.3, 2, 4)
-    c.evol_hamiltonian(h)
-    m = c.get_cmatrix(now_i=True, now_j=False)
-    m0 = c1.get_ot_cmatrix(h1)
-    np.testing.assert_allclose(m, m0, atol=1e-5)
-    m = c.get_cmatrix(now_i=False, now_j=True)
-    m0 = c1.get_ot_cmatrix(h1, now_i=False)
-    np.testing.assert_allclose(m, m0, atol=1e-5)
-
-
-@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
-def test_fgs_ad(backend, highp):
-    import optax
-
-    N = 18
-
-    def f(chi):
-        c = tc.fgs.FGSSimulator(N, [0])
-        for i in range(N - 1):
-            c.evol_hp(i, i + 1, chi[i])
-        cm = c.get_cmatrix()
-        return -tc.backend.real(1 - cm[N, N])
-
-    chi_vg = tc.backend.jit(tc.backend.value_and_grad(f))
-
-    if tc.backend.name == "tensorflow":
-        opt = tc.backend.optimizer(tf.keras.optimizers.Adam(1e-2))
-    else:
-        opt = tc.backend.optimizer(optax.adam(1e-2))
-
-    param = tc.backend.ones([N - 1], dtype="float64")
-    for _ in range(300):
-        vs, gs = chi_vg(param)
-        param = opt.update(gs, param)
-
-    np.testing.assert_allclose(vs, -0.9986, atol=1e-2)
-
-
-@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
-def test_hamiltonian_generation(backend):
-    hc = F.chemical_potential(0.8, 0, 3)
-    h1 = FT.get_hmatrix(hc, 3) + 0.4 * tc.backend.eye(2**3)
-    # constant shift
-    h2 = FT.chemical_potential_jw(0.8, 0, 3)
-    np.testing.assert_allclose(h1, h2, atol=1e-5)
-
-    hc = F.hopping(0.8j, 0, 1, 3)
-    h1 = FT.get_hmatrix(hc, 3)
-    h2 = FT.hopping_jw(0.8j, 0, 1, 3)
-    np.testing.assert_allclose(h1, h2, atol=1e-5)
-
-    hc = F.sc_pairing(0.8, 1, 0, 3)
-    h1 = FT.get_hmatrix(hc, 3)
-    h2 = FT.sc_pairing_jw(0.8, 1, 0, 3)
-    np.testing.assert_allclose(h1, h2, atol=1e-5)
-
-
-@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
-def test_ground_state(backend, highp):
-    N = 3
-    hc = (
-        F.hopping(1.0, 0, 1, N)
-        + F.hopping(1.0, 1, 2, N)
-        + F.chemical_potential(0.4, 0, N)
-        - F.chemical_potential(0.4, 1, N)
-        + F.sc_pairing(0.8, 0, 1, N)
-    )
-
-    c = F(N, hc=hc)
-    c1 = FT(N, hc=hc)
-
-    np.testing.assert_allclose(c.get_cmatrix(), c1.get_cmatrix(), atol=1e-6)
-    np.testing.assert_allclose(c1.entropy([0]), c.entropy([0]), atol=1e-5)
-
-
-@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
-def test_post_select(backend, highp):
-    for ind in [0, 1, 2]:
-        for keep in [0, 1]:
-            c = F(3, filled=[0, 2])
-            c.evol_hp(0, 1, 0.2)
-            c.evol_cp(0, 0.5)
-            c.evol_hp(1, 2, 0.3j)
-            c.evol_cp(1, -0.8)
-            c.evol_sp(0, 2, 0.3)
-            c.post_select(ind, keep)
-            c1 = FT(3, filled=[0, 2])
-            c1.evol_hp(0, 1, 0.2)
-            c1.evol_cp(0, 0.5)
-            c1.evol_hp(1, 2, 0.3j)
-            c1.evol_cp(1, -0.8)
-            c1.evol_sp(0, 2, 0.3)
-            c1.post_select(ind, keep)
-            np.testing.assert_allclose(c.get_cmatrix(), c1.get_cmatrix(), atol=1e-5)
-
-
-@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
-def test_post_select_cmatrix_refresh(backend, highp):
-    sim_post = tc.FGSSimulator(L=2, filled=[0])
-    sim_post.evol_hp(0, 1, np.pi / 4)
-    sim_post.post_select(0, keep=1)
-    np.testing.assert_allclose(1 - sim_post.get_cmatrix()[0, 0], 1, atol=1e-6)
-
-
-@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
-def test_jittable_measure(backend):
-    @tc.backend.jit
-    def get_cmatrix(status):
-        r = []
-        c = F(3, filled=[0])
-        c.evol_hp(0, 1, 0.2)
-        c.evol_cp(0, 0.5)
-        c.evol_hp(1, 2, 0.3j)
-        r.append(c.cond_measure(1, status[0]))
-        c.evol_cp(1, -0.8)
-        r.append(c.cond_measure(2, status[1]))
-        r.append(c.cond_measure(1, status[3]))
-        c.evol_sp(0, 2, 0.3)
-        c.evol_hp(2, 1, 0.2)
-        c.evol_hp(0, 1, 0.8)
-        return c.get_cmatrix(), tc.backend.stack(r)
-
-    def get_cmatrix_baseline(status):
-        r = []
-        c = FT(3, filled=[0])
-        c.evol_hp(0, 1, 0.2)
-        c.evol_cp(0, 0.5)
-        c.evol_hp(1, 2, 0.3j)
-        r.append(c.cond_measure(1, status[0]))
-        c.evol_cp(1, -0.8)
-        r.append(c.cond_measure(2, status[1]))
-        r.append(c.cond_measure(1, status[3]))
-        c.evol_sp(0, 2, 0.3)
-        c.evol_hp(2, 1, 0.2)
-        c.evol_hp(0, 1, 0.8)
-        return c.get_cmatrix(), np.array(r)
-
-    status = np.array([0.2, 0.83, 0.61, 0.07])
-    m, his = get_cmatrix(tc.backend.convert_to_tensor(status))
-    m1, his1 = get_cmatrix_baseline(status)
-    np.testing.assert_allclose(m, m1, atol=1e-5)
-    np.testing.assert_allclose(his, his1, atol=1e-5)
-
-
-@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
-def test_exp_2body(backend):
-    c = F(3, filled=[1])
-    np.testing.assert_allclose(c.expectation_2body(4, 1), 1.0, atol=1e-5)
-    np.testing.assert_allclose(c.expectation_2body(5, 2), 0.0, atol=1e-5)
-    np.testing.assert_allclose(c.expectation_2body(1, 4), 0.0, atol=1e-5)
-
-
-@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
-def test_exp_4body(backend):
-    c = F(4, filled=[0, 2])
-    c.evol_hp(0, 1, 0.3)
-    c.evol_hp(2, 3, 0.3)
-    c.evol_sp(0, 3, 0.5)
-    c.evol_sp(0, 2, 0.9)
-    c.evol_cp(0, -0.4)
-
-    c1 = FT(4, filled=[0, 2])
-    c1.evol_hp(0, 1, 0.3)
-    c1.evol_hp(2, 3, 0.3)
-    c1.evol_sp(0, 3, 0.5)
-    c1.evol_sp(0, 2, 0.9)
-    c1.evol_cp(0, -0.4)
-
-    np.testing.assert_allclose(
-        c.expectation_4body(0, 4, 1, 5), c.expectation_4body(0, 4, 1, 5), atol=1e-5
-    )
-    np.testing.assert_allclose(
-        c.expectation_4body(0, 1, 4, 5), c.expectation_4body(0, 1, 4, 5), atol=1e-5
-    )
-    np.testing.assert_allclose(
-        c.expectation_4body(0, 4, 2, 6), c.expectation_4body(0, 4, 2, 6), atol=1e-5
-    )
-    np.testing.assert_allclose(
-        c.expectation_4body(0, 2, 3, 1), c.expectation_4body(0, 2, 3, 1), atol=1e-5
-    )
-    np.testing.assert_allclose(
-        c.expectation_4body(0, 1, 4, 6), c.expectation_4body(0, 1, 4, 6), atol=1e-5
-    )
-    np.testing.assert_allclose(
-        c.expectation_4body(1, 0, 6, 7), c.expectation_4body(1, 0, 6, 7), atol=1e-5
-    )
-
-
-@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
-def test_overlap(backend):
-    def compute_overlap(FGScls):
-        c = FGScls(3, filled=[0, 2])
-        c.evol_hp(0, 1, 1.2)
-        c.evol_hp(1, 2, 0.3)
-        c.evol_cp(0, 0.5)
-        c.evol_sp(0, 2, 0.3)
-        c1 = FGScls(3, filled=[0, 2])
-        c1.evol_hp(0, 1, 0.2)
-        c1.evol_hp(1, 2, 0.6)
-        c1.evol_sp(1, 0, -1.1)
-        return tc.backend.abs(c.overlap(c1))
-
-    np.testing.assert_allclose(
-        compute_overlap(tc.FGSSimulator), compute_overlap(tc.fgs.FGSTestSimulator), 1e-5
-    )
-
-
-@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
-def test_entanglement_asymmetry(backend, highp):
-    def discrete_c(c):
-        c.evol_sp(0, 1, 0.3)
-        c.evol_sp(1, 3, 1.1)
-        c.evol_hp(0, 1, -0.2)
-        c.evol_hp(2, 3, 1.4)
-        c.evol_hp(3, 4, -0.6)
-        c.evol_hp(4, 5, -0.8)
-        c.evol_sp(2, 6, 2.1)
-        return c
-
-    c = discrete_c(tc.fgs.FGSSimulator(7, filled=[0, 2, 4, 6]))
-    c1 = discrete_c(tc.fgs.FGSTestSimulator(7, filled=[0, 2, 4, 6]))
-    for k in [2, 3, 4]:
-        for traceout in [[1, 2, 4, 5], [2, 3, 5], [0, 3, 4, 5, 6]]:
-            print(k, traceout)
-            sa_ans = c1.renyi_entropy(k, traceout)
-            saq_ans = c1.renyi_entanglement_asymmetry(k, traceout)
-            np.testing.assert_allclose(
-                np.log(c.charge_moment([0 for _ in range(k)], k, traceout))
-                * 1
-                / (1 - k),
-                sa_ans,
-                atol=2e-4,
-            )
-            saq, std = c.renyi_entanglement_asymmetry(
-                k, traceout, batch=400, with_std=True
-            )
-            np.testing.assert_allclose(
-                saq,
-                saq_ans,
-                atol=6 * std,
-            )
-            print(std)

tests/test_gates.py

@@ -1,156 +0,0 @@
-import sys
-import os
-import numpy as np
-
-thisfile = os.path.abspath(__file__)
-modulepath = os.path.dirname(os.path.dirname(thisfile))
-
-sys.path.insert(0, modulepath)
-import tensorcircuit as tc
-
-
-def test_rgate(highp):
-    np.testing.assert_almost_equal(
-        tc.gates.r_gate(1, 2, 3).tensor, tc.gates.rgate_theoretical(1, 2, 3).tensor
-    )
-
-
-def test_phase_gate():
-    c = tc.Circuit(1)
-    c.h(0)
-    c.phase(0, theta=np.pi / 2)
-    np.testing.assert_allclose(c.state()[1], 0.7071j, atol=1e-4)
-
-
-def test_cu_gate():
-    c = tc.Circuit(2)
-    c.cu(0, 1, theta=np.pi / 2, phi=-np.pi / 4, lbd=np.pi / 4)
-    m = c.matrix()
-    print(m)
-    np.testing.assert_allclose(m[2:, 2:], tc.gates._wroot_matrix, atol=1e-5)
-    np.testing.assert_allclose(m[:2, :2], np.eye(2), atol=1e-5)
-
-
-def test_get_u_parameter(highp):
-    for _ in range(6):
-        hermitian = np.random.uniform(size=[2, 2])
-        hermitian += np.conj(np.transpose(hermitian))
-        unitary = tc.backend.expm(hermitian * 1.0j)
-        params = tc.gates.get_u_parameter(unitary)
-        unitary2 = tc.gates.u_gate(theta=params[0], phi=params[1], lbd=params[2])
-        ans = unitary2.tensor
-        unitary = unitary / np.exp(1j * np.angle(unitary[0, 0]))
-        np.testing.assert_allclose(unitary, ans, atol=1e-3)
-
-
-def test_ided_gate():
-    g = tc.gates.rx.ided()
-    np.testing.assert_allclose(
-        tc.backend.reshapem(g(theta=0.3).tensor),
-        np.kron(np.eye(2), tc.gates.rx(theta=0.3).tensor),
-        atol=1e-5,
-    )
-    g1 = tc.gates.rx.ided(before=False)
-    np.testing.assert_allclose(
-        tc.backend.reshapem(g1(theta=0.3).tensor),
-        np.kron(tc.gates.rx(theta=0.3).tensor, np.eye(2)),
-        atol=1e-5,
-    )
-
-
-def test_fsim_gate():
-    theta = 0.2
-    phi = 0.3
-    c = tc.Circuit(2)
-    c.iswap(0, 1, theta=-theta)
-    c.cphase(0, 1, theta=-phi)
-    m = c.matrix()
-    ans = np.array(
-        [
-            [1.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
-            [0.0 + 0.0j, 0.95105654 + 0.0j, 0.0 - 0.309017j, 0.0 + 0.0j],
-            [0.0 + 0.0j, 0.0 - 0.309017j, 0.95105654 + 0.0j, 0.0 + 0.0j],
-            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.9553365 - 0.29552022j],
-        ]
-    )
-    np.testing.assert_allclose(m, ans, atol=1e-5)
-    print(m)
-
-
-def test_exp_gate():
-    c = tc.Circuit(2)
-    c.exp(
-        0,
-        1,
-        unitary=tc.gates.array_to_tensor(
-            np.array([[1.0, 0, 0, 0], [0, -1.0, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
-        ),
-        theta=tc.gates.num_to_tensor(np.pi / 2),
-    )
-    np.testing.assert_allclose(c.wavefunction()[0], -1j)
-
-
-def test_any_gate():
-    c = tc.Circuit(2)
-    c.any(0, unitary=np.eye(2))
-    np.testing.assert_allclose(c.expectation((tc.gates.z(), [0])), 1.0)
-
-
-def test_iswap_gate():
-    t = tc.gates.iswap_gate().tensor
-    ans = np.array([[1.0, 0, 0, 0], [0, 0, 1j, 0], [0, 1j, 0, 0], [0, 0, 0, 1.0]])
-    np.testing.assert_allclose(t, ans.reshape([2, 2, 2, 2]), atol=1e-5)
-    t = tc.gates.iswap_gate(theta=0).tensor
-    np.testing.assert_allclose(t, np.eye(4).reshape([2, 2, 2, 2]), atol=1e-5)
-
-
-def test_gate_list():
-    assert tc.Circuit.sgates == tc.abstractcircuit.sgates
-
-
-def test_controlled():
-    xgate = tc.gates.x
-    cxgate = xgate.controlled()
-    ccxgate = cxgate.controlled()
-    assert ccxgate.n == "ccx"
-    assert ccxgate.ctrl == [1, 1]
-    np.testing.assert_allclose(
-        ccxgate().tensor, tc.backend.reshape2(tc.gates._toffoli_matrix)
-    )
-    ocxgate = cxgate.ocontrolled()
-    c = tc.Circuit(3)
-    c.x(0)
-    c.any(1, 0, 2, unitary=ocxgate())
-    np.testing.assert_allclose(c.expectation([tc.gates.z(), [2]]), -1, atol=1e-5)
-    print(c.to_qir()[1])
-
-
-def test_variable_controlled():
-    crxgate = tc.gates.rx.controlled()
-    c = tc.Circuit(2)
-    c.x(0)
-    tc.Circuit.crx_my = tc.Circuit.apply_general_variable_gate_delayed(crxgate)
-    c.crx_my(0, 1, theta=0.3)
-    np.testing.assert_allclose(
-        c.expectation([tc.gates.z(), [1]]), 0.95533645, atol=1e-5
-    )
-    assert c.to_qir()[1]["name"] == "crx"
-
-
-def test_adjoint_gate():
-    np.testing.assert_allclose(
-        tc.gates.sd().tensor, tc.backend.adjoint(tc.gates._s_matrix)
-    )
-    assert tc.gates.td.n == "td"
-
-
-def test_rxx_gate():
-    c1 = tc.Circuit(3)
-    c1.rxx(0, 1, theta=1.0)
-    c1.ryy(0, 2, theta=0.5)
-    c1.rzz(0, 1, theta=-0.5)
-    c2 = tc.Circuit(3)
-    c2.exp1(0, 1, theta=1.0 / 2, unitary=tc.gates._xx_matrix)
-    c2.exp1(0, 2, theta=0.5 / 2, unitary=tc.gates._yy_matrix)
-    c2.exp1(0, 1, theta=-0.5 / 2, unitary=tc.gates._zz_matrix)
-    np.testing.assert_allclose(c1.state(), c2.state(), atol=1e-5)

tests/test_hamiltonians.py

@@ -1,159 +0,0 @@
-import pytest
-import numpy as np
-from pytest_lazyfixture import lazy_fixture as lf
-
-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
-
-    @pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
-    def test_two_sites_chain(self, backend):
-        """
-        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)
-        print(tc.backend.to_dense(h_generated))
-        assert np.allclose(tc.backend.to_dense(h_generated), h_expected, atol=1e-5)
-
-    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.")
-
-    @pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
-    def test_anisotropic_heisenberg(self, backend):
-        """
-        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)