Introduce type hints in the HF backend files

.github/workflows/ci.yaml

@@ -161,11 +161,8 @@ jobs:
           python-version: ${{ matrix.version }}
           auto-activate-base: false
       - name: Install Psi4 from c-f
-        # TODO: remove libxc-c from the install list once
-        # https://github.com/psi4/psi4/issues/3474
-        # is fixed
         run: |
-          mamba install psi4 pyddx "libxc-c=7.0.0" -c conda-forge -c conda-forge/label/libint_dev
+          mamba install psi4 pyddx -c conda-forge -c conda-forge/label/libint_dev
           psi4 --version
           python -c "import psi4"
       - name: Install system dependencies on macOS

adcc/backends/EriBuilder.py

@@ -20,22 +20,47 @@
 ## along with adcc. If not, see <http://www.gnu.org/licenses/>.
 ##
 ## ---------------------------------------------------------------------
+from dataclasses import dataclass
 from itertools import product
-from collections import namedtuple
+from typing import Literal, TypeAlias, TypeGuard
+import numpy as np
 
+IntSlice: TypeAlias = "slice[int, int, int]"
+IntSlice4D = tuple[IntSlice, IntSlice, IntSlice, IntSlice]
+Block = Literal["O", "V"]
+Block4D = tuple[Block, Block, Block, Block]
+Spin = Literal["a", "b"]
+Spin4D = tuple[Spin, Spin, Spin, Spin]
+Array4D = np.ndarray[tuple[int, int, int, int], np.dtype[np.float64]]
 
-def range_in(inner, full):
-    if inner.start is None:
-        inner = slice(0, inner.stop, 1)
-    if full.start is None:
-        full = slice(0, full.stop, 1)
-    return all(r in range(full.start, full.stop)
-               for r in range(inner.start, inner.stop))
+
+def is_int_slice(sl: slice) -> TypeGuard[IntSlice]:
+    return (
+        isinstance(sl.start, int)
+        and isinstance(sl.stop, int)
+        and isinstance(sl.step, int)
+    )
+
+
+@dataclass(frozen=True, slots=True)
+class SpinBlockSlice:
+    block: Block
+    spin: Spin
+    fromslice: IntSlice
+    toslice: IntSlice
 
 
-# Helper namedtuple for slices of spin blocks
-SpinBlockSlice = namedtuple('SpinBlockSlice',
-                            ['block', 'spin', 'fromslice', 'toslice'])
+@dataclass(frozen=True, slots=True)
+class SpinBlockSlice4D:
+    block: Block4D
+    spin: Spin4D
+    fromslice: IntSlice4D
+    toslice: IntSlice4D
+
+
+def range_in(inner: IntSlice, full: IntSlice) -> bool:
+    return all(r in range(full.start, full.stop)
+               for r in range(inner.start, inner.stop))
 
 
 class EriBuilder:
@@ -48,49 +73,57 @@ class EriBuilder:
           Gets passed the block as a string like 'OOVV' and the spin block as
           as string like 'abab'.
     """
-    def __init__(self, n_orbs, n_orbs_alpha, n_alpha, n_beta, restricted):
-        self.n_orbs = n_orbs
-        self.n_orbs_alpha = n_orbs_alpha
-        self.n_alpha = n_alpha
-        self.n_beta = n_beta
-        self.eri_cache = {}
-        self.restricted = restricted
-        self.block2slice = {
-            "oa": slice(0, self.n_alpha, 1),
-            "va": slice(self.n_alpha, self.n_orbs_alpha, 1),
-            "ob": slice(self.n_orbs_alpha, self.n_orbs_alpha + self.n_beta, 1),
-            "vb": slice(self.n_orbs_alpha + self.n_beta, self.n_orbs, 1),
+    def __init__(self, n_orbs: int, n_orbs_alpha: int, n_alpha: int, n_beta: int,
+                 restricted: bool):
+        self.n_orbs: int = n_orbs
+        self.n_orbs_alpha: int = n_orbs_alpha
+        self.n_alpha: int = n_alpha
+        self.n_beta: int = n_beta
+        self.eri_cache: dict[str, Array4D] = {}
+        self.restricted: bool = restricted
+        self.block2slice: dict[tuple[Block, Spin], IntSlice] = {
+            ("O", "a"): slice(0, self.n_alpha, 1),
+            ("V", "a"): slice(self.n_alpha, self.n_orbs_alpha, 1),
+            ("O", "b"): slice(self.n_orbs_alpha,
+                              self.n_orbs_alpha + self.n_beta, 1),
+            ("V", "b"): slice(self.n_orbs_alpha + self.n_beta, self.n_orbs, 1),
         }
 
-    def compute_mo_eri(self, blocks, spins):
+    def compute_mo_eri(self, blocks: Block4D, spins: Spin4D) -> Array4D:
         """
         Compute block of the ERI tensor in chemists' indexing
         """
         raise NotImplementedError("Implement compute_mo_eri")
 
-    def split_4d_slice(self, slices):
+    def split_4d_slice(self, slices: IntSlice4D) -> list[SpinBlockSlice4D]:
         """
         Split tuple of four slices into the block spin slices
         and their mapping to where elements are to be placed
         """
-        return [SpinBlockSlice(tpl[0][0] + tpl[1][0] + tpl[2][0] + tpl[3][0],
-                               tpl[0][1] + tpl[1][1] + tpl[2][1] + tpl[3][1],
-                               (tpl[0][2], tpl[1][2], tpl[2][2], tpl[3][2]),
-                               (tpl[0][3], tpl[1][3], tpl[2][3], tpl[3][3]))
-                for tpl in product(*(self.split_1d_slice(sl) for sl in slices))]
+        splitted = (self.split_1d_slice(sl) for sl in slices)
+        return [SpinBlockSlice4D(
+            (sl1.block, sl2.block, sl3.block, sl4.block),
+            (sl1.spin, sl2.spin, sl3.spin, sl4.spin),
+            (sl1.fromslice, sl2.fromslice, sl3.fromslice, sl4.fromslice),
+            (sl1.toslice, sl2.toslice, sl3.toslice, sl4.toslice)
+        ) for sl1, sl2, sl3, sl4 in product(*splitted)]
 
-    def split_1d_slice(self, sl):
+    def split_1d_slice(
+        self, sl: "slice[int | None, int, int | None]"
+    ) -> list[SpinBlockSlice]:
         """
         Split slice into block-slices or multiple block-slices
         """
         if sl.start is None:
             sl = slice(0, sl.stop, 1)
         if sl.step is None:
             sl = slice(sl.start, sl.stop, 1)
+        assert is_int_slice(sl)
 
-        ret = []
+        ret: list[SpinBlockSlice] = []
         for (block, bslice) in self.block2slice.items():
-            fromslice = toslice = None
+            fromslice: tuple[int, int] | None = None
+            toslice: tuple[int, int] | None = None
             if range_in(sl, bslice):
                 fromslice = (sl.start - bslice.start, sl.stop - bslice.start)
                 toslice = (0, sl.stop - sl.start)
@@ -107,20 +140,30 @@ def split_1d_slice(self, sl):
                 toslice = (bslice.start - sl.start, sl.stop - sl.start)
             if fromslice is None or toslice is None:
                 continue   # Not found
-            ret.append(SpinBlockSlice(block[0].upper(), block[1],
-                                      slice(*fromslice), slice(*toslice)))
+            ret.append(SpinBlockSlice(
+                block[0], block[1], slice(*fromslice, 1), slice(*toslice, 1)
+            ))
         assert len(ret) > 0
         return ret
 
-    def fill_slice_symm(self, slices, out):
+    def fill_slice_symm(self, slices: IntSlice4D, out: Array4D) -> None:
+        non_zero_spin_blocks: list[Spin4D] = [  # chemist notation
+            ("a", "a", "a", "a"),
+            ("a", "a", "b", "b"),
+            ("b", "b", "a", "a"),
+            ("b", "b", "b", "b"),
+        ]
         for sbslices in self.split_4d_slice(slices):
-            blocks, spins, fromslices, toslices = sbslices
-            if spins not in ["aaaa", "aabb", "bbaa", "bbbb"]:
+            blocks: Block4D = sbslices.block
+            spins: Spin4D = sbslices.spin
+            fromslices: IntSlice4D = sbslices.fromslice
+            toslices: IntSlice4D = sbslices.toslice
+            if spins not in non_zero_spin_blocks:
                 out[toslices] = 0  # Zero by symmetry
                 continue
             if self.restricted:
                 # For restricted spins in chem eri do not matter
-                spins = "aaaa"
+                spins = ("a", "a", "a", "a")
 
             cache_key = "".join(blocks) + "".join(spins)
             if cache_key in self.eri_cache:
@@ -131,5 +174,5 @@ def fill_slice_symm(self, slices, out):
 
             out[toslices] = eri[fromslices]
 
-    def flush_cache(self):
+    def flush_cache(self) -> None:
         self.eri_cache = {}