linting

ebcc/backend/__init__.py

@@ -94,9 +94,9 @@ def _put(
         # TODO MPI has to be manually managed here
         if isinstance(indices, tuple):
             indices_flat = np.ravel_multi_index(indices, array.shape)
-            array.write(indices_flat, values)
+            array.write(indices_flat, values)  # type: ignore
         else:
-            array.write(indices, values)
+            array.write(indices, values)  # type: ignore
         return array
     else:
         raise NotImplementedError(f"Backend {BACKEND} _put not implemented.")

ebcc/backend/_ctf.py

@@ -4,7 +4,6 @@
 import ctf
 import numpy
 import opt_einsum
-import warnings
 
 
 def __getattr__(name):
@@ -16,17 +15,17 @@ class FakeLinalg:
     """Fake linalg module for CTF."""
 
     def __getattr__(self, name):
+        """Get the attribute from CTF's linalg module."""
         return getattr(ctf.linalg, name)
 
-    def eigh(self, a):
+    def eigh(self, a):  # noqa: D102
         # TODO Need to determine if SCALAPACK is available
         w, v = numpy.linalg.eigh(a.to_nparray())
         w = ctf.astensor(w)
         v = ctf.astensor(v)
         return w, v
 
-    #norm = ctf.norm
-    def norm(self, a, ord=None):
+    def norm(self, a, ord=None):  # noqa: D102
         return ctf.norm(a, ord=ord)
 
 
@@ -40,48 +39,49 @@ def norm(self, a, ord=None):
 
 _array = ctf.array
 
-def array(obj, **kwargs):
+
+def array(obj, **kwargs):  # noqa: D103
     if isinstance(obj, ctf.tensor):
         return obj
     return _array(numpy.asarray(obj), **kwargs)
 
 
-def astype(obj, dtype):
+def astype(obj, dtype):  # noqa: D103
     return obj.astype(dtype)
 
 
-def zeros_like(obj):
+def zeros_like(obj):  # noqa: D103
     return ctf.zeros(obj.shape).astype(obj.dtype)
 
 
-def ones_like(obj):
+def ones_like(obj):  # noqa: D103
     return ctf.ones(obj.shape).astype(obj.dtype)
 
 
-def arange(start, stop=None, step=1, dtype=None):
+def arange(start, stop=None, step=1, dtype=None):  # noqa: D103
     if stop is None:
         stop = start
         start = 0
     return ctf.arange(start, stop, step=step, dtype=dtype)
 
 
-def argmin(obj):
+def argmin(obj):  # noqa: D103
     return ctf.to_nparray(obj).argmin()
 
 
-def argmax(obj):
+def argmax(obj):  # noqa: D103
     return ctf.to_nparray(obj).argmax()
 
 
-def bitwise_and(a, b):
+def bitwise_and(a, b):  # noqa: D103
     return a * b
 
 
-def bitwise_not(a):
+def bitwise_not(a):  # noqa: D103
     return ones_like(a) - a
 
 
-def concatenate(arrays, axis=None):
+def concatenate(arrays, axis=None):  # noqa: D103
     if axis is None:
         axis = 0
     if axis < 0:
@@ -106,15 +106,15 @@ def concatenate(arrays, axis=None):
     return result
 
 
-def _block_recursive(arrays, max_depth, depth=0):
+def _block_recursive(arrays, max_depth, depth=0):  # noqa: D103
     if depth < max_depth:
         arrs = [_block_recursive(arr, max_depth, depth + 1) for arr in arrays]
         return concatenate(arrs, axis=-(max_depth - depth))
     else:
         return arrays
 
 
-def block(arrays):
+def block(arrays):  # noqa: D103
     def _get_max_depth(arrays):
         if isinstance(arrays, list):
             return 1 + max([_get_max_depth(arr) for arr in arrays])

ebcc/backend/_jax.py

@@ -15,7 +15,8 @@ def __getattr__(name):
 
 _jax_ix_ = jax.numpy.ix_
 
-def ix_(*args):
+
+def ix_(*args):  # noqa: D103
     args_ = []
     for arg in args:
         if isinstance(arg, jax.numpy.ndarray) and arg.dtype == jax.numpy.bool_:

ebcc/backend/_numpy.py

@@ -8,15 +8,15 @@
 
 if TYPE_CHECKING:
     from numpy import generic
-    from numpy.typing import DType, NDArray
+    from numpy.typing import NDArray
 
 
 def __getattr__(name: str) -> object:
     """Get the attribute from NumPy."""
     return getattr(numpy, name)
 
 
-def astype(obj: NDArray[generic], dtype: DType) -> NDArray[generic]:
+def astype(obj: NDArray[generic], dtype: type[generic]) -> NDArray[generic]:
     """Cast the array to the specified type.
 
     Args:

ebcc/backend/_tensorflow.py

@@ -15,19 +15,19 @@ def __getattr__(name):
     return getattr(tensorflow.experimental.numpy, name)
 
 
-def astype(obj, dtype):
+def astype(obj, dtype):  # noqa: D103
     return obj.astype(dtype)
 
 
-def _block_recursive(arrays, max_depth, depth=0):
+def _block_recursive(arrays, max_depth, depth=0):  # noqa: D103
     if depth < max_depth:
         arrs = [_block_recursive(arr, max_depth, depth + 1) for arr in arrays]
         return tensorflow.experimental.numpy.concatenate(arrs, axis=-(max_depth - depth))
     else:
         return arrays
 
 
-def block(arrays):
+def block(arrays):  # noqa: D103
     def _get_max_depth(arrays):
         if isinstance(arrays, list):
             return 1 + max([_get_max_depth(arr) for arr in arrays])
@@ -36,7 +36,7 @@ def _get_max_depth(arrays):
     return _block_recursive(arrays, _get_max_depth(arrays))
 
 
-def ravel_multi_index(multi_index, dims, mode="raise", order="C"):
+def ravel_multi_index(multi_index, dims, mode="raise", order="C"):  # noqa: D103
     if mode != "raise":
         raise NotImplementedError("Only 'raise' mode is implemented")
     if order != "C":
@@ -51,7 +51,7 @@ def ravel_multi_index(multi_index, dims, mode="raise", order="C"):
     return flat_index
 
 
-def indices(dimensions, dtype=tf.int32, sparse=False):
+def indices(dimensions, dtype=tf.int32, sparse=False):  # noqa: D103
     # Generate a range of indices for each dimension
     ranges = [tf.range(dim, dtype=dtype) for dim in dimensions]
 

ebcc/cc/base.py

@@ -4,7 +4,7 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 from ebcc import default_log, init_logging
 from ebcc import numpy as np
@@ -256,7 +256,9 @@ def kernel(self, eris: Optional[ERIsInputType] = None) -> float:
                 vector = self.amplitudes_to_vector(amplitudes)
                 vector = diis.update(vector)
                 amplitudes = self.vector_to_amplitudes(vector)
-                dt = np.linalg.norm(np.abs(vector - self.amplitudes_to_vector(amplitudes_prev)), ord=np.inf)
+                dt = np.linalg.norm(
+                    np.abs(vector - self.amplitudes_to_vector(amplitudes_prev)), ord=np.inf
+                )
 
                 # Update the energy and calculate change:
                 e_prev = e_cc
@@ -951,7 +953,7 @@ def const(self) -> float:
         """
         if self.options.shift:
             assert self.omega is not None
-            return np.ravel(util.einsum("I,I->", self.omega, self.xi**2.0))[0]
+            return cast(float, np.ravel(util.einsum("I,I->", self.omega, self.xi**2.0))[0])
         return 0.0
 
     @property

ebcc/cc/gebcc.py

@@ -192,7 +192,8 @@ def from_uebcc(cls, ucc: UEBCC) -> GEBCC:
                                     amplitudes[name] = _put(
                                         amplitudes[name],
                                         mask,
-                                        amplitudes[name][mask] + np.transpose(amp, transpose) * sign,
+                                        amplitudes[name][mask]
+                                        + np.transpose(amp, transpose) * sign,
                                     )
                             done.add(combn)
 
@@ -234,7 +235,8 @@ def from_uebcc(cls, ucc: UEBCC) -> GEBCC:
                                     amplitudes[name] = _put(
                                         amplitudes[name],
                                         mask,
-                                        amplitudes[name][mask] + np.transpose(amp, transpose) * sign,
+                                        amplitudes[name][mask]
+                                        + np.transpose(amp, transpose) * sign,
                                     )
                             done.add(combn)
 

ebcc/cc/uebcc.py

@@ -243,7 +243,9 @@ def init_amps(self, eris: Optional[ERIsInputType] = None) -> Namespace[SpinArray
                 elif n == 2:
                     comb_t = comb[0] + comb[2] + comb[1] + comb[3]
                     key_t = key[0] + key[2] + key[1] + key[3]
-                    tn[comb] = np.transpose(eris[comb_t][key_t], (0, 2, 1, 3)) / self.energy_sum(key, comb)
+                    tn[comb] = np.transpose(eris[comb_t][key_t], (0, 2, 1, 3)) / self.energy_sum(
+                        key, comb
+                    )
                     if comb in ("aaaa", "bbbb"):
                         # TODO generalise:
                         tn[comb] = (tn[comb] - np.transpose(tn[comb], (1, 0, 2, 3))) * 0.5

ebcc/core/damping.py

@@ -112,10 +112,12 @@ def update(self, x: NDArray[T], xerr: Optional[NDArray[T]] = None) -> NDArray[T]
         # this looks crazy, but it's just updating the `self._index`th row and
         # column with the new errors, it's just done this way to avoid using
         # calls to `__setitem__` in immutable backends
-        m_i = np.array([
-            np.ravel(np.dot(np.conj(np.ravel(x1)), np.ravel(self._errors[i])))[0]
-            for i in range(nd)
-        ])
+        m_i = np.array(
+            [
+                np.ravel(np.dot(np.conj(np.ravel(x1)), np.ravel(self._errors[i])))[0]
+                for i in range(nd)
+            ]
+        )
         m_i = np.concatenate([np.array([1.0]), m_i, np.zeros(self.space - nd)])
         m_i = np.reshape(m_i, (-1, 1))
         m_j = np.conj(np.transpose(m_i))
@@ -166,7 +168,7 @@ def extrapolate(self, nd: Optional[int] = None) -> NDArray[T]:
         if np.any(np.abs(w) < 1e-14):
             mask = np.abs(w) > 1e-14
             w, v = w[mask], v[:, mask]
-        c = util.einsum("pi,qi,i,q->p", v, np.conj(v), w ** -1, g)
+        c = util.einsum("pi,qi,i,q->p", v, np.conj(v), w**-1.0, g)
 
         # Construct the new vector
         xnew: NDArray[T] = np.zeros_like(self._arrays[0])

ebcc/eom/base.py

@@ -17,7 +17,7 @@
 if TYPE_CHECKING:
     from typing import Any, Callable, Optional
 
-    from numpy import float64, int64
+    from numpy import float64
     from numpy.typing import NDArray
 
     from ebcc.cc.base import BaseEBCC, ERIsInputType, SpaceType, SpinArrayType

ebcc/eom/geom.py

@@ -11,7 +11,7 @@
 if TYPE_CHECKING:
     from typing import Optional
 
-    from numpy import float64, int64
+    from numpy import float64
     from numpy.typing import NDArray
 
     from ebcc.cc.gebcc import GEBCC, ERIsInputType, SpinArrayType
@@ -32,7 +32,7 @@ class GEOM(BaseEOM):
 class IP_GEOM(GEOM, BaseIP_EOM):
     """Generalised ionisation potential equation-of-motion coupled cluster."""
 
-    def _argsort_guesses(self, diag: NDArray[T]) -> NDArray[int64]:
+    def _argsort_guesses(self, diag: NDArray[T]) -> list[int]:
         """Sort the diagonal to inform the initial guesses."""
         if self.options.koopmans:
             r1 = self.vector_to_amplitudes(diag)["r1"]
@@ -128,7 +128,7 @@ def vector_to_amplitudes(self, vector: NDArray[T]) -> Namespace[SpinArrayType]:
 class EA_GEOM(GEOM, BaseEA_EOM):
     """Generalised electron affinity equation-of-motion coupled cluster."""
 
-    def _argsort_guesses(self, diag: NDArray[T]) -> NDArray[int64]:
+    def _argsort_guesses(self, diag: NDArray[T]) -> list[int]:
         """Sort the diagonal to inform the initial guesses."""
         if self.options.koopmans:
             r1 = self.vector_to_amplitudes(diag)["r1"]
@@ -224,7 +224,7 @@ def vector_to_amplitudes(self, vector: NDArray[T]) -> Namespace[SpinArrayType]:
 class EE_GEOM(GEOM, BaseEE_EOM):
     """Generalised electron-electron equation-of-motion coupled cluster."""
 
-    def _argsort_guesses(self, diag: NDArray[T]) -> NDArray[int64]:
+    def _argsort_guesses(self, diag: NDArray[T]) -> list[int]:
         """Sort the diagonal to inform the initial guesses."""
         if self.options.koopmans:
             r1 = self.vector_to_amplitudes(diag)["r1"]

ebcc/eom/reom.py

@@ -11,7 +11,7 @@
 if TYPE_CHECKING:
     from typing import Optional
 
-    from numpy import float64, int64
+    from numpy import float64
     from numpy.typing import NDArray
 
     from ebcc.cc.rebcc import REBCC, ERIsInputType, SpinArrayType
@@ -220,7 +220,7 @@ def vector_to_amplitudes(self, vector: NDArray[T]) -> Namespace[SpinArrayType]:
 class EE_REOM(REOM, BaseEE_EOM):
     """Restricted electron-electron equation-of-motion coupled cluster."""
 
-    def _argsort_guesses(self, diag: NDArray[T]) -> NDArray[int64]:
+    def _argsort_guesses(self, diag: NDArray[T]) -> list[int]:
         """Sort the diagonal to inform the initial guesses."""
         if self.options.koopmans:
             r1 = self.vector_to_amplitudes(diag)["r1"]

ebcc/eom/ueom.py

@@ -11,7 +11,7 @@
 if TYPE_CHECKING:
     from typing import Optional
 
-    from numpy import float64, int64
+    from numpy import float64
     from numpy.typing import NDArray
 
     from ebcc.cc.uebcc import UEBCC, ERIsInputType, SpinArrayType
@@ -32,7 +32,7 @@ class UEOM(BaseEOM):
 class IP_UEOM(UEOM, BaseIP_EOM):
     """Unrestricted ionisation potential equation-of-motion coupled cluster."""
 
-    def _argsort_guesses(self, diag: NDArray[T]) -> NDArray[int64]:
+    def _argsort_guesses(self, diag: NDArray[T]) -> list[int]:
         """Sort the diagonal to inform the initial guesses."""
         if self.options.koopmans:
             r1 = self.vector_to_amplitudes(diag)["r1"]
@@ -147,7 +147,7 @@ def vector_to_amplitudes(self, vector: NDArray[T]) -> Namespace[SpinArrayType]:
 class EA_UEOM(UEOM, BaseEA_EOM):
     """Unrestricted electron affinity equation-of-motion coupled cluster."""
 
-    def _argsort_guesses(self, diag: NDArray[T]) -> NDArray[int64]:
+    def _argsort_guesses(self, diag: NDArray[T]) -> list[int]:
         """Sort the diagonal to inform the initial guesses."""
         if self.options.koopmans:
             r1 = self.vector_to_amplitudes(diag)["r1"]
@@ -262,7 +262,7 @@ def vector_to_amplitudes(self, vector: NDArray[T]) -> Namespace[SpinArrayType]:
 class EE_UEOM(UEOM, BaseEE_EOM):
     """Unrestricted electron-electron equation-of-motion coupled cluster."""
 
-    def _argsort_guesses(self, diag: NDArray[T]) -> NDArray[int64]:
+    def _argsort_guesses(self, diag: NDArray[T]) -> list[int]:
         """Sort the diagonal to inform the initial guesses."""
         if self.options.koopmans:
             r1 = self.vector_to_amplitudes(diag)["r1"]

ebcc/ham/cderis.py

@@ -57,7 +57,9 @@ def __getitem__(self, key: str, e2: Optional[bool] = False) -> NDArray[T]:
         if key_e2 not in self._members:
             s = 0 if not e2 else 2
             coeffs = [
-                numpy.asarray(self.mo_coeff[i + s][:, self.space[i + s].slice(k)], dtype=numpy.float64)
+                numpy.asarray(
+                    self.mo_coeff[i + s][:, self.space[i + s].slice(k)], dtype=numpy.float64
+                )
                 for i, k in enumerate(key)
             ]
             ijslice = (

ebcc/ham/elbos.py

@@ -4,8 +4,8 @@
 
 from typing import TYPE_CHECKING
 
-from ebcc.ham.base import BaseElectronBoson, BaseGHamiltonian, BaseRHamiltonian, BaseUHamiltonian
 from ebcc import numpy as np
+from ebcc.ham.base import BaseElectronBoson, BaseGHamiltonian, BaseRHamiltonian, BaseUHamiltonian
 
 if TYPE_CHECKING:
     from numpy import float64