diff --git a/tardis/plasma/properties/radiative_properties.py b/tardis/plasma/properties/radiative_properties.py index c9eb98428aa..3d26f8cd9ec 100644 --- a/tardis/plasma/properties/radiative_properties.py +++ b/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 diff --git a/tardis/plasma/properties/util/macro_atom.pyx b/tardis/plasma/properties/util/macro_atom.pyx index 61c7f0e371d..e1c6fcceaf2 100644 --- a/tardis/plasma/properties/util/macro_atom.pyx +++ b/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, diff --git a/tardis_env3.yml b/tardis_env3.yml index c757e460d5b..5509954eaa0 100644 --- a/tardis_env3.yml +++ b/tardis_env3.yml @@ -11,6 +11,7 @@ dependencies: - pandas=0.24 - astropy=3 +- numba=0.43 - numexpr - Cython=0.29