sanity check all pickled non-interacting Hamiltonians in h0 directory

impurityModel/ed/get_spectra.py

@@ -396,7 +396,7 @@ def get_hamiltonian_operator(nBaths, nValBaths, slaterCondon, SOCs, DCinfo, hFie
     hHfieldOperator = finite.gethHfieldop(hx, hy, hz, l=2)
 
     # Read the non-relativistic non-interacting Hamiltonian operator from file.
-    h0_operator = get_h0_operator(h0_filename, nBaths)
+    h0_operator = read_pickled_file(h0_filename)
 
     # Add Hamiltonian terms to one operator.
     hOperator = finite.addOps(
@@ -416,37 +416,10 @@ def get_hamiltonian_operator(nBaths, nValBaths, slaterCondon, SOCs, DCinfo, hFie
     return hOp
 
 
-def get_h0_operator(h0_filename, nBaths):
-    """
-    Return h0 operator.
-
-    Parameters
-    ----------
-    h0_filename : str
-        Filename of non-interacting, non-relativistic operator.
-    nBaths : dict
-        Number of bath states for each angular momentum.
-
-    Returns
-    -------
-    h0_operator : dict
-        The non-relativistic non-interacting Hamiltonian in operator form.
-        Hamiltonian describes 3d orbitals and bath orbitals.
-        tuple : complex,
-        where each tuple describes a process of two steps (annihilation and then creation).
-        Each step is described by a tuple of the form:
-        (spin_orb, 'c') or (spin_orb, 'a'),
-        where spin_orb is a tuple of the form (l, s, m) or (l, b) or ((l_a, l_b), b).
-
-    """
-    with open(h0_filename, "rb") as handle:
-        h0_operator = pickle.loads(handle.read())
-    # Sanity check
-    for process in h0_operator.keys():
-        for event in process:
-            if len(event[0]) == 2:
-                assert nBaths[event[0][0]] > event[0][1]
-    return h0_operator
+def read_pickled_file(filename: str):
+    with open(filename, "rb") as handle:
+        content = pickle.load(handle)
+    return content
 
 
 if __name__ == "__main__":

impurityModel/test/test_h0.py

@@ -0,0 +1,58 @@
+import os
+from glob import glob
+
+from impurityModel.ed.get_spectra import read_pickled_file
+
+DIR_PATH = os.path.dirname(os.path.realpath(__file__))
+
+
+def test_read_all_h0_pickle_files():
+    for h0_filename in glob(os.path.join(DIR_PATH, "../../h0/h0*.pickle")):
+        h0 = read_pickled_file(h0_filename)
+
+        string = os.path.basename(h0_filename).split(".")[0].split("_")[-1].split("bath")[0]
+        nBaths_tot = sum([int(nbath) for nbath in string.split("p")])
+        # So far, in all the non-interacting Hamiltonians the angular momentum is equal to two
+        # for the correlated orbitals that are coupled to the bath states.
+        nBaths = {2: nBaths_tot}
+
+        sanity_check_non_interacting_hamiltonian(h0, nBaths)
+
+
+def sanity_check_non_interacting_hamiltonian(h0: dict[tuple, complex], nBaths: dict[int, int]):
+    """
+    Sanity check non-interacting Hamiltonian operator.
+
+    Parameters
+    -------
+    h0:
+        Non-interacting Hamiltonian operator.
+        Describes impurity orbitals and bath orbitals.
+        Each tuple key describes a process of two steps: annihilation followed by creation.
+        Each step is described by a tuple of the form:
+        (spin_orb, 'c') for creation, and
+        (spin_orb, 'a') for annihilation,
+        where spin_orb is a tuple of the form (l, s, m) or (l, b) or ((l_a, l_b), b).
+    nBaths:
+        angular momentum: number of bath states coupled to the correlated orbitals with this angular momentum.
+    """
+    for process, value in h0.items():
+        assert isinstance(value, complex)
+        assert len(process) == 2  # two events in non-interacting Hamiltonian
+        assert process[1][1] == "a"  # First event is annihilation
+        assert process[0][1] == "c"  # Second event is creation
+        for event in process[::-1]:
+            assert len(event) == 2  # spin-orbit and create or remove
+            spin_orbit_info = event[0]
+            if len(spin_orbit_info) == 2:
+                l, bath_index = spin_orbit_info
+                assert 0 <= l <= 3
+                assert 0 <= bath_index < nBaths[l]
+            elif len(spin_orbit_info) == 3:
+                l, s, m = spin_orbit_info
+                assert l in nBaths
+                assert 0 <= l <= 3
+                assert s in (0, 1)
+                assert m in range(-l, l + 1)
+            else:
+                raise ValueError(f"{spin_orbit_info}")