diff --git a/benchmarks/run_stardis.py b/benchmarks/run_stardis.py index ed0bcb34..bf4ebb25 100644 --- a/benchmarks/run_stardis.py +++ b/benchmarks/run_stardis.py @@ -11,9 +11,9 @@ from stardis.io.model.marcs import read_marcs_model from stardis.plasma import create_stellar_plasma -from stardis.radiation_field.opacities.opacities_solvers import calc_alphas -from stardis.radiation_field.radiation_field_solvers import raytrace -from stardis.radiation_field.opacities.opacities_solvers import calc_alpha_line_at_nu +from stardis.opacities import calc_alphas +from stardis.transport import raytrace +from stardis.opacities import calc_alpha_line_at_nu class BenchmarkStardis: diff --git a/stardis/base.py b/stardis/base.py index ade53c61..0c0e3797 100644 --- a/stardis/base.py +++ b/stardis/base.py @@ -9,9 +9,8 @@ from astropy import units as u from stardis.plasma import create_stellar_plasma -from stardis.radiation_field.opacities.opacities_solvers import calc_alphas -from stardis.radiation_field.radiation_field_solvers import raytrace -from stardis.radiation_field import RadiationField +from stardis.opacities import calc_alphas +from stardis.transport import raytrace base_dir = os.path.abspath(os.path.dirname(__file__)) @@ -73,34 +72,20 @@ def run_stardis(config_fname, tracing_lambdas_or_nus): # plasma stellar_plasma = create_stellar_plasma(stellar_model, adata) - if True: # change to checking source function from config - from stardis.radiation_field.source_functions.blackbody import ( - blackbody_flux_at_nu, - ) - - stellar_radiation_field = RadiationField( - tracing_nus, blackbody_flux_at_nu, stellar_model - ) - - calc_alphas( + # Below becomes radiation field + alphas, gammas, doppler_widths = calc_alphas( stellar_plasma=stellar_plasma, stellar_model=stellar_model, - stellar_radiation_field=stellar_radiation_field, + tracing_nus=tracing_nus, opacity_config=config.opacity, ) - raytrace(stellar_model, stellar_radiation_field, no_of_thetas=config.no_of_thetas) + F_nu = raytrace( + stellar_model, alphas, tracing_nus, no_of_thetas=config.no_of_thetas + ) return STARDISOutput( - stellar_plasma, - stellar_model, - stellar_radiation_field.opacities.opacities_dict, - stellar_radiation_field.opacities.opacities_dict["alpha_line_at_nu_gammas"], - stellar_radiation_field.opacities.opacities_dict[ - "alpha_line_at_nu_doppler_widths" - ], - stellar_radiation_field.F_nu, - stellar_radiation_field.frequencies, + stellar_plasma, stellar_model, alphas, gammas, doppler_widths, F_nu, tracing_nus ) @@ -149,9 +134,6 @@ class STARDISOutput: Numpy array of wavelengths used for spectrum with units of Angstroms. """ - ###TODO: Instead of returning all these various quantities of the radiation, simply return - # the radiation field with class properties that return useful quantities such as spectrum lambda and lambdas. - def __init__( self, stellar_plasma, stellar_model, alphas, gammas, doppler_widths, F_nu, nus ): diff --git a/stardis/opacities/__init__.py b/stardis/opacities/__init__.py new file mode 100644 index 00000000..8694eac3 --- /dev/null +++ b/stardis/opacities/__init__.py @@ -0,0 +1 @@ +from stardis.opacities.base import * diff --git a/stardis/radiation_field/opacities/opacities_solvers/base.py b/stardis/opacities/base.py similarity index 91% rename from stardis/radiation_field/opacities/opacities_solvers/base.py rename to stardis/opacities/base.py index b328ca78..ef56eb5f 100644 --- a/stardis/radiation_field/opacities/opacities_solvers/base.py +++ b/stardis/opacities/base.py @@ -3,19 +3,15 @@ import numba +from tardis.io.config_reader import ConfigurationNameSpace + from astropy import units as u, constants as const -from stardis.radiation_field.opacities.opacities_solvers.broadening import ( - calculate_broadening, -) -from stardis.radiation_field.opacities.opacities_solvers.voigt import voigt_profile -from stardis.radiation_field.opacities.opacities_solvers.util import ( - sigma_file, - map_items_to_indices, - get_number_density, -) - -# constants +from stardis.opacities.broadening import calculate_broadening +from stardis.opacities.voigt import voigt_profile +from stardis.opacities.util import sigma_file, map_items_to_indices, get_number_density + + VACUUM_ELECTRIC_PERMITTIVITY = 1 / (4 * np.pi) BF_CONSTANT = ( 64 @@ -532,18 +528,18 @@ def calc_alan_entries( def calc_alphas( stellar_plasma, stellar_model, - stellar_radiation_field, + tracing_nus, opacity_config, ): """ - Calculates each opacity and adds it to the opacity dictionary contained in the radiation field. + Calculates total opacity. Parameters ---------- stellar_plasma : tardis.plasma.base.BasePlasma stellar_model : stardis.model.base.StellarModel - stellar_radiation_field stardis.radiation_field.base.RadiationField - Contains the frequencies at which opacities are calculated. Also holds the resulting opacity information. + tracing_nus : astropy.unit.quantity.Quantity + Numpy array of frequencies used for ray tracing with units of Hz. opacity_config : tardis.io.config_reader.Configuration Opacity section of the STARDIS configuration. @@ -552,61 +548,59 @@ def calc_alphas( alphas : numpy.ndarray Array of shape (no_of_depth_points, no_of_frequencies). Total opacity at each depth point for each frequency in tracing_nus. + gammas : numpy.ndarray + Array of shape (no_of_lines, no_of_depth_points). Collisional broadening + parameter of each line at each depth point. + doppler_widths : numpy.ndarray + Array of shape (no_of_lines, no_of_depth_points). Doppler width of each + line at each depth point. """ alpha_file = calc_alpha_file( stellar_plasma, stellar_model, - stellar_radiation_field.frequencies, + tracing_nus, opacity_config.file, ) - stellar_radiation_field.opacities.opacities_dict["alpha_file"] = alpha_file - alpha_bf = calc_alpha_bf( stellar_plasma, stellar_model, - stellar_radiation_field.frequencies, + tracing_nus, opacity_config.bf, ) - stellar_radiation_field.opacities.opacities_dict["alpha_bf"] = alpha_bf - alpha_ff = calc_alpha_ff( stellar_plasma, stellar_model, - stellar_radiation_field.frequencies, + tracing_nus, opacity_config.ff, ) - stellar_radiation_field.opacities.opacities_dict["alpha_ff"] = alpha_ff - alpha_rayleigh = calc_alpha_rayleigh( stellar_plasma, stellar_model, - stellar_radiation_field.frequencies, + tracing_nus, opacity_config.rayleigh, ) - stellar_radiation_field.opacities.opacities_dict["alpha_rayleigh"] = alpha_rayleigh - alpha_electron = calc_alpha_electron( stellar_plasma, stellar_model, - stellar_radiation_field.frequencies, + tracing_nus, opacity_config.disable_electron_scattering, ) - stellar_radiation_field.opacities.opacities_dict["alpha_electron"] = alpha_electron - alpha_line_at_nu, gammas, doppler_widths = calc_alpha_line_at_nu( stellar_plasma, stellar_model, - stellar_radiation_field.frequencies, + tracing_nus, opacity_config.line, ) - stellar_radiation_field.opacities.opacities_dict[ - "alpha_line_at_nu" - ] = alpha_line_at_nu - stellar_radiation_field.opacities.opacities_dict["alpha_line_at_nu_gammas"] = gammas - stellar_radiation_field.opacities.opacities_dict[ - "alpha_line_at_nu_doppler_widths" - ] = doppler_widths - alphas = stellar_radiation_field.opacities.calc_total_alphas() - - return alphas + + alphas = ( + alpha_file + + alpha_bf + + alpha_ff + + alpha_rayleigh + + alpha_electron + + alpha_line_at_nu + ) + + ### TODO create opacity_dict to return + return alphas, gammas, doppler_widths diff --git a/stardis/radiation_field/opacities/opacities_solvers/broadening.py b/stardis/opacities/broadening.py similarity index 96% rename from stardis/radiation_field/opacities/opacities_solvers/broadening.py rename to stardis/opacities/broadening.py index 689617c9..0fd17a0a 100644 --- a/stardis/radiation_field/opacities/opacities_solvers/broadening.py +++ b/stardis/opacities/broadening.py @@ -30,7 +30,7 @@ def _calc_doppler_width(nu_line, temperature, atomic_mass): nu_line : float Frequency of line being considered. temperature : float - Temperature of depth points being considered. + Temperature of depth point being considered. atomic_mass : float Atomic mass of element being considered in grams. @@ -188,7 +188,7 @@ def calc_gamma_linear_stark(n_eff_upper, n_eff_lower, electron_density): n_eff_lower : float Effective principal quantum number of lower level of transition. electron_density : float - Electron density at depth point being considered. + Electron density in depth point being considered. Returns ------- @@ -232,7 +232,7 @@ def calc_gamma_quadratic_stark( electron_density : float Electron density at depth point being considered. temperature : float - Temperature of depth points being considered. + Temperature of depth point being considered. Returns ------- @@ -280,7 +280,7 @@ def calc_gamma_van_der_waals( n_eff_lower : float Effective principal quantum number of lower level of transition. temperature : float - Temperature of depth points being considered. + Temperature of depth point being considered. h_density : float Number density of Hydrogen at depth point being considered. h_mass : float @@ -329,7 +329,7 @@ def calc_gamma( ): """ Calculates total collision broadening parameter for a specific line - and depth points. + and depth point. Parameters ---------- @@ -348,7 +348,7 @@ def calc_gamma( electron_density : float Electron density at depth point being considered. temperature : float - Temperature of depth points being considered. + Temperature of depth point being considered. h_density : float Number density of Hydrogen at depth point being considered. h_mass : float @@ -440,7 +440,7 @@ def calculate_broadening( line_cols : dict Matches the name of a quantity to its column index in lines_array. no_depth_points : int - Number of depth pointss. + Number of depth points. atomic_masses : numpy.ndarray Atomic mass of all elements included in the simulation. electron_densities : numpy.ndarray @@ -467,10 +467,10 @@ def calculate_broadening( line_nus : numpy.ndarray Frequency of each line. gammas : numpy.ndarray - Array of shape (no_of_lines, no_depth_points). Collisional broadening + Array of shape (no_of_lines, no_of_depth_points). Collisional broadening parameter of each line at each depth point. doppler_widths : numpy.ndarray - Array of shape (no_of_lines, no_depth_points). Doppler width of each + Array of shape (no_of_lines, no_of_depth_points). Doppler width of each line at each depth point. """ diff --git a/stardis/radiation_field/opacities/tests/__init__.py b/stardis/opacities/tests/__init__.py similarity index 100% rename from stardis/radiation_field/opacities/tests/__init__.py rename to stardis/opacities/tests/__init__.py diff --git a/stardis/radiation_field/opacities/tests/test_broadening.py b/stardis/opacities/tests/test_broadening.py similarity index 99% rename from stardis/radiation_field/opacities/tests/test_broadening.py rename to stardis/opacities/tests/test_broadening.py index ffc27185..be63407c 100644 --- a/stardis/radiation_field/opacities/tests/test_broadening.py +++ b/stardis/opacities/tests/test_broadening.py @@ -3,7 +3,7 @@ from astropy import constants as const from numba import cuda -from stardis.radiation_field.opacities.opacities_solvers.broadening import ( +from stardis.opacities.broadening import ( calc_doppler_width, _calc_doppler_width_cuda, calc_doppler_width_cuda, diff --git a/stardis/opacities/tests/test_voigt.py b/stardis/opacities/tests/test_voigt.py new file mode 100644 index 00000000..d9460c41 --- /dev/null +++ b/stardis/opacities/tests/test_voigt.py @@ -0,0 +1,282 @@ +import pytest +import numpy as np +from math import sqrt +from numba import cuda + +from stardis.opacities.voigt import ( + faddeeva, + _faddeeva_cuda, + faddeeva_cuda, + voigt_profile, + _voigt_profile_cuda, + voigt_profile_cuda, +) + +# Test cases must also take into account use of a GPU to run. If there is no GPU then the test cases will fail. +GPUs_available = cuda.is_available() + +if GPUs_available: + import cupy as cp + + +@pytest.mark.parametrize( + "faddeeva_sample_values_input, faddeeva_sample_values_expected_result", + [ + (0, 1 + 0j), + (0.0, 1.0 + 0.0j), + (np.array([0.0]), np.array([1.0 + 0.0j])), + (np.array([0, 0]), np.array([1 + 0j, 1 + 0j])), + ], +) +def test_faddeeva_sample_values( + faddeeva_sample_values_input, faddeeva_sample_values_expected_result +): + assert np.allclose( + faddeeva(faddeeva_sample_values_input), + faddeeva_sample_values_expected_result, + ) + + +@pytest.mark.skipif( + not GPUs_available, reason="No GPU is available to test CUDA function" +) +@pytest.mark.parametrize( + "faddeeva_cuda_unwrapped_sample_values_input, faddeeva_cuda_unwrapped_sample_values_expected_result", + [ + (np.array([0.0], dtype=complex), np.array([1.0 + 0.0j])), + (np.array([0, 0], dtype=complex), np.array([1 + 0j, 1 + 0j])), + ], +) +def test_faddeeva_cuda_unwrapped_sample_values( + faddeeva_cuda_unwrapped_sample_values_input, + faddeeva_cuda_unwrapped_sample_values_expected_result, +): + test_values = cp.asarray(faddeeva_cuda_unwrapped_sample_values_input) + result_values = cp.empty_like(test_values) + + nthreads = 256 + length = len(faddeeva_cuda_unwrapped_sample_values_input) + nblocks = 1 + (length // nthreads) + + _faddeeva_cuda[nblocks, nthreads](result_values, test_values) + + assert np.allclose( + cp.asnumpy(result_values), + faddeeva_cuda_unwrapped_sample_values_expected_result, + ) + + +@pytest.mark.skipif( + not GPUs_available, reason="No GPU is available to test CUDA function" +) +@pytest.mark.parametrize( + "faddeeva_cuda_wrapped_sample_numpy_values_input, faddeeva_cuda_wrapped_sample_numpy_values_expected_result", + [ + (np.array([0.0]), np.array([1.0 + 0.0j])), + (np.array([0, 0]), np.array([1 + 0j, 1 + 0j])), + ], +) +def test_faddeeva_cuda_wrapped_sample_numpy_values( + faddeeva_cuda_wrapped_sample_numpy_values_input, + faddeeva_cuda_wrapped_sample_numpy_values_expected_result, +): + assert np.allclose( + faddeeva_cuda(faddeeva_cuda_wrapped_sample_numpy_values_input), + faddeeva_cuda_wrapped_sample_numpy_values_expected_result, + ) + + +@pytest.mark.skipif( + not GPUs_available, reason="No GPU is available to test CUDA function" +) +@pytest.mark.parametrize( + "faddeeva_cuda_wrapped_sample_cuda_values_input, faddeeva_cuda_wrapped_sample_cuda_values_expected_result", + [ + ( + np.array([0, 0], dtype=complex), + np.array([1 + 0j, 1 + 0j]), + ), + ], +) +def test_faddeeva_cuda_wrapped_sample_cuda_values( + faddeeva_cuda_wrapped_sample_cuda_values_input, + faddeeva_cuda_wrapped_sample_cuda_values_expected_result, +): + assert np.allclose( + faddeeva_cuda(cp.asarray(faddeeva_cuda_wrapped_sample_cuda_values_input)), + faddeeva_cuda_wrapped_sample_cuda_values_expected_result, + ) + + +test_voigt_profile_division_by_zero_test_values = [ + -100, + -5, + -1, + 0, + 0.0, + 1.2, + 3, + 100, + np.array( + [ + 0, + -1.0, + 1, + ] + ), +] + + +@pytest.mark.parametrize( + "voigt_profile_division_by_zero_input_delta_nu", + test_voigt_profile_division_by_zero_test_values, +) +@pytest.mark.parametrize( + "voigt_profile_division_by_zero_input_gamma", + test_voigt_profile_division_by_zero_test_values, +) +def test_voigt_profile_division_by_zero( + voigt_profile_division_by_zero_input_delta_nu, + voigt_profile_division_by_zero_input_gamma, +): + arg_list = ( + voigt_profile_division_by_zero_input_delta_nu, + 0, + voigt_profile_division_by_zero_input_gamma, + ) + with pytest.raises(ZeroDivisionError): + _ = voigt_profile(*arg_list) + + +@pytest.mark.parametrize( + "voigt_profile_sample_values_input_delta_nu, voigt_profile_sample_values_input_doppler_width, voigt_profile_sample_values_input_gamma, voigt_profile_sample_values_expected_result", + [ + (0, 1, 0, 1 / sqrt(np.pi)), + (0, 2, 0, 1 / (sqrt(np.pi) * 2)), + ( + np.array([0, 0]), + np.array([1, 2]), + np.array([0, 0]), + np.array([1 / sqrt(np.pi), 1 / (sqrt(np.pi) * 2)]), + ), + ], +) +def test_voigt_profile_sample_values_sample_values( + voigt_profile_sample_values_input_delta_nu, + voigt_profile_sample_values_input_doppler_width, + voigt_profile_sample_values_input_gamma, + voigt_profile_sample_values_expected_result, +): + arg_list = ( + voigt_profile_sample_values_input_delta_nu, + voigt_profile_sample_values_input_doppler_width, + voigt_profile_sample_values_input_gamma, + ) + assert np.allclose( + voigt_profile(*arg_list), + voigt_profile_sample_values_expected_result, + ) + + +@pytest.mark.skipif( + not GPUs_available, reason="No GPU is available to test CUDA function" +) +@pytest.mark.parametrize( + "voigt_profile_cuda_sample_values_input_delta_nu, voigt_profile_cuda_sample_values_input_doppler_width, voigt_profile_cuda_sample_values_input_gamma, voigt_profile_cuda_unwrapped_sample_values_expected_result", + [ + ( + np.array([0, 0]), + np.array([1, 2]), + np.array([0, 0]), + np.array([1 / sqrt(np.pi), 1 / (sqrt(np.pi) * 2)]), + ), + ], +) +def test_voigt_profile_cuda_unwrapped_sample_values( + voigt_profile_cuda_sample_values_input_delta_nu, + voigt_profile_cuda_sample_values_input_doppler_width, + voigt_profile_cuda_sample_values_input_gamma, + voigt_profile_cuda_unwrapped_sample_values_expected_result, +): + arg_list = ( + voigt_profile_cuda_sample_values_input_delta_nu, + voigt_profile_cuda_sample_values_input_doppler_width, + voigt_profile_cuda_sample_values_input_gamma, + ) + shortest_arg_idx = np.argmin(map(len, arg_list)) + size = len(arg_list[shortest_arg_idx]) + test_values = map(cp.asarray, arg_list) + result_values = cp.empty_like(arg_list[shortest_arg_idx], dtype=float) + + nthreads = 256 + nblocks = 1 + (size // nthreads) + + _voigt_profile_cuda[nblocks, nthreads](result_values, *arg_list) + + assert np.allclose( + cp.asnumpy(result_values), + voigt_profile_cuda_unwrapped_sample_values_expected_result, + ) + + +@pytest.mark.skipif( + not GPUs_available, reason="No GPU is available to test CUDA function" +) +@pytest.mark.parametrize( + "voigt_profile_cuda_sample_values_input_delta_nu, voigt_profile_cuda_sample_values_input_doppler_width, voigt_profile_cuda_sample_values_input_gamma, voigt_profile_cuda_wrapped_sample_numpy_values_expected_result", + [ + ( + np.array([0, 0]), + np.array([1, 2]), + np.array([0, 0]), + np.array([1 / sqrt(np.pi), 1 / (sqrt(np.pi) * 2)]), + ), + ], +) +def test_voigt_profile_cuda_wrapped_sample_numpy_values( + voigt_profile_cuda_sample_values_input_delta_nu, + voigt_profile_cuda_sample_values_input_doppler_width, + voigt_profile_cuda_sample_values_input_gamma, + voigt_profile_cuda_wrapped_sample_numpy_values_expected_result, +): + arg_list = ( + voigt_profile_cuda_sample_values_input_delta_nu, + voigt_profile_cuda_sample_values_input_doppler_width, + voigt_profile_cuda_sample_values_input_gamma, + ) + + assert np.allclose( + voigt_profile_cuda(*arg_list), + voigt_profile_cuda_wrapped_sample_numpy_values_expected_result, + ) + + +@pytest.mark.skipif( + not GPUs_available, reason="No GPU is available to test CUDA function" +) +@pytest.mark.parametrize( + "voigt_profile_cuda_sample_values_input_delta_nu, voigt_profile_cuda_sample_values_input_doppler_width, voigt_profile_cuda_sample_values_input_gamma, voigt_profile_cuda_wrapped_sample_cuda_values_expected_result", + [ + ( + np.array([0, 0]), + np.array([1, 2]), + np.array([0, 0]), + np.array([1 / sqrt(np.pi), 1 / (sqrt(np.pi) * 2)]), + ), + ], +) +def test_voigt_profile_cuda_wrapped_sample_cuda_values( + voigt_profile_cuda_sample_values_input_delta_nu, + voigt_profile_cuda_sample_values_input_doppler_width, + voigt_profile_cuda_sample_values_input_gamma, + voigt_profile_cuda_wrapped_sample_cuda_values_expected_result, +): + arg_list = ( + voigt_profile_cuda_sample_values_input_delta_nu, + voigt_profile_cuda_sample_values_input_doppler_width, + voigt_profile_cuda_sample_values_input_gamma, + ) + assert np.allclose( + voigt_profile_cuda(*map(cp.asarray, arg_list)), + voigt_profile_cuda_wrapped_sample_cuda_values_expected_result, + ) diff --git a/stardis/radiation_field/opacities/opacities_solvers/util.py b/stardis/opacities/util.py similarity index 100% rename from stardis/radiation_field/opacities/opacities_solvers/util.py rename to stardis/opacities/util.py diff --git a/stardis/opacities/voigt.py b/stardis/opacities/voigt.py new file mode 100644 index 00000000..824b09e6 --- /dev/null +++ b/stardis/opacities/voigt.py @@ -0,0 +1,182 @@ +import numpy as np +import numba +from numba import cuda +import cmath + +GPUs_available = cuda.is_available() + +if GPUs_available: + import cupy as cp + +SQRT_PI = np.sqrt(np.pi, dtype=float) +SQRT_2 = np.sqrt(2, dtype=float) +PI = float(np.pi) + + +@numba.njit() +def _faddeeva(z): + """ + The Faddeeva function. Code adapted from + https://github.com/tiagopereira/Transparency.jl/blob/966fb46c21/src/voigt.jl#L13. + + Parameters + ---------- + z : complex + + Returns + ------- + w : complex + """ + z = complex(z) + x = float(z.real) + y = float(z.imag) + t = y - 1j * x + s = abs(x) + y + w = complex(0.0) + u = t * t + + IN_REGION_I = s > 15.0 + IN_REGION_II = (not IN_REGION_I) and (s > 5.5) + IN_REGION_III = ( + (not IN_REGION_I) and (not IN_REGION_II) and (y >= 0.195 * abs(x) - 0.176) + ) + IN_REGION_IV = (not IN_REGION_I) and (not IN_REGION_II) and (not IN_REGION_III) + + # If in Region I + w = 1j * 1 / SQRT_PI * z / (z**2 - 0.5) if IN_REGION_I else w + + # If in Region II + w = ( + 1j + * (z * (z**2 * 1 / SQRT_PI - 1.4104739589)) + / (0.75 + z**2 * (z**2 - 3.0)) + if IN_REGION_II + else w + ) + + # If in Region III + w = ( + (16.4955 + t * (20.20933 + t * (11.96482 + t * (3.778987 + 0.5642236 * t)))) + / ( + 16.4955 + + t * (38.82363 + t * (39.27121 + t * (21.69274 + t * (6.699398 + t)))) + ) + if IN_REGION_III + else w + ) + + # If in Region IV + numerator = t * ( + 36183.31 + - u + * ( + 3321.99 + - u + * (1540.787 - u * (219.031 - u * (35.7668 - u * (1.320522 - u * 0.56419)))) + ) + ) + denominator = 32066.6 - u * ( + 24322.8 + - u + * (9022.23 - u * (2186.18 - u * (364.219 - u * (61.5704 - u * (1.84144 - u))))) + ) + w = (cmath.exp(u) - numerator / denominator) if IN_REGION_IV else w + + return w + + +@numba.vectorize(nopython=True) +def faddeeva(z): + return _faddeeva(z) + + +@cuda.jit +def _faddeeva_cuda(res, z): + tid = cuda.grid(1) + size = len(res) + + if tid < size: + res[tid] = _faddeeva(z[tid]) + + +def faddeeva_cuda(z, nthreads=256, ret_np_ndarray=False): + size = len(z) + nblocks = 1 + (size // nthreads) + z = cp.asarray(z, dtype=complex) + res = cp.empty_like(z) + + _faddeeva_cuda[nblocks, nthreads](res, z) + return cp.asnumpy(res) if ret_np_ndarray else res + + +@numba.njit +def _voigt_profile(delta_nu, doppler_width, gamma): + """ + Calculates the Voigt profile, the convolution of a Lorentz profile + and a Gaussian profile. + + Parameters + ---------- + delta_nu : float + Difference between the frequency that the profile is being evaluated at + and the line's resonance frequency. + doppler_width : float + Doppler width for Gaussian profile. + gamma : float + Broadening parameter for Lorentz profile. + + Returns + ------- + phi : float + Value of Voigt profile. + """ + delta_nu, doppler_width, gamma = float(delta_nu), float(doppler_width), float(gamma) + + z = complex(delta_nu, gamma / (4 * PI)) / doppler_width + phi = _faddeeva(z).real / (SQRT_PI * doppler_width) + return phi + + +@numba.vectorize(nopython=True) +def voigt_profile(delta_nu, doppler_width, gamma): + return _voigt_profile(delta_nu, doppler_width, gamma) + + +@cuda.jit +def _voigt_profile_cuda(res, delta_nu, doppler_width, gamma): + tid = cuda.grid(1) + + size = len(res) + + if tid < size: + res[tid] = _voigt_profile(delta_nu[tid], doppler_width[tid], gamma[tid]) + + +def voigt_profile_cuda( + delta_nu, + doppler_width, + gamma, + nthreads=256, + ret_np_ndarray=False, + dtype=float, +): + arg_list = ( + delta_nu, + doppler_width, + gamma, + ) + shortest_arg_idx = np.argmin(map(len, arg_list)) + size = len(arg_list[shortest_arg_idx]) + + nblocks = 1 + (size // nthreads) + + arg_list = tuple(map(lambda v: cp.array(v, dtype=dtype), arg_list)) + + res = cp.empty_like(arg_list[shortest_arg_idx], dtype=dtype) + + _voigt_profile_cuda[nblocks, nthreads]( + res, + *arg_list, + ) + + return cp.asnumpy(res) if ret_np_ndarray else res diff --git a/stardis/radiation_field/__init__.py b/stardis/radiation_field/__init__.py deleted file mode 100644 index b0295c66..00000000 --- a/stardis/radiation_field/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from stardis.radiation_field.base import RadiationField diff --git a/stardis/radiation_field/base.py b/stardis/radiation_field/base.py deleted file mode 100644 index b37f760b..00000000 --- a/stardis/radiation_field/base.py +++ /dev/null @@ -1,33 +0,0 @@ -from stardis.radiation_field.opacities import Opacities -import numpy as np - - -class RadiationField: - """ - Class containing information about the radiation field. - ###TODO Radiation field temperature should be a separate attribute, for the case of differing gas and radiation. - ###TODO Implement a source function class. Doesn't need to be done until we have another source function than just blackbody. - - Parameters - ---------- - frequencies : astronopy.units.Quantity - source_function : stardis.radiation_field.source_function - opacities : stardis.radiation_field.opacities - - Attributes - ---------- - frequencies : astropy.units.Quantity - Frequencies of the radiation field. - source_function : stardis.radiation_field.source_function - Source function of the radiation field. - opacities : stardis.radiation_field.opacities - Composition of the model. Includes density and atomic mass fractions. - F_nu : numpy.ndarray - Radiation field fluxes at each frequency at each depth point. Initialized as zeros and calculated by a solver. - """ - - def __init__(self, frequencies, source_function, stellar_model): - self.frequencies = frequencies - self.source_function = source_function - self.opacities = Opacities() - self.F_nu = np.zeros((len(stellar_model.geometry.r), len(frequencies))) diff --git a/stardis/radiation_field/opacities/__init__.py b/stardis/radiation_field/opacities/__init__.py deleted file mode 100644 index 9a1b27e8..00000000 --- a/stardis/radiation_field/opacities/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from stardis.radiation_field.opacities.base import Opacities diff --git a/stardis/radiation_field/opacities/base.py b/stardis/radiation_field/opacities/base.py deleted file mode 100644 index 80739aaf..00000000 --- a/stardis/radiation_field/opacities/base.py +++ /dev/null @@ -1,25 +0,0 @@ -class Opacities: - """ - Holds opacity information. - - Paramaters - ---------- - opacities_dict : dict - Python dictionary to contain each of the sources of opacity by name as well as their contribution at each frequency specified in the radiation field. - total_alphas : numpy.ndarray - Array of the total opacity at each frequency specified in the radiation field at each depth points. - Added as an attribute when calc_total_alphas() is called. - """ - - def __init__(self): - self.opacities_dict = {} - - ###TODO: Better implementation for this - def calc_total_alphas(self): - for i, item in enumerate(self.opacities_dict.items()): - if "gammas" not in item[0] and "doppler" not in item[0]: - if i == 0: - self.total_alphas = item[1] - else: - self.total_alphas += item[1] - return self.total_alphas diff --git a/stardis/radiation_field/opacities/opacities_solvers/__init__.py b/stardis/radiation_field/opacities/opacities_solvers/__init__.py deleted file mode 100644 index 192853c2..00000000 --- a/stardis/radiation_field/opacities/opacities_solvers/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from stardis.radiation_field.opacities.opacities_solvers.base import * diff --git a/stardis/radiation_field/opacities/opacities_solvers/voigt.py b/stardis/radiation_field/opacities/opacities_solvers/voigt.py deleted file mode 100644 index 3852711d..00000000 --- a/stardis/radiation_field/opacities/opacities_solvers/voigt.py +++ /dev/null @@ -1,99 +0,0 @@ -import numpy as np -import numba - - -@numba.njit -def faddeeva(z): - """ - The Faddeeva function. Code adapted from - https://github.com/tiagopereira/Transparency.jl/blob/966fb46c21/src/voigt.jl#L13. - - Parameters - ---------- - z : complex - - Returns - ------- - w : complex - """ - s = abs(np.real(z)) + np.imag(z) - - if s > 15.0: - # Region I - w = 1j * 1 / np.sqrt(np.pi) * z / (z**2 - 0.5) - - elif s > 5.5: - # Region II - w = ( - 1j - * (z * (z**2 * 1 / np.sqrt(np.pi) - 1.4104739589)) - / (0.75 + z**2 * (z**2 - 3.0)) - ) - - else: - x = np.real(z) - y = np.imag(z) - t = y - 1j * x - - if y >= 0.195 * abs(x) - 0.176: - # Region III - w = ( - 16.4955 - + t * (20.20933 + t * (11.96482 + t * (3.778987 + 0.5642236 * t))) - ) / ( - 16.4955 - + t * (38.82363 + t * (39.27121 + t * (21.69274 + t * (6.699398 + t)))) - ) - - else: - # Region IV - u = t * t - numerator = t * ( - 36183.31 - - u - * ( - 3321.99 - - u - * ( - 1540.787 - - u * (219.031 - u * (35.7668 - u * (1.320522 - u * 0.56419))) - ) - ) - ) - denominantor = 32066.6 - u * ( - 24322.8 - - u - * ( - 9022.23 - - u * (2186.18 - u * (364.219 - u * (61.5704 - u * (1.84144 - u)))) - ) - ) - w = np.exp(u) - numerator / denominantor - - return w - - -@numba.njit -def voigt_profile(delta_nu, doppler_width, gamma): - """ - Calculates the Voigt profile, the convolution of a Lorentz profile - and a Gaussian profile. - - Parameters - ---------- - delta_nu : float - Difference between the frequency that the profile is being evaluated at - and the line's resonance frequency. - doppler_width : float - Doppler width for Gaussian profile. - gamma : float - Broadening parameter for Lorentz profile. - - Returns - ------- - phi : float - Value of Voigt profile. - """ - z = (delta_nu + (gamma / (4 * np.pi)) * 1j) / doppler_width - phi = np.real(faddeeva(z)) / (np.sqrt(np.pi) * doppler_width) - return phi diff --git a/stardis/radiation_field/opacities/tests/test_voigt.py b/stardis/radiation_field/opacities/tests/test_voigt.py deleted file mode 100644 index 6fef44d3..00000000 --- a/stardis/radiation_field/opacities/tests/test_voigt.py +++ /dev/null @@ -1,80 +0,0 @@ -import pytest -from numpy import allclose, pi as PI -from math import sqrt - -from stardis.radiation_field.opacities.opacities_solvers.voigt import ( - faddeeva, - voigt_profile, -) - - -@pytest.mark.parametrize( - "faddeeva_sample_values_input, faddeeva_sample_values_expected_result", - [ - (0, 1 + 0j), - (0.0, 1.0 + 0.0j), - ], -) -def test_faddeeva_sample_values( - faddeeva_sample_values_input, faddeeva_sample_values_expected_result -): - assert allclose( - faddeeva(faddeeva_sample_values_input), - faddeeva_sample_values_expected_result, - ) - - -test_voigt_profile_division_by_zero_test_values = [ - -100, - -5, - -1, - 0, - 0.0, - 1j, - 1.2, - 3, - 100, -] - - -@pytest.mark.parametrize( - "voigt_profile_division_by_zero_input_delta_nu", - test_voigt_profile_division_by_zero_test_values, -) -@pytest.mark.parametrize( - "voigt_profile_division_by_zero_input_gamma", - test_voigt_profile_division_by_zero_test_values, -) -def test_voigt_profile_division_by_zero( - voigt_profile_division_by_zero_input_delta_nu, - voigt_profile_division_by_zero_input_gamma, -): - with pytest.raises(ZeroDivisionError): - _ = voigt_profile( - voigt_profile_division_by_zero_input_delta_nu, - 0, - voigt_profile_division_by_zero_input_gamma, - ) - - -@pytest.mark.parametrize( - "voigt_profile_sample_values_input_delta_nu, voigt_profile_sample_values_input_doppler_width, voigt_profile_sample_values_input_gamma, voigt_profile_sample_values_expected_result", - [ - (0, 1, 0, 1 / sqrt(PI)), - (0, 2, 0, 1 / (sqrt(PI) * 2)), - ], -) -def test_voigt_profile_sample_values_sample_values( - voigt_profile_sample_values_input_delta_nu, - voigt_profile_sample_values_input_doppler_width, - voigt_profile_sample_values_input_gamma, - voigt_profile_sample_values_expected_result, -): - assert allclose( - voigt_profile( - voigt_profile_sample_values_input_delta_nu, - voigt_profile_sample_values_input_doppler_width, - voigt_profile_sample_values_input_gamma, - ), - voigt_profile_sample_values_expected_result, - ) diff --git a/stardis/radiation_field/radiation_field_solvers/__init__.py b/stardis/radiation_field/radiation_field_solvers/__init__.py deleted file mode 100644 index bda3ea69..00000000 --- a/stardis/radiation_field/radiation_field_solvers/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from stardis.radiation_field.radiation_field_solvers.base import * diff --git a/stardis/radiation_field/source_functions/blackbody.py b/stardis/radiation_field/source_functions/blackbody.py deleted file mode 100644 index 958b80aa..00000000 --- a/stardis/radiation_field/source_functions/blackbody.py +++ /dev/null @@ -1,35 +0,0 @@ -import numpy as np -from astropy import units as u, constants as const -import numba - -H_CGS = const.h.cgs.value -C_CGS = const.c.cgs.value -K_B_CGS = const.k_B.cgs.value - - -@numba.njit -def blackbody_flux_at_nu(tracing_nus, temps): - """ - Planck blackbody intensity distribution w.r.t. frequency. - - Parameters - ---------- - tracing_nus : astropy.unit.quantity.Quantity - Numpy array of frequencies used for ray tracing with units of Hz. - temps : numpy.ndarray - Temperatures in K of all depth points. Note that array must - be transposed. Must not have astropy units when passed into the function. - - Returns - ------- - bb : astropy.unit.quantity.Quantity - Numpy array of shape (no_of_depth_points, no_of_frequencies) with units - of erg/(s cm^2 Hz). Blackbody specific intensity at each depth point - for each frequency in tracing_nus. - """ - - bb_prefactor = (2 * H_CGS * tracing_nus**3) / C_CGS**2 - bb_flux_nu = bb_prefactor / ( - np.exp(((H_CGS * tracing_nus) / (K_B_CGS * temps))) - 1 - ) - return bb_flux_nu diff --git a/stardis/transport/__init__.py b/stardis/transport/__init__.py new file mode 100644 index 00000000..513a6125 --- /dev/null +++ b/stardis/transport/__init__.py @@ -0,0 +1 @@ +from stardis.transport.base import * diff --git a/stardis/radiation_field/radiation_field_solvers/base.py b/stardis/transport/base.py similarity index 71% rename from stardis/radiation_field/radiation_field_solvers/base.py rename to stardis/transport/base.py index b27ccf2b..20d36c9e 100644 --- a/stardis/radiation_field/radiation_field_solvers/base.py +++ b/stardis/transport/base.py @@ -1,7 +1,36 @@ import numba import numpy as np +from astropy import units as u, constants as const -# from stardis.radiation_field.source_functions.blackbody import blackbody_flux_at_nu +H_CGS = const.h.cgs.value +C_CGS = const.c.cgs.value +K_B_CGS = const.k_B.cgs.value + + +@numba.njit +def bb_nu(tracing_nus, temps): + """ + Planck blackbody intensity distribution w.r.t. frequency. + + Parameters + ---------- + tracing_nus : astropy.unit.quantity.Quantity + Numpy array of frequencies used for ray tracing with units of Hz. + temps : numpy.ndarray + Temperatures in K of all depth points. Note that array must + be transposed. + + Returns + ------- + bb : astropy.unit.quantity.Quantity + Numpy array of shape (no_of_depth_points, no_of_frequencies) with units + of erg/(s cm^2 Hz). Blackbody specific intensity at each depth point + for each frequency in tracing_nus. + """ + + bb_prefactor = (2 * H_CGS * tracing_nus**3) / C_CGS**2 + bb = bb_prefactor / (np.exp(((H_CGS * tracing_nus) / (K_B_CGS * temps))) - 1) + return bb @numba.njit @@ -42,7 +71,6 @@ def single_theta_trace( alphas, tracing_nus, theta, - source_function, ): """ Performs ray tracing at an angle following van Noort 2001 eq 14. @@ -73,15 +101,16 @@ def single_theta_trace( taus = mean_alphas.T * geometry_dist_to_next_depth_point / np.cos(theta) no_of_depth_gaps = len(geometry_dist_to_next_depth_point) - ###TODO: Generalize this for source functions other than blackbody that may require args other than frequency and temperature - source = source_function(tracing_nus, temps) - delta_source = source[1:] - source[:-1] + bb = bb_nu(tracing_nus, temps) + source = bb + delta_source = bb[1:] - bb[:-1] I_nu_theta = np.ones((no_of_depth_gaps + 1, len(tracing_nus))) * np.nan - I_nu_theta[0] = source[0] # the innermost depth point is the photosphere + I_nu_theta[0] = bb[0] # the innermost depth point is the photosphere for i in range(len(tracing_nus)): # iterating over nus (columns) for j in range(no_of_depth_gaps): # iterating over depth_gaps (rows) curr_tau = taus[i, j] + next_tau = taus[i, j + 1] w0, w1, w2 = calc_weights(curr_tau) @@ -90,7 +119,6 @@ def single_theta_trace( else: second_term = w1 * delta_source[j, i] / curr_tau if j < no_of_depth_gaps - 1: - next_tau = taus[i, j + 1] third_term = w2 * ( ( (delta_source[j + 1, i] / next_tau) @@ -114,7 +142,7 @@ def single_theta_trace( return I_nu_theta -def raytrace(stellar_model, stellar_radiation_field, no_of_thetas=20): +def raytrace(stellar_model, alphas, tracing_nus, no_of_thetas=20): """ Raytraces over many angles and integrates to get flux using the midpoint rule. @@ -122,8 +150,11 @@ def raytrace(stellar_model, stellar_radiation_field, no_of_thetas=20): Parameters ---------- stellar_model : stardis.model.base.StellarModel - stellar_radiation_field : stardis.radiation_field.base.StellarRadiationField - Contains temperatures, frequencies, and opacities needed to calculate F_nu. + alphas : numpy.ndarray + Array of shape (no_of_depth_points, no_of_frequencies). Total opacity at + each depth point for each frequency in tracing_nus. + tracing_nus : astropy.unit.quantity.Quantity + Numpy array of frequencies used for ray tracing with units of Hz. no_of_thetas : int, optional Number of angles to sample for ray tracing, by default 20. @@ -138,17 +169,16 @@ def raytrace(stellar_model, stellar_radiation_field, no_of_thetas=20): start_theta = dtheta / 2 end_theta = (np.pi / 2) - (dtheta / 2) thetas = np.linspace(start_theta, end_theta, no_of_thetas) + F_nu = np.zeros((len(stellar_model.geometry.r), len(tracing_nus))) - ###TODO: Thetas should probably be held by the model? Then can be passed in from there. for theta in thetas: weight = 2 * np.pi * dtheta * np.sin(theta) * np.cos(theta) - stellar_radiation_field.F_nu += weight * single_theta_trace( + F_nu += weight * single_theta_trace( stellar_model.geometry.dist_to_next_depth_point, stellar_model.temperatures.value.reshape(-1, 1), - stellar_radiation_field.opacities.total_alphas, - stellar_radiation_field.frequencies, + alphas, + tracing_nus, theta, - stellar_radiation_field.source_function, ) - return stellar_radiation_field.F_nu + return F_nu