Migrate beta_sobolev calculation from cython to numba

tardis/plasma/properties/radiative_properties.py

@@ -5,6 +5,7 @@
 import numexpr as ne
 from astropy import units as u
 from tardis import constants as const
+from numba import jit, prange
 
 from tardis.plasma.properties.base import ProcessingPlasmaProperty
 from tardis.plasma.properties.util import macro_atom
@@ -116,6 +117,7 @@ def calculate(self, lines, level_number_density, lines_lower_level_index,
         return pd.DataFrame(tau_sobolevs, index=lines.index,
             columns=np.array(level_number_density.columns))
 
+
 class BetaSobolev(ProcessingPlasmaProperty):
     """
     Attributes
@@ -137,11 +139,24 @@ def calculate(self, tau_sobolevs):
                 columns=tau_sobolevs.columns
                 )
 
-        macro_atom.calculate_beta_sobolev(
-            tau_sobolevs.values.ravel(),
-            beta_sobolev.values.ravel())
+        self.calculate_beta_sobolev(tau_sobolevs.values.ravel(),
+                                    beta_sobolev.values.ravel())
         return beta_sobolev
 
+    @staticmethod
+    @jit(nopython=True, parallel=True)
+    def calculate_beta_sobolev(tau_sobolevs, beta_sobolevs):
+        for i in prange(len(tau_sobolevs)):
+            if tau_sobolevs[i] > 1e3:
+                beta_sobolevs[i] = tau_sobolevs[i]**-1
+            elif tau_sobolevs[i] < 1e-4:
+                beta_sobolevs[i] = 1 - 0.5 * tau_sobolevs[i]
+            else:
+                beta_sobolevs[i] = (1 - np.exp(-tau_sobolevs[i])) / (
+                    tau_sobolevs[i])
+        return beta_sobolevs
+
+
 class TransitionProbabilities(ProcessingPlasmaProperty):
     """
     Attributes

tardis/plasma/properties/util/macro_atom.pyx

@@ -25,22 +25,6 @@ cdef double kb = constants.k_B.cgs.value
 cdef double inv_c2 = 1 / (c ** 2)
 
 
-def calculate_beta_sobolev(np.ndarray[double, ndim=1] tau_sobolevs, np.ndarray[double, ndim=1] beta_sobolevs):
-    cdef double beta_sobolev
-    cdef double tau_sobolev
-    cdef int i
-
-    for i in range(len(tau_sobolevs)):
-        tau_sobolev = tau_sobolevs[i]
-
-        if tau_sobolev > 1e3:
-            beta_sobolev = 1 / tau_sobolev
-        elif tau_sobolev < 1e-4:
-            beta_sobolev = 1 - 0.5 * tau_sobolev
-        else:
-            beta_sobolev = (1 - exp(-tau_sobolev)) / tau_sobolev
-        beta_sobolevs[i] = beta_sobolev
-
 def calculate_transition_probabilities(
         double [:] transition_probability_coef,
         double [:, ::1] beta_sobolev, double [:, ::1] j_blues,

tardis_env3.yml

@@ -11,6 +11,7 @@ dependencies:
 - pandas=0.24
 - astropy=3
 
+- numba=0.43
 - numexpr
 - Cython=0.29