Update and fix pre-commit hooks

Author: greschd

.github/workflows/ci.yml

@@ -1,5 +1,4 @@
 name: Continuous Integration
-
 on: [push, pull_request]
 jobs:
   docs:
@@ -11,29 +10,24 @@ jobs:
     steps:
     - name: Checkout code
       uses: actions/checkout@v2
-
     - name: Install apt dependencies
       run: |
         sudo apt update -qq
         sudo apt install -qq liblapack3 liblapack-dev
-
     - uses: actions/cache@v2
       with:
         path: ~/.cache/pip
         key: pip-${{ matrix.python-version }}-${{ matrix.install-type }}-${{ hashFiles('**/setup.json')
           }}
         restore-keys: pip-${{ matrix.python-version }}-${{ matrix.install-type }}
-
     - name: Set up Python
       uses: actions/setup-python@v2
       with:
         python-version: ${{ matrix.python-version }}
-
     - name: Install the python project
       env:
         INSTALL_TYPE: ${{ matrix.install-type }}
       run: .ci/install_script.sh
-
     - name: Build documentation
       env:
         READTHEDOCS: 'True'
@@ -42,7 +36,6 @@ jobs:
       with:
         name: doc-build
         path: doc/build/html
-
   pre-commit:
     runs-on: ubuntu-latest
     strategy:
@@ -52,29 +45,24 @@ jobs:
     steps:
     - name: Checkout code
       uses: actions/checkout@v2
-
     - name: Install apt dependencies
       run: |
         sudo apt update -qq
         sudo apt install -qq liblapack3 liblapack-dev
-
     - uses: actions/cache@v2
       with:
         path: ~/.cache/pip
         key: pip-${{ matrix.python-version }}-${{ matrix.install-type }}-${{ hashFiles('**/setup.json')
           }}
         restore-keys: pip-${{ matrix.python-version }}-${{ matrix.install-type }}
-
     - name: Set up Python
       uses: actions/setup-python@v2
       with:
         python-version: ${{ matrix.python-version }}
-
     - name: Install the python project
       env:
         INSTALL_TYPE: ${{ matrix.install-type }}
       run: .ci/install_script.sh
-
     - name: Run pre-commit
       run: pre-commit run --all-files || ( git status --short ; git diff ; exit 1
         )
@@ -92,29 +80,24 @@ jobs:
     steps:
     - name: Checkout code
       uses: actions/checkout@v2
-
     - name: Install apt dependencies
       run: |
         sudo apt update -qq
         sudo apt install -qq liblapack3 liblapack-dev
-
     - uses: actions/cache@v2
       with:
         path: ~/.cache/pip
         key: pip-${{ matrix.python-version }}-${{ matrix.install-type }}-${{ hashFiles('**/setup.json')
           }}
         restore-keys: pip-${{ matrix.python-version }}-${{ matrix.install-type }}
-
     - name: Set up Python
       uses: actions/setup-python@v2
       with:
         python-version: ${{ matrix.python-version }}
-
     - name: Install the python project
       env:
         INSTALL_TYPE: ${{ matrix.install-type }}
       run: .ci/install_script.sh
-
     - name: Run pytest
       run: pytest --cov=tbmodels --cov-config=.coveragerc
     - name: Run codecov

mypy.ini

@@ -2,6 +2,7 @@
 
 [mypy]
 python_version = 3.7
+plugins = numpy.typing.mypy_plugin
 
 ; Strictness settings
 ; disallow_any_unimported = True

setup.cfg

@@ -30,7 +30,7 @@ python_requires = >=3.7
 setup_requires =
     fastentrypoints
 install_requires =
-    numpy
+    numpy>=1.20
     scipy>=0.15
     h5py>=3.2.0
     fsc.export

tbmodels/_tb_model.py

@@ -20,6 +20,7 @@
 
 import h5py
 import numpy as np
+import numpy.typing as npt
 import scipy.linalg as la
 from scipy.special import factorial
 from fsc.export import export
@@ -96,7 +97,7 @@ def __init__(
         dim: ty.Optional[int] = None,
         occ: ty.Optional[int] = None,
         pos: ty.Optional[ty.Sequence[ty.Sequence[float]]] = None,
-        uc: ty.Optional[np.ndarray] = None,
+        uc: ty.Optional[npt.ArrayLike] = None,
         contains_cc: bool = True,
         cc_check_tolerance: float = 1e-12,
         sparse: bool = False,
@@ -638,14 +639,14 @@ def from_wannier_files(  # pylint: disable=too-many-locals
                 atom_pos_cartesian = np.array([a.pos for a in atom_list_cartesian])
                 if pos_kind == "wannier":
                     pos_cartesian: ty.Union[
-                        ty.List[np.ndarray], np.ndarray
+                        ty.List[npt.NDArray[np.float_]], npt.NDArray[np.float_]
                     ] = wannier_pos_cartesian
                 elif pos_kind == "nearest_atom":
                     if distance_ratio_threshold < 1:
                         raise ValueError(
                             "Invalid value for 'distance_ratio_threshold': must be >= 1."
                         )
-                    pos_cartesian = ty.cast(ty.List[np.ndarray], [])
+                    pos_cartesian = ty.cast(ty.List[npt.NDArray[np.float_]], [])
                     for p in wannier_pos_cartesian:
                         p_reduced = la.solve(kwargs["uc"].T, np.array(p).T).T
                         T_base = np.floor(p_reduced)
@@ -1079,7 +1080,7 @@ def hamilton(
         self,
         k: ty.Union[ty.Sequence[float], ty.Sequence[ty.Sequence[float]]],
         convention: int = 2,
-    ) -> np.ndarray:
+    ) -> npt.NDArray[np.complex_]:
         """
         Calculates the Hamilton matrix for a given k-point or list of
         k-points.
@@ -1130,12 +1131,12 @@ def hamilton(
             H = pos_exponential.conjugate().transpose((0, 2, 1)) * H * pos_exponential
 
         if single_point:
-            return ty.cast(np.ndarray, H[0])
+            return ty.cast(npt.NDArray[np.complex_], H[0])
         return H
 
     def eigenval(
         self, k: ty.Union[ty.Sequence[float], ty.Sequence[ty.Sequence[float]]]
-    ) -> ty.Union[np.ndarray, ty.List[np.ndarray]]:
+    ) -> ty.Union[npt.NDArray[np.float_], ty.List[npt.NDArray[np.float_]]]:
         """
         Returns the eigenvalues at a given k point, or list of k-points.
 
@@ -1149,7 +1150,7 @@ def eigenval(
         hamiltonians = self.hamilton(k)
         if hamiltonians.ndim == 3:
             return [la.eigvalsh(ham) for ham in hamiltonians]
-        return ty.cast(np.ndarray, la.eigvalsh(hamiltonians))
+        return ty.cast(npt.NDArray[np.float_], la.eigvalsh(hamiltonians))
 
     # -------------------MODIFYING THE MODEL ----------------------------#
     def add_hop(
@@ -1442,12 +1443,12 @@ def _apply_operation(  # pylint: disable=too-many-locals
             ty.Tuple[int, ...], ty.Tuple[ty.List[int], ty.List[int]]
         ] = co.defaultdict(lambda: ([], []))
         for (i, Ti), (j, Tj) in itertools.product(enumerate(uc_shift), repeat=2):
-            shift = tuple(np.array(Tj) - np.array(Ti))
+            shift_tuple = tuple(np.array(Tj) - np.array(Ti))
             for idx1, idx2 in itertools.product(
                 sublattices[i].indices, sublattices[j].indices
             ):
-                hop_shifts_idx[shift][0].append(idx1)
-                hop_shifts_idx[shift][1].append(idx2)
+                hop_shifts_idx[shift_tuple][0].append(idx1)
+                hop_shifts_idx[shift_tuple][1].append(idx2)
 
         # create hoppings with shifted R (by uc_shift[j] - uc_shift[i])
         new_hop: HoppingType = co.defaultdict(self._empty_matrix)
@@ -1456,8 +1457,8 @@ def _apply_operation(  # pylint: disable=too-many-locals
                 np.rint(np.dot(symmetry_operation.rotation_matrix, R)),
                 dtype=int,
             )
-            for shift, (idx1, idx2) in hop_shifts_idx.items():
-                new_R = tuple(np.array(R_transformed) + np.array(shift))
+            for shift_tuple, (idx1, idx2) in hop_shifts_idx.items():
+                new_R = tuple(np.array(R_transformed) + np.array(shift_tuple))
                 new_hop[new_R][idx1, idx2] += mat[idx1, idx2]
 
         # apply D(g) ... D(g)^-1 (since D(g) is unitary: D(g)^-1 == D(g)^H)
@@ -1592,7 +1593,7 @@ def change_unit_cell(  # pylint: disable=too-many-branches
                 )
             # convert to reduced coordinates
             if uc is None:
-                new_uc: ty.Optional[np.ndarray] = self.uc
+                new_uc: ty.Optional[npt.NDArray[np.float_]] = self.uc
                 uc_reduced = np.eye(self.dim)
             else:
                 new_uc = np.array(uc)
@@ -1668,7 +1669,7 @@ def supercell(  # pylint: disable=too-many-locals
             new_uc = (self.uc.T * size_array).T
 
         # the new positions, normalized to the supercell
-        new_pos: ty.List[np.ndarray] = []
+        new_pos: ty.List[npt.NDArray[np.float_]] = []
         reduced_pos = np.array([p / size_array for p in self.pos])
         uc_offsets = list(
             np.array(offset)

tbmodels/helpers.py

@@ -8,6 +8,8 @@
 import typing as ty
 
 import numpy as np
+
+import numpy.typing as npt
 from fsc.export import export
 
 
@@ -17,7 +19,7 @@ def matrix_to_hop(
     orbitals: ty.Optional[ty.Sequence[int]] = None,
     R: ty.Sequence[int] = (0, 0, 0),
     multiplier: float = 1.0,
-) -> ty.List[ty.Tuple[complex, int, int, np.ndarray]]:
+) -> ty.List[ty.Tuple[complex, int, int, npt.NDArray[np.int_]]]:
     r"""
     Turns a square matrix into a series of hopping terms.
 

tbmodels/kdotp.py

@@ -8,6 +8,7 @@
 import typing as ty
 
 import numpy as np
+import numpy.typing as npt
 import scipy.linalg as la
 
 from fsc.export import export
@@ -49,7 +50,7 @@ def __init__(
 
     def hamilton(
         self, k: ty.Union[ty.Sequence[float], ty.Sequence[ty.Sequence[float]]]
-    ) -> np.ndarray:
+    ) -> npt.NDArray[np.complex_]:
         """
         Calculates the Hamilton matrix for a given k-point or list of
         k-points.
@@ -69,20 +70,20 @@ def hamilton(
             single_point = False
 
         ham = ty.cast(
-            np.ndarray,
+            npt.NDArray[np.complex_],
             sum(
                 np.prod(k_array ** k_powers, axis=-1).reshape(-1, 1, 1)
                 * mat[np.newaxis, :, :]
                 for k_powers, mat in self.taylor_coefficients.items()
             ),
         )
         if single_point:
-            return ty.cast(np.ndarray, ham[0])
+            return ty.cast(npt.NDArray[np.complex_], ham[0])
         return ham
 
     def eigenval(
         self, k: ty.Union[ty.Sequence[float], ty.Sequence[ty.Sequence[float]]]
-    ) -> ty.Union[np.ndarray, ty.List[np.ndarray]]:
+    ) -> ty.Union[npt.NDArray[np.float_], ty.List[npt.NDArray[np.float_]]]:
         """
         Returns the eigenvalues at a given k point, or list of k-points.
 
@@ -96,4 +97,4 @@ def eigenval(
         hamiltonians = self.hamilton(k)
         if hamiltonians.ndim == 3:
             return [la.eigvalsh(ham) for ham in hamiltonians]
-        return ty.cast(np.ndarray, la.eigvalsh(hamiltonians))
+        return ty.cast(npt.NDArray[np.float_], la.eigvalsh(hamiltonians))