Added type hints

src/mdevaluate/functions.py

@@ -1,22 +1,23 @@
 import numpy as np
+from numpy.typing import ArrayLike
 from scipy.special import gamma as spgamma
 from scipy.integrate import quad as spquad
 
 
-def kww(t, A, tau, beta):
+def kww(t: ArrayLike, A: float, tau: float, beta: float) -> ArrayLike:
     return A * np.exp(-((t / tau) ** beta))
 
 
-def kww_1e(A, tau, beta):
+def kww_1e(A: float, tau: float, beta: float) -> float:
     return tau * (-np.log(1 / (np.e * A))) ** (1 / beta)
 
 
-def cole_davidson(omega, A, beta, tau):
+def cole_davidson(omega: ArrayLike, A: float, beta: float, tau: float) -> ArrayLike:
     P = np.arctan(omega * tau)
     return A * np.cos(P) ** beta * np.sin(beta * P)
 
 
-def cole_cole(omega, A, beta, tau):
+def cole_cole(omega: ArrayLike, A: float, beta: float, tau: float) -> ArrayLike:
     return (
         A
         * (omega * tau) ** beta
@@ -29,7 +30,9 @@ def cole_cole(omega, A, beta, tau):
     )
 
 
-def havriliak_negami(omega, A, beta, alpha, tau):
+def havriliak_negami(
+    omega: ArrayLike, A: float, beta: float, alpha: float, tau: float
+) -> ArrayLike:
     r"""
     Imaginary part of the Havriliak-Negami function.
 
@@ -40,26 +43,26 @@ def havriliak_negami(omega, A, beta, alpha, tau):
 
 
 # fits decay of correlation times, e.g. with distance to pore walls
-def colen(d, X, t8, A):
-    return t8 * np.exp(A * np.exp(-d / X))
+def colen(d: ArrayLike, X: float, tau_pc: float, A: float) -> ArrayLike:
+    return tau_pc * np.exp(A * np.exp(-d / X))
 
 
 # fits decay of the plateau height of the overlap function,
 # e.g. with distance to pore walls
-def colenQ(d, X, Qb, g):
+def colenQ(d: ArrayLike, X: float, Qb: float, g: float) -> ArrayLike:
     return (1 - Qb) * np.exp(-((d / X) ** g)) + Qb
 
 
-def vft(T, tau_0, B, T_inf):
+def vft(T: ArrayLike, tau_0: float, B: float, T_inf: float) -> ArrayLike:
     return tau_0 * np.exp(B / (T - T_inf))
 
 
-def arrhenius(T, tau_0, E_a):
+def arrhenius(T: ArrayLike, tau_0: float, E_a: float) -> ArrayLike:
     return tau_0 * np.exp(E_a / T)
 
 
-def MLfit(t, tau, A, alpha):
-    def MLf(z, a):
+def MLfit(t: ArrayLike, tau: float, A: float, alpha: float) -> ArrayLike:
+    def MLf(z: ArrayLike, a: float) -> ArrayLike:
         """Mittag-Leffler function"""
         z = np.atleast_1d(z)
         if a == 0:
@@ -71,7 +74,7 @@ def MLfit(t, tau, A, alpha):
             return np.polynomial.polynomial.polyval(z, 1 / spgamma(a * k + 1))
 
         # a helper for tricky case, from Gorenflo, Loutchko & Luchko
-        def _MLf(z, a):
+        def _MLf(z: float, a: float) -> ArrayLike:
             if z < 0:
                 f = lambda x: (
                     np.exp(-x * (-z) ** (1 / a))