diff --git a/tardis/energy_input/calculate_opacity.py b/tardis/energy_input/calculate_opacity.py index 7b2d11f756d..8a52d198bd8 100644 --- a/tardis/energy_input/calculate_opacity.py +++ b/tardis/energy_input/calculate_opacity.py @@ -1,5 +1,6 @@ import numpy as np import astropy.units as u +from numba import njit from tardis import constants as const from tardis.energy_input.util import kappa_calculation @@ -10,7 +11,8 @@ M_P = const.m_p.to(u.g).value SIGMA_T = const.sigma_T.cgs.value -# TODO: add units for completeness + +@njit def compton_opacity_calculation(energy, ejecta_density): """Calculate the Compton scattering opacity for a given energy (Rybicki & Lightman, 1979) @@ -54,6 +56,7 @@ def compton_opacity_calculation(energy, ejecta_density): return ejecta_density / (M_P * 2) * sigma_KN +@njit def photoabsorption_opacity_calculation( energy, ejecta_density, iron_group_fraction ): @@ -93,6 +96,7 @@ def photoabsorption_opacity_calculation( return si_opacity + fe_opacity +@njit def pair_creation_opacity_calculation( energy, ejecta_density, iron_group_fraction ): diff --git a/tardis/energy_input/energy_source.py b/tardis/energy_input/energy_source.py index 939c5ecc40c..a272ac9f7d2 100644 --- a/tardis/energy_input/energy_source.py +++ b/tardis/energy_input/energy_source.py @@ -1,6 +1,7 @@ import numpy as np import pandas as pd from nuclear.ejecta import Ejecta +from numba import njit def decay_nuclides(shell_mass, initial_composition, epoch): @@ -10,6 +11,7 @@ def decay_nuclides(shell_mass, initial_composition, epoch): return new_fractions +@njit def create_energy_cdf(energy, intensity): """Creates a CDF of given intensities @@ -37,6 +39,7 @@ def create_energy_cdf(energy, intensity): return energy, cdf +@njit def sample_energy_distribution(energy_sorted, cdf): """Randomly samples a CDF of energies diff --git a/tardis/energy_input/gamma_ray_grid.py b/tardis/energy_input/gamma_ray_grid.py index 71d56345799..abc76f15863 100644 --- a/tardis/energy_input/gamma_ray_grid.py +++ b/tardis/energy_input/gamma_ray_grid.py @@ -3,6 +3,7 @@ from nuclear.io.nndc import get_decay_radiation_database, store_decay_radiation import pandas as pd import astropy.units as u +from numba import njit from tardis.energy_input.util import ( solve_quadratic_equation, diff --git a/tardis/energy_input/gamma_ray_interactions.py b/tardis/energy_input/gamma_ray_interactions.py index a766d5021d9..696fb133cb5 100644 --- a/tardis/energy_input/gamma_ray_interactions.py +++ b/tardis/energy_input/gamma_ray_interactions.py @@ -1,5 +1,6 @@ import copy import numpy as np +from numba import njit from tardis.energy_input.util import ( kappa_calculation, @@ -15,6 +16,7 @@ from tardis.energy_input.GXPhoton import GXPhotonStatus +@njit def get_compton_angle(energy): """ Computes the compton angle from the Klein-Nishina equation. @@ -65,7 +67,9 @@ def compton_scatter(photon, compton_angle): Photon phi direction """ # transform original direction vector to cartesian coordinates - original_direction = normalize_vector(photon.direction.cartesian_coords) + original_direction = normalize_vector( + np.array(photon.direction.cartesian_coords) + ) # compute an arbitrary perpendicular vector to the original direction orthogonal_vector = get_perpendicular_vector(original_direction) # determine a random vector with compton_angle to the original direction @@ -126,6 +130,7 @@ def pair_creation(photon): return photon, backward_ray +@njit def scatter_type(compton_opacity, photoabsorption_opacity, total_opacity): """ Determines the scattering type based on process opacities diff --git a/tardis/energy_input/util.py b/tardis/energy_input/util.py index 974abc9e66c..71e33d984c6 100644 --- a/tardis/energy_input/util.py +++ b/tardis/energy_input/util.py @@ -1,6 +1,7 @@ import astropy.units as u import tardis.constants as const import numpy as np +from numba import njit R_ELECTRON_SQUARED = (const.a0.cgs.value * const.alpha.cgs.value ** 2.0) ** 2.0 ELECTRON_MASS_ENERGY_KEV = (const.m_e * const.c ** 2.0).to("keV").value @@ -47,6 +48,7 @@ def cartesian_coords(self): return x, y, z +@njit def spherical_to_cartesian(r, theta, phi): x = r * np.cos(phi) * np.sin(theta) y = r * np.sin(phi) * np.sin(theta) @@ -54,6 +56,7 @@ def spherical_to_cartesian(r, theta, phi): return x, y, z +@njit def cartesian_to_spherical(x, y, z): r = np.sqrt(x ** 2 + y ** 2 + z ** 2) theta = np.arccos(z / r) @@ -61,6 +64,7 @@ def cartesian_to_spherical(x, y, z): return r, theta, phi +@njit def kappa_calculation(energy): """ Calculates kappa for various other calculations @@ -79,6 +83,7 @@ def kappa_calculation(energy): return energy / ELECTRON_MASS_ENERGY_KEV +@njit def euler_rodrigues(theta, direction): """ Calculates the Euler-Rodrigues rotation matrix @@ -116,6 +121,7 @@ def euler_rodrigues(theta, direction): ) +@njit def solve_quadratic_equation(x, y, z, x_dir, y_dir, z_dir, radius_velocity): """ Solves the quadratic equation for the distance to the shell boundary @@ -145,6 +151,7 @@ def solve_quadratic_equation(x, y, z, x_dir, y_dir, z_dir, radius_velocity): return solution_1, solution_2 +@njit def klein_nishina(energy, theta_C): """ Calculates the Klein-Nishina equation @@ -181,6 +188,7 @@ def klein_nishina(energy, theta_C): ) +@njit def compton_theta_distribution(energy, sample_resolution=100): """ Calculates the cumulative distribution function of theta angles @@ -205,6 +213,7 @@ def compton_theta_distribution(energy, sample_resolution=100): return theta_angles, norm_theta_distribution +@njit def get_random_theta_photon(): """Get a random theta direction between 0 and pi Returns @@ -215,6 +224,7 @@ def get_random_theta_photon(): return np.arccos(1.0 - 2.0 * np.random.random()) +@njit def get_random_phi_photon(): """Get a random phi direction between 0 and 2 * pi @@ -252,6 +262,7 @@ def convert_half_life_to_astropy_units(half_life_string): return half_life_with_unit.to(u.s) +@njit def normalize_vector(vector): """ Normalizes a vector in cartesian coordinates @@ -269,6 +280,7 @@ def normalize_vector(vector): return vector / np.linalg.norm(vector) +@njit def get_perpendicular_vector(original_direction): """ Computes a vector which is perpendicular to the input vector