diff --git a/src/mdevaluate/functions.py b/src/mdevaluate/functions.py index 7348d0c..8132e75 100644 --- a/src/mdevaluate/functions.py +++ b/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))