C function for energy distribution spectral density

src/nmreval/distributions/energy.py

@@ -1,13 +1,18 @@
+import ctypes
 from itertools import product
 
 import numpy as np
+from scipy import LowLevelCallable
 from scipy.integrate import quad, simps as simpson
 
 from .base import Distribution
 from ..lib.utils import ArrayLike
 from ..utils.constants import kB
 
+from .helper import HAS_C_FUNCS, lib
 
+
+# noinspection PyMethodOverriding
 class EnergyBarriers(Distribution):
     name = 'Energy barriers'
     parameter = [r'\tau_{0}', r'E_{m}', r'\Delta E']
@@ -51,7 +56,7 @@ class EnergyBarriers(Distribution):
         omega = np.atleast_1d(omega)
         temperature = np.atleast_1d(temperature)
 
-        e_axis = np.linspace(max(0, e_m-50*e_b), e_m+50*e_b, num=5001)
+        e_axis = np.linspace(max(0., e_m-50*e_b), e_m+50*e_b, num=5001)
         ret_val = []
         for o, tt in product(omega, temperature):
             ret_val.append(simpson(_integrand_freq_real(e_axis, o, tau0, e_m, e_b, tt), e_axis) -
@@ -60,23 +65,41 @@ class EnergyBarriers(Distribution):
         return np.array(ret_val)
 
     @staticmethod
-    def specdens(omega, temperature, *args):
+    def specdens(omega, temperature, tau0: float, e_m: float, e_b: float):
         # in contrast to other spectral densities, it's omega and temperature
-        tau0, e_m, e_b = args
-
-        def integrand(e_a, w, t0, mu, sigma, t):
-            r = EnergyBarriers.rate(t0, e_a, t)
-            return r/(r**2 + w**2) * EnergyBarriers.energydistribution(e_a, mu, sigma)
 
         omega = np.atleast_1d(omega)
         temperature = np.atleast_1d(temperature)
 
-        e_axis = np.linspace(max(0, e_m-50*e_b), e_m+50*e_b, num=5001)
+        if HAS_C_FUNCS:
+            ret_val = EnergyBarriers.spec_dens_c(omega, temperature, tau0, e_m, e_b)
+        else:
+            ret_val = EnergyBarriers.spec_dens_py(omega, temperature, tau0, e_m, e_b)
 
-        ret_val = np.array([simpson(integrand(e_axis, o, tau0, e_m, e_b, tt), e_axis)
-                            for o in omega for tt in temperature])
+        return ret_val.squeeze()
 
-        return ret_val
+    @staticmethod
+    def spec_dens_c(omega: np.ndarray, temperature: np.ndarray, tau0: float, e_m: float, e_b: float) -> np.ndarray:
+        res = []
+        for o, t in product(omega, temperature):
+            c = (ctypes.c_double * 5)(o, tau0, e_m, e_b, t)
+            user_data = ctypes.cast(ctypes.pointer(c), ctypes.c_void_p)
+            area = quad(LowLevelCallable(lib.energyDist_SD, user_data), 0, np.infty, epsabs=1e-10)[0]
+            res.append(area)
+
+        return np.array(res)
+
+    @staticmethod
+    def spec_dens_py(omega: np.ndarray, temperature: np.ndarray, tau0: float, e_m: float, e_b: float) -> np.ndarray:
+        def integrand(e_a, w, t0, mu, sigma, t):
+            r = EnergyBarriers.rate(t0, e_a, t)
+            return r/(r**2 + w**2) * EnergyBarriers.energydistribution(e_a, mu, sigma)
+
+        e_axis = np.linspace(max(0., e_m-50*e_b), e_m+50*e_b, num=5001)
+
+        ret_val = [simpson(integrand(e_axis, o, tau0, e_m, e_b, tt), e_axis) for o in omega for tt in temperature]
+
+        return np.array(ret_val)
 
     @staticmethod
     def mean(*args):