From c28deb09c9bd74abdfb3f63dbbd5b6bb6a605d85 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 5 Mar 2026 21:07:15 +0000
Subject: [PATCH 1/4] Initial plan


From f4b89cae4637f86025530f6d4d795e9c7743552c Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 5 Mar 2026 21:19:57 +0000
Subject: [PATCH 2/4] Add tests for hamiltonian.py in tests/ folder

Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com>
---
 tests/test_hamiltonian.py | 292 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 292 insertions(+)
 create mode 100644 tests/test_hamiltonian.py

diff --git a/tests/test_hamiltonian.py b/tests/test_hamiltonian.py
new file mode 100644
index 0000000..fede404
--- /dev/null
+++ b/tests/test_hamiltonian.py
@@ -0,0 +1,292 @@
+"""
+Tests for qmp.hamiltonian.Hamiltonian.
+
+We use a 2-qubit fermionic Hamiltonian (n_qubytes=1):
+
+    H = c†₀c₁ + c†₁c₀ + n₀ + n₁
+
+where c†ᵢ / cᵢ are creation / annihilation operators and nᵢ = c†ᵢcᵢ is the
+number operator at site i.
+
+Operator-key convention in the Hamiltonian dict:
+    Key   ((s₀, k₀), (s₁, k₁), ...)  →  operator  O_{k₀,s₀} O_{k₁,s₁}
+    with  k=0: annihilation,  k=1: creation.
+    The operators are applied right-to-left to the quantum state, i.e. the
+    rightmost operator in the expression acts on the state first.
+
+Occupation-number basis for 2 qubits (bit 0 = site 0, bit 1 = site 1):
+    |00⟩  →  uint8 value 0
+    |10⟩  →  uint8 value 1  (site 0 occupied)
+    |01⟩  →  uint8 value 2  (site 1 occupied)
+    |11⟩  →  uint8 value 3  (both sites occupied)
+
+In this basis, the Hamiltonian matrix is:
+    H = [[0, 0, 0, 0],
+         [0, 1, 1, 0],
+         [0, 1, 1, 0],
+         [0, 0, 0, 2]]
+"""
+
+import pytest
+import torch
+
+from qmp.hamiltonian import Hamiltonian
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="session")
+def two_qubit_hamiltonian() -> Hamiltonian:
+    """
+    2-qubit fermionic Hamiltonian: H = c†₀c₁ + c†₁c₀ + n₀ + n₁.
+
+    Compiled and cached once per test session to avoid repeated JIT overhead.
+    """
+    return Hamiltonian(
+        {
+            ((0, 1), (1, 0)): 1.0,  # c†₀c₁
+            ((1, 1), (0, 0)): 1.0,  # c†₁c₀
+            ((0, 1), (0, 0)): 1.0,  # n₀ = c†₀c₀
+            ((1, 1), (1, 0)): 1.0,  # n₁ = c†₁c₁
+        },
+        kind="fermi",
+    )
+
+
+# ---------------------------------------------------------------------------
+# Helpers – configs for 2-qubit system (n_qubytes = 1)
+# ---------------------------------------------------------------------------
+
+C_00 = torch.tensor([[0]], dtype=torch.uint8)  # |00⟩
+C_10 = torch.tensor([[1]], dtype=torch.uint8)  # |10⟩  (site 0 occupied)
+C_01 = torch.tensor([[2]], dtype=torch.uint8)  # |01⟩  (site 1 occupied)
+C_11 = torch.tensor([[3]], dtype=torch.uint8)  # |11⟩  (both sites occupied)
+
+ALL_CONFIGS = torch.tensor([[0], [1], [2], [3]], dtype=torch.uint8)
+
+
+# ---------------------------------------------------------------------------
+# Tests for diagonal_term
+# ---------------------------------------------------------------------------
+
+
+class TestDiagonalTerm:
+    """diagonal_term(configs) returns the diagonal element H[i,i] for each config."""
+
+    def test_vacuum_diagonal_is_zero(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        result = two_qubit_hamiltonian.diagonal_term(C_00)
+        expected = torch.tensor([0.0 + 0j], dtype=torch.complex128)
+        torch.testing.assert_close(result, expected)
+
+    def test_single_particle_site0(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        result = two_qubit_hamiltonian.diagonal_term(C_10)
+        expected = torch.tensor([1.0 + 0j], dtype=torch.complex128)
+        torch.testing.assert_close(result, expected)
+
+    def test_single_particle_site1(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        result = two_qubit_hamiltonian.diagonal_term(C_01)
+        expected = torch.tensor([1.0 + 0j], dtype=torch.complex128)
+        torch.testing.assert_close(result, expected)
+
+    def test_two_particle_sector(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        result = two_qubit_hamiltonian.diagonal_term(C_11)
+        expected = torch.tensor([2.0 + 0j], dtype=torch.complex128)
+        torch.testing.assert_close(result, expected)
+
+    def test_all_configs_batch(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """Batched evaluation matches individual results."""
+        result = two_qubit_hamiltonian.diagonal_term(ALL_CONFIGS)
+        expected = torch.tensor([0.0 + 0j, 1.0 + 0j, 1.0 + 0j, 2.0 + 0j], dtype=torch.complex128)
+        torch.testing.assert_close(result, expected)
+
+
+# ---------------------------------------------------------------------------
+# Tests for apply_within
+# ---------------------------------------------------------------------------
+
+
+class TestApplyWithin:
+    """apply_within(configs_i, psi_i, configs_j) computes (H·ψ) restricted to configs_j."""
+
+    def test_apply_to_basis_state_10(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """H|10⟩ = |10⟩ + |01⟩ in the 1-particle sector."""
+        configs = torch.cat([C_10, C_01])
+        psi_i = torch.tensor([1.0 + 0j, 0.0 + 0j], dtype=torch.complex128)
+        result = two_qubit_hamiltonian.apply_within(configs, psi_i, configs)
+        expected = torch.tensor([1.0 + 0j, 1.0 + 0j], dtype=torch.complex128)
+        torch.testing.assert_close(result, expected)
+
+    def test_apply_to_basis_state_01(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """H|01⟩ = |10⟩ + |01⟩ in the 1-particle sector."""
+        configs = torch.cat([C_10, C_01])
+        psi_i = torch.tensor([0.0 + 0j, 1.0 + 0j], dtype=torch.complex128)
+        result = two_qubit_hamiltonian.apply_within(configs, psi_i, configs)
+        expected = torch.tensor([1.0 + 0j, 1.0 + 0j], dtype=torch.complex128)
+        torch.testing.assert_close(result, expected)
+
+    def test_apply_to_two_particle_state(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """H|11⟩ = 2|11⟩ (both number operators contribute; hopping is blocked)."""
+        psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
+        result = two_qubit_hamiltonian.apply_within(C_11, psi_i, C_11)
+        expected = torch.tensor([2.0 + 0j], dtype=torch.complex128)
+        torch.testing.assert_close(result, expected)
+
+    def test_no_transition_to_vacuum(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """H maps no configuration to |00⟩."""
+        psi_i = torch.ones(4, dtype=torch.complex128)
+        result = two_qubit_hamiltonian.apply_within(ALL_CONFIGS, psi_i, C_00)
+        expected = torch.tensor([0.0 + 0j], dtype=torch.complex128)
+        torch.testing.assert_close(result, expected)
+
+    def test_off_diagonal_10_to_01(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """c†₁c₀|10⟩ → |01⟩ with amplitude 1."""
+        psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
+        result = two_qubit_hamiltonian.apply_within(C_10, psi_i, C_01)
+        expected = torch.tensor([1.0 + 0j], dtype=torch.complex128)
+        torch.testing.assert_close(result, expected)
+
+    def test_off_diagonal_01_to_10(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """c†₀c₁|01⟩ → |10⟩ with amplitude 1."""
+        psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
+        result = two_qubit_hamiltonian.apply_within(C_01, psi_i, C_10)
+        expected = torch.tensor([1.0 + 0j], dtype=torch.complex128)
+        torch.testing.assert_close(result, expected)
+
+    def test_symmetric_state_is_eigenstate(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """(|10⟩+|01⟩)/√2 is an eigenstate of H with eigenvalue 2."""
+        configs = torch.cat([C_10, C_01])
+        inv_sqrt2 = (0.5**0.5) + 0j
+        psi_i = torch.tensor([inv_sqrt2, inv_sqrt2], dtype=torch.complex128)
+        result = two_qubit_hamiltonian.apply_within(configs, psi_i, configs)
+        torch.testing.assert_close(result, 2 * psi_i)
+
+    def test_antisymmetric_state_has_zero_eigenvalue(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """(|10⟩−|01⟩)/√2 is an eigenstate of H with eigenvalue 0."""
+        configs = torch.cat([C_10, C_01])
+        inv_sqrt2 = (0.5**0.5) + 0j
+        psi_i = torch.tensor([inv_sqrt2, -inv_sqrt2], dtype=torch.complex128)
+        result = two_qubit_hamiltonian.apply_within(configs, psi_i, configs)
+        expected = torch.zeros(2, dtype=torch.complex128)
+        torch.testing.assert_close(result, expected)
+
+    def test_complex_amplitude(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """With imaginary input amplitude, the output is also imaginary."""
+        psi_i = torch.tensor([1.0j], dtype=torch.complex128)
+        result = two_qubit_hamiltonian.apply_within(C_10, psi_i, C_01)
+        expected = torch.tensor([1.0j], dtype=torch.complex128)
+        torch.testing.assert_close(result, expected)
+
+
+# ---------------------------------------------------------------------------
+# Tests for find_relative
+# ---------------------------------------------------------------------------
+
+
+class TestFindRelative:
+    """find_relative returns new configurations connected to the input by H (excluding configs_exclude)."""
+
+    def test_from_10_finds_01(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """From |10⟩ (excluding itself), the only off-diagonal connection is |01⟩."""
+        psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
+        result = two_qubit_hamiltonian.find_relative(C_10, psi_i, count_selected=5)
+        # result is a uint8 tensor of shape [n_found, n_qubytes]
+        assert result.shape[1] == 1
+        assert result.size(0) >= 1
+        found_values = result[:, 0].tolist()
+        assert 2 in found_values  # |01⟩ = uint8 2
+
+    def test_from_01_finds_10(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """From |01⟩ (excluding itself), the only off-diagonal connection is |10⟩."""
+        psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
+        result = two_qubit_hamiltonian.find_relative(C_01, psi_i, count_selected=5)
+        assert result.shape[1] == 1
+        assert result.size(0) >= 1
+        found_values = result[:, 0].tolist()
+        assert 1 in found_values  # |10⟩ = uint8 1
+
+    def test_from_11_no_off_diagonal(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """From |11⟩ hopping is Pauli-blocked; excluding itself, no connections exist."""
+        psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
+        result = two_qubit_hamiltonian.find_relative(C_11, psi_i, count_selected=5)
+        assert result.size(0) == 0
+
+    def test_custom_exclude(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """With a custom exclude that contains |01⟩, no off-diagonal result is returned."""
+        psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
+        # Exclude both |10⟩ and |01⟩ so that the only off-diagonal partner is gone
+        exclude = torch.cat([C_10, C_01])
+        result = two_qubit_hamiltonian.find_relative(C_10, psi_i, count_selected=5, configs_exclude=exclude)
+        # c†₁c₀|10⟩ → |01⟩ is excluded; no other off-diagonal terms apply
+        found_values = result[:, 0].tolist() if result.size(0) > 0 else []
+        assert 2 not in found_values
+
+
+# ---------------------------------------------------------------------------
+# Tests for list_relative
+# ---------------------------------------------------------------------------
+
+
+class TestListRelative:
+    """list_relative returns (configs_j, psi_j): all new configurations and their accumulated amplitudes."""
+
+    def test_from_10_finds_01_with_amplitude(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """c†₁c₀|10⟩ → |01⟩; excluding |10⟩ (default), |01⟩ is returned with amplitude 1."""
+        psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
+        configs_j, psi_j = two_qubit_hamiltonian.list_relative(C_10, psi_i)
+        assert configs_j.shape == (1, 1)
+        assert torch.equal(configs_j, C_01)
+        torch.testing.assert_close(psi_j, torch.tensor([1.0 + 0j], dtype=torch.complex128))
+
+    def test_from_01_finds_10_with_amplitude(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """c†₀c₁|01⟩ → |10⟩; excluding |01⟩ (default), |10⟩ is returned with amplitude 1."""
+        psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
+        configs_j, psi_j = two_qubit_hamiltonian.list_relative(C_01, psi_i)
+        assert configs_j.shape == (1, 1)
+        assert torch.equal(configs_j, C_10)
+        torch.testing.assert_close(psi_j, torch.tensor([1.0 + 0j], dtype=torch.complex128))
+
+    def test_from_11_empty_result(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """From |11⟩, hopping is Pauli-blocked; no off-diagonal connections exist."""
+        psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
+        configs_j, psi_j = two_qubit_hamiltonian.list_relative(C_11, psi_i)
+        assert configs_j.size(0) == 0
+        assert psi_j.size(0) == 0
+
+    def test_complex_amplitude_preserved(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """With imaginary input amplitude i, the output amplitude is also i."""
+        psi_i = torch.tensor([1.0j], dtype=torch.complex128)
+        configs_j, psi_j = two_qubit_hamiltonian.list_relative(C_10, psi_i)
+        assert torch.equal(configs_j, C_01)
+        torch.testing.assert_close(psi_j, torch.tensor([1.0j], dtype=torch.complex128))
+
+    def test_custom_exclude_reveals_diagonal_term(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """Excluding |01⟩ instead of |10⟩ lets the diagonal term n₀|10⟩=|10⟩ appear."""
+        psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
+        # exclude |01⟩: the off-diagonal c†₁c₀|10⟩→|01⟩ is excluded, but n₀|10⟩→|10⟩ is not
+        configs_j, psi_j = two_qubit_hamiltonian.list_relative(C_10, psi_i, configs_exclude=C_01)
+        assert torch.equal(configs_j, C_10)
+        torch.testing.assert_close(psi_j, torch.tensor([1.0 + 0j], dtype=torch.complex128))
+
+    def test_accumulated_amplitude_from_multiple_input_configs(self, two_qubit_hamiltonian: Hamiltonian) -> None:
+        """
+        Starting from both |10⟩ and |01⟩ with amplitudes a and b respectively,
+        n₀|10⟩→|10⟩ and c†₀c₁|01⟩→|10⟩ both contribute to the output |10⟩.
+        Similarly for |01⟩.  Excluding only |11⟩ so all 1-particle configs can appear.
+        """
+        configs_i = torch.cat([C_10, C_01])
+        psi_i = torch.tensor([2.0 + 0j, 3.0 + 0j], dtype=torch.complex128)
+        # Exclude |11⟩ so that both |10⟩ and |01⟩ can appear in the output
+        configs_j, psi_j = two_qubit_hamiltonian.list_relative(configs_i, psi_i, configs_exclude=C_11)
+
+        # Build expected: for each output config, sum contributions from all terms
+        # |10⟩ receives: n₀|10⟩ → coef=1, psi=2 → +2; c†₀c₁|01⟩ → coef=1, psi=3 → +3; total=5
+        # |01⟩ receives: c†₁c₀|10⟩ → coef=1, psi=2 → +2; n₁|01⟩ → coef=1, psi=3 → +3; total=5
+        expected_configs = {1: 5.0 + 0j, 2: 5.0 + 0j}
+
+        assert configs_j.size(0) == len(expected_configs)
+        for idx in range(configs_j.size(0)):
+            config_val = configs_j[idx, 0].item()
+            assert config_val in expected_configs
+            torch.testing.assert_close(psi_j[idx], torch.tensor(expected_configs[config_val], dtype=torch.complex128))

From 9e09079e14fb9640898b89eb690eb0020df96c55 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 5 Mar 2026 22:03:56 +0000
Subject: [PATCH 3/4] Switch test model to 4-site 1D spinless free fermion with
 distinct n/t coefficients

Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com>
---
 tests/test_hamiltonian.py | 402 +++++++++++++++++++-------------------
 1 file changed, 206 insertions(+), 196 deletions(-)

diff --git a/tests/test_hamiltonian.py b/tests/test_hamiltonian.py
index fede404..723da54 100644
--- a/tests/test_hamiltonian.py
+++ b/tests/test_hamiltonian.py
@@ -1,30 +1,35 @@
 """
 Tests for qmp.hamiltonian.Hamiltonian.
 
-We use a 2-qubit fermionic Hamiltonian (n_qubytes=1):
+We use a 4-site 1D spinless free-fermion Hamiltonian (n_qubytes=1):
 
-    H = c†₀c₁ + c†₁c₀ + n₀ + n₁
+    H = n + t
+      = Σᵢ nᵢ  +  2 · Σ_{⟨i,j⟩} (c†ᵢcⱼ + c†ⱼcᵢ)
 
-where c†ᵢ / cᵢ are creation / annihilation operators and nᵢ = c†ᵢcᵢ is the
-number operator at site i.
+where nᵢ = c†ᵢcᵢ is the number operator (coefficient 1) and the hopping is
+nearest-neighbor only with coefficient 2. The distinct coefficients make it
+easy to verify that the diagonal (n) and off-diagonal (t) parts are both
+computed correctly.
 
 Operator-key convention in the Hamiltonian dict:
     Key   ((s₀, k₀), (s₁, k₁), ...)  →  operator  O_{k₀,s₀} O_{k₁,s₁}
     with  k=0: annihilation,  k=1: creation.
-    The operators are applied right-to-left to the quantum state, i.e. the
-    rightmost operator in the expression acts on the state first.
-
-Occupation-number basis for 2 qubits (bit 0 = site 0, bit 1 = site 1):
-    |00⟩  →  uint8 value 0
-    |10⟩  →  uint8 value 1  (site 0 occupied)
-    |01⟩  →  uint8 value 2  (site 1 occupied)
-    |11⟩  →  uint8 value 3  (both sites occupied)
-
-In this basis, the Hamiltonian matrix is:
-    H = [[0, 0, 0, 0],
-         [0, 1, 1, 0],
-         [0, 1, 1, 0],
-         [0, 0, 0, 2]]
+    The operators are applied right-to-left, i.e. the last listed operator
+    acts on the state first.
+
+Config encoding (n_qubytes=1): bit k of the uint8 value = occupation of site k.
+    config = Σᵢ nᵢ · 2ⁱ  (bit 0 = site 0, …, bit 3 = site 3)
+
+Selected basis states used in tests:
+    0  = 0b0000  – vacuum
+    1  = 0b0001  – site 0 occupied
+    2  = 0b0010  – site 1 occupied
+    4  = 0b0100  – site 2 occupied
+    8  = 0b1000  – site 3 occupied
+    5  = 0b0101  – sites 0,2 occupied
+    6  = 0b0110  – sites 1,2 occupied
+    10 = 0b1010  – sites 1,3 occupied
+    15 = 0b1111  – all sites occupied
 """
 
 import pytest
@@ -33,38 +38,53 @@
 from qmp.hamiltonian import Hamiltonian
 
 # ---------------------------------------------------------------------------
-# Fixtures
+# Fixture
 # ---------------------------------------------------------------------------
 
 
 @pytest.fixture(scope="session")
-def two_qubit_hamiltonian() -> Hamiltonian:
+def hamiltonian() -> Hamiltonian:
     """
-    2-qubit fermionic Hamiltonian: H = c†₀c₁ + c†₁c₀ + n₀ + n₁.
+    4-site 1D spinless free-fermion Hamiltonian: H = Σ nᵢ + 2·Σ_{⟨i,j⟩}(c†ᵢcⱼ + h.c.)
 
+    Number-operator coefficient = 1, nearest-neighbor hopping coefficient = 2.
     Compiled and cached once per test session to avoid repeated JIT overhead.
     """
     return Hamiltonian(
         {
-            ((0, 1), (1, 0)): 1.0,  # c†₀c₁
-            ((1, 1), (0, 0)): 1.0,  # c†₁c₀
-            ((0, 1), (0, 0)): 1.0,  # n₀ = c†₀c₀
-            ((1, 1), (1, 0)): 1.0,  # n₁ = c†₁c₁
+            # Number operators (coefficient 1)
+            ((0, 1), (0, 0)): 1.0,  # n₀
+            ((1, 1), (1, 0)): 1.0,  # n₁
+            ((2, 1), (2, 0)): 1.0,  # n₂
+            ((3, 1), (3, 0)): 1.0,  # n₃
+            # Nearest-neighbor hopping (coefficient 2)
+            ((0, 1), (1, 0)): 2.0,  # c†₀c₁
+            ((1, 1), (0, 0)): 2.0,  # c†₁c₀
+            ((1, 1), (2, 0)): 2.0,  # c†₁c₂
+            ((2, 1), (1, 0)): 2.0,  # c†₂c₁
+            ((2, 1), (3, 0)): 2.0,  # c†₂c₃
+            ((3, 1), (2, 0)): 2.0,  # c†₃c₂
         },
         kind="fermi",
     )
 
 
 # ---------------------------------------------------------------------------
-# Helpers – configs for 2-qubit system (n_qubytes = 1)
+# Config tensors (shape [batch, n_qubytes=1])
 # ---------------------------------------------------------------------------
 
-C_00 = torch.tensor([[0]], dtype=torch.uint8)  # |00⟩
-C_10 = torch.tensor([[1]], dtype=torch.uint8)  # |10⟩  (site 0 occupied)
-C_01 = torch.tensor([[2]], dtype=torch.uint8)  # |01⟩  (site 1 occupied)
-C_11 = torch.tensor([[3]], dtype=torch.uint8)  # |11⟩  (both sites occupied)
+C_vacuum = torch.tensor([[0]], dtype=torch.uint8)   # |0000⟩
+C_site0 = torch.tensor([[1]], dtype=torch.uint8)    # |1000⟩ – site 0 occupied
+C_site1 = torch.tensor([[2]], dtype=torch.uint8)    # |0100⟩ – site 1 occupied
+C_site2 = torch.tensor([[4]], dtype=torch.uint8)    # |0010⟩ – site 2 occupied
+C_site3 = torch.tensor([[8]], dtype=torch.uint8)    # |0001⟩ – site 3 occupied
+C_sites02 = torch.tensor([[5]], dtype=torch.uint8)  # |1010⟩ – sites 0,2 occupied
+C_sites12 = torch.tensor([[6]], dtype=torch.uint8)  # |0110⟩ – sites 1,2 occupied
+C_sites13 = torch.tensor([[10]], dtype=torch.uint8) # |0101⟩ – sites 1,3 occupied
+C_all = torch.tensor([[15]], dtype=torch.uint8)     # |1111⟩ – all sites occupied
 
-ALL_CONFIGS = torch.tensor([[0], [1], [2], [3]], dtype=torch.uint8)
+# All 1-particle configs (sites 0–3)
+C_1particle = torch.tensor([[1], [2], [4], [8]], dtype=torch.uint8)
 
 
 # ---------------------------------------------------------------------------
@@ -73,32 +93,32 @@ def two_qubit_hamiltonian() -> Hamiltonian:
 
 
 class TestDiagonalTerm:
-    """diagonal_term(configs) returns the diagonal element H[i,i] for each config."""
-
-    def test_vacuum_diagonal_is_zero(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        result = two_qubit_hamiltonian.diagonal_term(C_00)
-        expected = torch.tensor([0.0 + 0j], dtype=torch.complex128)
-        torch.testing.assert_close(result, expected)
-
-    def test_single_particle_site0(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        result = two_qubit_hamiltonian.diagonal_term(C_10)
-        expected = torch.tensor([1.0 + 0j], dtype=torch.complex128)
-        torch.testing.assert_close(result, expected)
-
-    def test_single_particle_site1(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        result = two_qubit_hamiltonian.diagonal_term(C_01)
-        expected = torch.tensor([1.0 + 0j], dtype=torch.complex128)
-        torch.testing.assert_close(result, expected)
-
-    def test_two_particle_sector(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        result = two_qubit_hamiltonian.diagonal_term(C_11)
-        expected = torch.tensor([2.0 + 0j], dtype=torch.complex128)
-        torch.testing.assert_close(result, expected)
-
-    def test_all_configs_batch(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """Batched evaluation matches individual results."""
-        result = two_qubit_hamiltonian.diagonal_term(ALL_CONFIGS)
-        expected = torch.tensor([0.0 + 0j, 1.0 + 0j, 1.0 + 0j, 2.0 + 0j], dtype=torch.complex128)
+    """diagonal_term returns the diagonal matrix element H[config, config]."""
+
+    def test_vacuum_is_zero(self, hamiltonian: Hamiltonian) -> None:
+        """Vacuum has no particles; all number operators vanish."""
+        result = hamiltonian.diagonal_term(C_vacuum)
+        torch.testing.assert_close(result, torch.tensor([0.0 + 0j], dtype=torch.complex128))
+
+    def test_single_particle_equals_n_coefficient(self, hamiltonian: Hamiltonian) -> None:
+        """Single-particle states give diagonal = 1 (the number-operator coefficient)."""
+        result = hamiltonian.diagonal_term(C_site0)
+        torch.testing.assert_close(result, torch.tensor([1.0 + 0j], dtype=torch.complex128))
+
+    def test_two_particles_sums_n_coefficients(self, hamiltonian: Hamiltonian) -> None:
+        """Two non-adjacent occupied sites: diagonal = 2·1 = 2 (hopping doesn't contribute)."""
+        result = hamiltonian.diagonal_term(C_sites02)
+        torch.testing.assert_close(result, torch.tensor([2.0 + 0j], dtype=torch.complex128))
+
+    def test_all_occupied_diagonal_is_four(self, hamiltonian: Hamiltonian) -> None:
+        """All 4 sites occupied: diagonal = 4·1 = 4 (hopping blocked by Pauli exclusion)."""
+        result = hamiltonian.diagonal_term(C_all)
+        torch.testing.assert_close(result, torch.tensor([4.0 + 0j], dtype=torch.complex128))
+
+    def test_batch_all_single_particle_states(self, hamiltonian: Hamiltonian) -> None:
+        """All four 1-particle basis states give diagonal = 1 in a single batched call."""
+        result = hamiltonian.diagonal_term(C_1particle)
+        expected = torch.ones(4, dtype=torch.complex128)
         torch.testing.assert_close(result, expected)
 
 
@@ -110,72 +130,66 @@ def test_all_configs_batch(self, two_qubit_hamiltonian: Hamiltonian) -> None:
 class TestApplyWithin:
     """apply_within(configs_i, psi_i, configs_j) computes (H·ψ) restricted to configs_j."""
 
-    def test_apply_to_basis_state_10(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """H|10⟩ = |10⟩ + |01⟩ in the 1-particle sector."""
-        configs = torch.cat([C_10, C_01])
-        psi_i = torch.tensor([1.0 + 0j, 0.0 + 0j], dtype=torch.complex128)
-        result = two_qubit_hamiltonian.apply_within(configs, psi_i, configs)
-        expected = torch.tensor([1.0 + 0j, 1.0 + 0j], dtype=torch.complex128)
-        torch.testing.assert_close(result, expected)
-
-    def test_apply_to_basis_state_01(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """H|01⟩ = |10⟩ + |01⟩ in the 1-particle sector."""
-        configs = torch.cat([C_10, C_01])
-        psi_i = torch.tensor([0.0 + 0j, 1.0 + 0j], dtype=torch.complex128)
-        result = two_qubit_hamiltonian.apply_within(configs, psi_i, configs)
-        expected = torch.tensor([1.0 + 0j, 1.0 + 0j], dtype=torch.complex128)
-        torch.testing.assert_close(result, expected)
+    def test_n_and_t_coefficients_are_distinguishable(self, hamiltonian: Hamiltonian) -> None:
+        """H|site1⟩ in the 1-particle sector gives [2, 1, 2] for configs [site0, site1, site2].
 
-    def test_apply_to_two_particle_state(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """H|11⟩ = 2|11⟩ (both number operators contribute; hopping is blocked)."""
+        The diagonal element is 1 (n coefficient), off-diagonal hops are 2 (t coefficient).
+        """
+        configs_j = torch.tensor([[1], [2], [4]], dtype=torch.uint8)  # sites 0, 1, 2
         psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
-        result = two_qubit_hamiltonian.apply_within(C_11, psi_i, C_11)
-        expected = torch.tensor([2.0 + 0j], dtype=torch.complex128)
-        torch.testing.assert_close(result, expected)
-
-    def test_no_transition_to_vacuum(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """H maps no configuration to |00⟩."""
-        psi_i = torch.ones(4, dtype=torch.complex128)
-        result = two_qubit_hamiltonian.apply_within(ALL_CONFIGS, psi_i, C_00)
-        expected = torch.tensor([0.0 + 0j], dtype=torch.complex128)
+        result = hamiltonian.apply_within(C_site1, psi_i, configs_j)
+        expected = torch.tensor([2.0 + 0j, 1.0 + 0j, 2.0 + 0j], dtype=torch.complex128)
         torch.testing.assert_close(result, expected)
 
-    def test_off_diagonal_10_to_01(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """c†₁c₀|10⟩ → |01⟩ with amplitude 1."""
+    def test_diagonal_only_from_all_occupied(self, hamiltonian: Hamiltonian) -> None:
+        """H|1111⟩ = 4·|1111⟩: hopping is Pauli-blocked, only n terms contribute."""
         psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
-        result = two_qubit_hamiltonian.apply_within(C_10, psi_i, C_01)
-        expected = torch.tensor([1.0 + 0j], dtype=torch.complex128)
-        torch.testing.assert_close(result, expected)
-
-    def test_off_diagonal_01_to_10(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """c†₀c₁|01⟩ → |10⟩ with amplitude 1."""
+        result = hamiltonian.apply_within(C_all, psi_i, C_all)
+        torch.testing.assert_close(result, torch.tensor([4.0 + 0j], dtype=torch.complex128))
+
+    def test_no_transition_to_vacuum(self, hamiltonian: Hamiltonian) -> None:
+        """H never maps any configuration to the vacuum |0000⟩."""
+        configs_i = torch.tensor([[0], [1], [2], [4], [8], [5], [6], [15]], dtype=torch.uint8)
+        psi_i = torch.ones(8, dtype=torch.complex128)
+        result = hamiltonian.apply_within(configs_i, psi_i, C_vacuum)
+        torch.testing.assert_close(result, torch.tensor([0.0 + 0j], dtype=torch.complex128))
+
+    def test_hopping_only_between_adjacent_sites(self, hamiltonian: Hamiltonian) -> None:
+        """c†₂c₃ connects site3 to site2 (adjacent); site0 is not connected to site3."""
         psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
-        result = two_qubit_hamiltonian.apply_within(C_01, psi_i, C_10)
-        expected = torch.tensor([1.0 + 0j], dtype=torch.complex128)
-        torch.testing.assert_close(result, expected)
-
-    def test_symmetric_state_is_eigenstate(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """(|10⟩+|01⟩)/√2 is an eigenstate of H with eigenvalue 2."""
-        configs = torch.cat([C_10, C_01])
-        inv_sqrt2 = (0.5**0.5) + 0j
-        psi_i = torch.tensor([inv_sqrt2, inv_sqrt2], dtype=torch.complex128)
-        result = two_qubit_hamiltonian.apply_within(configs, psi_i, configs)
-        torch.testing.assert_close(result, 2 * psi_i)
-
-    def test_antisymmetric_state_has_zero_eigenvalue(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """(|10⟩−|01⟩)/√2 is an eigenstate of H with eigenvalue 0."""
-        configs = torch.cat([C_10, C_01])
-        inv_sqrt2 = (0.5**0.5) + 0j
-        psi_i = torch.tensor([inv_sqrt2, -inv_sqrt2], dtype=torch.complex128)
-        result = two_qubit_hamiltonian.apply_within(configs, psi_i, configs)
-        expected = torch.zeros(2, dtype=torch.complex128)
-        torch.testing.assert_close(result, expected)
+        # site0 should receive 0 (non-adjacent to site3)
+        result_non_adjacent = hamiltonian.apply_within(C_site3, psi_i, C_site0)
+        torch.testing.assert_close(result_non_adjacent, torch.tensor([0.0 + 0j], dtype=torch.complex128))
+        # site2 should receive 2 (adjacent to site3, hopping coefficient)
+        result_adjacent = hamiltonian.apply_within(C_site3, psi_i, C_site2)
+        torch.testing.assert_close(result_adjacent, torch.tensor([2.0 + 0j], dtype=torch.complex128))
+
+    def test_hopping_amplitude_in_two_particle_sector(self, hamiltonian: Hamiltonian) -> None:
+        """c†₀c₁ connects sites12 (config 6) to sites02 (config 5) with amplitude 2."""
+        psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
+        result = hamiltonian.apply_within(C_sites12, psi_i, C_sites02)
+        torch.testing.assert_close(result, torch.tensor([2.0 + 0j], dtype=torch.complex128))
 
-    def test_complex_amplitude(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """With imaginary input amplitude, the output is also imaginary."""
+    def test_complex_amplitude_preserved(self, hamiltonian: Hamiltonian) -> None:
+        """With imaginary input amplitude, the output amplitude is imaginary too."""
         psi_i = torch.tensor([1.0j], dtype=torch.complex128)
-        result = two_qubit_hamiltonian.apply_within(C_10, psi_i, C_01)
-        expected = torch.tensor([1.0j], dtype=torch.complex128)
+        result = hamiltonian.apply_within(C_site0, psi_i, C_site1)
+        torch.testing.assert_close(result, torch.tensor([2.0j], dtype=torch.complex128))
+
+    def test_full_1particle_sector_matrix(self, hamiltonian: Hamiltonian) -> None:
+        """Apply H to the full 1-particle sector and verify the tridiagonal matrix.
+
+        In the 1-particle sector the Hamiltonian matrix is tridiagonal:
+            [[1, 2, 0, 0],
+             [2, 1, 2, 0],
+             [0, 2, 1, 2],
+             [0, 0, 2, 1]]
+        where 1 = n coefficient, 2 = t coefficient.
+        """
+        psi_i = torch.tensor([1.0 + 0j, 1.0 + 0j, 1.0 + 0j, 1.0 + 0j], dtype=torch.complex128)
+        result = hamiltonian.apply_within(C_1particle, psi_i, C_1particle)
+        # Row sums of the tridiagonal matrix above
+        expected = torch.tensor([3.0 + 0j, 5.0 + 0j, 5.0 + 0j, 3.0 + 0j], dtype=torch.complex128)
         torch.testing.assert_close(result, expected)
 
 
@@ -185,42 +199,37 @@ def test_complex_amplitude(self, two_qubit_hamiltonian: Hamiltonian) -> None:
 
 
 class TestFindRelative:
-    """find_relative returns new configurations connected to the input by H (excluding configs_exclude)."""
+    """find_relative(configs_i, psi_i, count_selected) returns at most count_selected
+    new configurations with the highest estimated weights, excluding configs_exclude."""
 
-    def test_from_10_finds_01(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """From |10⟩ (excluding itself), the only off-diagonal connection is |01⟩."""
+    def test_finds_both_neighbours_of_site1(self, hamiltonian: Hamiltonian) -> None:
+        """From site1 (config 2), connected configs are site0 and site2."""
         psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
-        result = two_qubit_hamiltonian.find_relative(C_10, psi_i, count_selected=5)
-        # result is a uint8 tensor of shape [n_found, n_qubytes]
-        assert result.shape[1] == 1
-        assert result.size(0) >= 1
-        found_values = result[:, 0].tolist()
-        assert 2 in found_values  # |01⟩ = uint8 2
-
-    def test_from_01_finds_10(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """From |01⟩ (excluding itself), the only off-diagonal connection is |10⟩."""
+        result = hamiltonian.find_relative(C_site1, psi_i, count_selected=10)
+        found = set(result[:, 0].tolist())
+        assert found == {1, 4}  # site0 and site2
+
+    def test_count_selected_limits_results(self, hamiltonian: Hamiltonian) -> None:
+        """count_selected=1 returns at most 1 config even when 2 neighbours exist."""
         psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
-        result = two_qubit_hamiltonian.find_relative(C_01, psi_i, count_selected=5)
-        assert result.shape[1] == 1
-        assert result.size(0) >= 1
-        found_values = result[:, 0].tolist()
-        assert 1 in found_values  # |10⟩ = uint8 1
-
-    def test_from_11_no_off_diagonal(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """From |11⟩ hopping is Pauli-blocked; excluding itself, no connections exist."""
+        result_limited = hamiltonian.find_relative(C_site1, psi_i, count_selected=1)
+        result_full = hamiltonian.find_relative(C_site1, psi_i, count_selected=10)
+        assert result_limited.size(0) == 1
+        assert result_full.size(0) == 2
+
+    def test_no_off_diagonal_from_all_occupied(self, hamiltonian: Hamiltonian) -> None:
+        """From |1111⟩, hopping is Pauli-blocked; no off-diagonal connections exist."""
         psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
-        result = two_qubit_hamiltonian.find_relative(C_11, psi_i, count_selected=5)
+        result = hamiltonian.find_relative(C_all, psi_i, count_selected=10)
         assert result.size(0) == 0
 
-    def test_custom_exclude(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """With a custom exclude that contains |01⟩, no off-diagonal result is returned."""
+    def test_custom_exclude_removes_neighbour(self, hamiltonian: Hamiltonian) -> None:
+        """Excluding site2 (config 4) from search, only site0 is returned from site1."""
         psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
-        # Exclude both |10⟩ and |01⟩ so that the only off-diagonal partner is gone
-        exclude = torch.cat([C_10, C_01])
-        result = two_qubit_hamiltonian.find_relative(C_10, psi_i, count_selected=5, configs_exclude=exclude)
-        # c†₁c₀|10⟩ → |01⟩ is excluded; no other off-diagonal terms apply
-        found_values = result[:, 0].tolist() if result.size(0) > 0 else []
-        assert 2 not in found_values
+        exclude = torch.cat([C_site1, C_site2])  # exclude self and site2
+        result = hamiltonian.find_relative(C_site1, psi_i, count_selected=10, configs_exclude=exclude)
+        found = set(result[:, 0].tolist())
+        assert found == {1}  # only site0
 
 
 # ---------------------------------------------------------------------------
@@ -229,64 +238,65 @@ def test_custom_exclude(self, two_qubit_hamiltonian: Hamiltonian) -> None:
 
 
 class TestListRelative:
-    """list_relative returns (configs_j, psi_j): all new configurations and their accumulated amplitudes."""
+    """list_relative returns ALL unique new configurations and their accumulated amplitudes,
+    unlike find_relative which returns at most count_selected."""
 
-    def test_from_10_finds_01_with_amplitude(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """c†₁c₀|10⟩ → |01⟩; excluding |10⟩ (default), |01⟩ is returned with amplitude 1."""
+    def test_lists_both_neighbours_of_site1(self, hamiltonian: Hamiltonian) -> None:
+        """From site1, list_relative finds BOTH site0 and site2 with amplitude 2 each."""
         psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
-        configs_j, psi_j = two_qubit_hamiltonian.list_relative(C_10, psi_i)
-        assert configs_j.shape == (1, 1)
-        assert torch.equal(configs_j, C_01)
-        torch.testing.assert_close(psi_j, torch.tensor([1.0 + 0j], dtype=torch.complex128))
-
-    def test_from_01_finds_10_with_amplitude(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """c†₀c₁|01⟩ → |10⟩; excluding |01⟩ (default), |10⟩ is returned with amplitude 1."""
+        configs_j, psi_j = hamiltonian.list_relative(C_site1, psi_i)
+        assert configs_j.size(0) == 2
+        result = {configs_j[i, 0].item(): psi_j[i] for i in range(configs_j.size(0))}
+        assert set(result.keys()) == {1, 4}  # site0 and site2
+        for amp in result.values():
+            torch.testing.assert_close(amp, torch.tensor(2.0 + 0j, dtype=torch.complex128))
+
+    def test_lists_all_connections_from_two_particle_state(self, hamiltonian: Hamiltonian) -> None:
+        """From sites12 (config 6), both off-diagonal connections carry amplitude 2."""
         psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
-        configs_j, psi_j = two_qubit_hamiltonian.list_relative(C_01, psi_i)
-        assert configs_j.shape == (1, 1)
-        assert torch.equal(configs_j, C_10)
-        torch.testing.assert_close(psi_j, torch.tensor([1.0 + 0j], dtype=torch.complex128))
-
-    def test_from_11_empty_result(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """From |11⟩, hopping is Pauli-blocked; no off-diagonal connections exist."""
+        configs_j, psi_j = hamiltonian.list_relative(C_sites12, psi_i)
+        assert configs_j.size(0) == 2
+        result = {configs_j[i, 0].item(): psi_j[i] for i in range(configs_j.size(0))}
+        assert set(result.keys()) == {5, 10}  # sites02 and sites13
+        for amp in result.values():
+            torch.testing.assert_close(amp, torch.tensor(2.0 + 0j, dtype=torch.complex128))
+
+    def test_empty_result_for_fully_occupied(self, hamiltonian: Hamiltonian) -> None:
+        """From |1111⟩, hopping is Pauli-blocked; list_relative returns empty tensors."""
         psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
-        configs_j, psi_j = two_qubit_hamiltonian.list_relative(C_11, psi_i)
+        configs_j, psi_j = hamiltonian.list_relative(C_all, psi_i)
         assert configs_j.size(0) == 0
         assert psi_j.size(0) == 0
 
-    def test_complex_amplitude_preserved(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """With imaginary input amplitude i, the output amplitude is also i."""
+    def test_complex_amplitude_preserved(self, hamiltonian: Hamiltonian) -> None:
+        """With imaginary input amplitude, the output amplitudes are imaginary."""
         psi_i = torch.tensor([1.0j], dtype=torch.complex128)
-        configs_j, psi_j = two_qubit_hamiltonian.list_relative(C_10, psi_i)
-        assert torch.equal(configs_j, C_01)
-        torch.testing.assert_close(psi_j, torch.tensor([1.0j], dtype=torch.complex128))
-
-    def test_custom_exclude_reveals_diagonal_term(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """Excluding |01⟩ instead of |10⟩ lets the diagonal term n₀|10⟩=|10⟩ appear."""
-        psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
-        # exclude |01⟩: the off-diagonal c†₁c₀|10⟩→|01⟩ is excluded, but n₀|10⟩→|10⟩ is not
-        configs_j, psi_j = two_qubit_hamiltonian.list_relative(C_10, psi_i, configs_exclude=C_01)
-        assert torch.equal(configs_j, C_10)
-        torch.testing.assert_close(psi_j, torch.tensor([1.0 + 0j], dtype=torch.complex128))
-
-    def test_accumulated_amplitude_from_multiple_input_configs(self, two_qubit_hamiltonian: Hamiltonian) -> None:
-        """
-        Starting from both |10⟩ and |01⟩ with amplitudes a and b respectively,
-        n₀|10⟩→|10⟩ and c†₀c₁|01⟩→|10⟩ both contribute to the output |10⟩.
-        Similarly for |01⟩.  Excluding only |11⟩ so all 1-particle configs can appear.
+        configs_j, psi_j = hamiltonian.list_relative(C_site1, psi_i)
+        assert configs_j.size(0) == 2
+        result = {configs_j[i, 0].item(): psi_j[i] for i in range(configs_j.size(0))}
+        for amp in result.values():
+            torch.testing.assert_close(amp, torch.tensor(2.0j, dtype=torch.complex128))
+
+    def test_amplitudes_accumulate_from_multiple_input_configs(self, hamiltonian: Hamiltonian) -> None:
+        """Starting from site0 and site2 (psi=[1,1]), list_relative accumulates contributions:
+        - site0 → (c†₁c₀) → site1 with amplitude 2
+        - site2 → (c†₁c₂) → site1 with amplitude 2  (accumulated: site1 total = 4)
+        - site2 → (c†₃c₂) → site3 with amplitude 2  (site3 total = 2)
         """
-        configs_i = torch.cat([C_10, C_01])
-        psi_i = torch.tensor([2.0 + 0j, 3.0 + 0j], dtype=torch.complex128)
-        # Exclude |11⟩ so that both |10⟩ and |01⟩ can appear in the output
-        configs_j, psi_j = two_qubit_hamiltonian.list_relative(configs_i, psi_i, configs_exclude=C_11)
-
-        # Build expected: for each output config, sum contributions from all terms
-        # |10⟩ receives: n₀|10⟩ → coef=1, psi=2 → +2; c†₀c₁|01⟩ → coef=1, psi=3 → +3; total=5
-        # |01⟩ receives: c†₁c₀|10⟩ → coef=1, psi=2 → +2; n₁|01⟩ → coef=1, psi=3 → +3; total=5
-        expected_configs = {1: 5.0 + 0j, 2: 5.0 + 0j}
-
-        assert configs_j.size(0) == len(expected_configs)
-        for idx in range(configs_j.size(0)):
-            config_val = configs_j[idx, 0].item()
-            assert config_val in expected_configs
-            torch.testing.assert_close(psi_j[idx], torch.tensor(expected_configs[config_val], dtype=torch.complex128))
+        configs_i = torch.cat([C_site0, C_site2])
+        psi_i = torch.tensor([1.0 + 0j, 1.0 + 0j], dtype=torch.complex128)
+        configs_j, psi_j = hamiltonian.list_relative(configs_i, psi_i)
+
+        assert configs_j.size(0) == 2
+        result = {configs_j[i, 0].item(): psi_j[i] for i in range(configs_j.size(0))}
+        assert set(result.keys()) == {2, 8}  # site1 and site3
+        torch.testing.assert_close(result[2], torch.tensor(4.0 + 0j, dtype=torch.complex128))
+        torch.testing.assert_close(result[8], torch.tensor(2.0 + 0j, dtype=torch.complex128))
+
+    def test_custom_exclude_removes_neighbour(self, hamiltonian: Hamiltonian) -> None:
+        """Excluding site2 (config 4), list_relative returns only site0 from site1."""
+        psi_i = torch.tensor([1.0 + 0j], dtype=torch.complex128)
+        exclude = torch.cat([C_site1, C_site2])  # exclude self and site2
+        configs_j, psi_j = hamiltonian.list_relative(C_site1, psi_i, configs_exclude=exclude)
+        assert configs_j.size(0) == 1
+        assert configs_j[0, 0].item() == 1  # only site0

From 3f4c74bb68f91e1c3ab0d0109b1d311dc615e094 Mon Sep 17 00:00:00 2001
From: Hao Zhang <hzhangxyz@outlook.com>
Date: Fri, 6 Mar 2026 06:21:18 +0800
Subject: [PATCH 4/4] Format and rename the test.

---
 tests/{test_hamiltonian.py => hamiltonian_test.py} | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)
 rename tests/{test_hamiltonian.py => hamiltonian_test.py} (96%)

diff --git a/tests/test_hamiltonian.py b/tests/hamiltonian_test.py
similarity index 96%
rename from tests/test_hamiltonian.py
rename to tests/hamiltonian_test.py
index 723da54..89fd3e4 100644
--- a/tests/test_hamiltonian.py
+++ b/tests/hamiltonian_test.py
@@ -73,15 +73,15 @@ def hamiltonian() -> Hamiltonian:
 # Config tensors (shape [batch, n_qubytes=1])
 # ---------------------------------------------------------------------------
 
-C_vacuum = torch.tensor([[0]], dtype=torch.uint8)   # |0000⟩
-C_site0 = torch.tensor([[1]], dtype=torch.uint8)    # |1000⟩ – site 0 occupied
-C_site1 = torch.tensor([[2]], dtype=torch.uint8)    # |0100⟩ – site 1 occupied
-C_site2 = torch.tensor([[4]], dtype=torch.uint8)    # |0010⟩ – site 2 occupied
-C_site3 = torch.tensor([[8]], dtype=torch.uint8)    # |0001⟩ – site 3 occupied
+C_vacuum = torch.tensor([[0]], dtype=torch.uint8)  # |0000⟩
+C_site0 = torch.tensor([[1]], dtype=torch.uint8)  # |1000⟩ – site 0 occupied
+C_site1 = torch.tensor([[2]], dtype=torch.uint8)  # |0100⟩ – site 1 occupied
+C_site2 = torch.tensor([[4]], dtype=torch.uint8)  # |0010⟩ – site 2 occupied
+C_site3 = torch.tensor([[8]], dtype=torch.uint8)  # |0001⟩ – site 3 occupied
 C_sites02 = torch.tensor([[5]], dtype=torch.uint8)  # |1010⟩ – sites 0,2 occupied
 C_sites12 = torch.tensor([[6]], dtype=torch.uint8)  # |0110⟩ – sites 1,2 occupied
-C_sites13 = torch.tensor([[10]], dtype=torch.uint8) # |0101⟩ – sites 1,3 occupied
-C_all = torch.tensor([[15]], dtype=torch.uint8)     # |1111⟩ – all sites occupied
+C_sites13 = torch.tensor([[10]], dtype=torch.uint8)  # |0101⟩ – sites 1,3 occupied
+C_all = torch.tensor([[15]], dtype=torch.uint8)  # |1111⟩ – all sites occupied
 
 # All 1-particle configs (sites 0–3)
 C_1particle = torch.tensor([[1], [2], [4], [8]], dtype=torch.uint8)