diff --git a/stardis/radiation_field/opacities/opacities_solvers/broadening.py b/stardis/radiation_field/opacities/opacities_solvers/broadening.py index 689617c9..65f19b62 100644 --- a/stardis/radiation_field/opacities/opacities_solvers/broadening.py +++ b/stardis/radiation_field/opacities/opacities_solvers/broadening.py @@ -176,7 +176,7 @@ def calc_n_effective_cuda( @numba.njit -def calc_gamma_linear_stark(n_eff_upper, n_eff_lower, electron_density): +def _calc_gamma_linear_stark(n_eff_upper, n_eff_lower, electron_density): """ Calculates broadening parameter for linear Stark broadening. https://ui.adsabs.harvard.edu/abs/1978JQSRT..20..333S/ @@ -196,21 +196,73 @@ def calc_gamma_linear_stark(n_eff_upper, n_eff_lower, electron_density): Broadening parameter for linear Stark broadening. """ - if n_eff_upper - n_eff_lower < 1.5: - a1 = 0.642 - else: - a1 = 1 + n_eff_upper, n_eff_lower, electron_density = ( + float(n_eff_upper), + float(n_eff_lower), + float(electron_density), + ) + + a1 = 0.642 if (n_eff_upper - n_eff_lower < 1.5) else 1.0 gamma_linear_stark = ( 0.51 * a1 * (n_eff_upper**2 - n_eff_lower**2) - * (electron_density ** (2 / 3)) + * (electron_density ** (2.0 / 3.0)) ) return gamma_linear_stark +@numba.vectorize(nopython=True) +def calc_gamma_linear_stark(n_eff_upper, n_eff_lower, electron_density): + return _calc_gamma_linear_stark(n_eff_upper, n_eff_lower, electron_density) + + +@cuda.jit +def _calc_gamma_linear_stark_cuda(res, n_eff_upper, n_eff_lower, electron_density): + tid = cuda.grid(1) + size = len(res) + + if tid < size: + res[tid] = _calc_gamma_linear_stark( + n_eff_upper[tid], + n_eff_lower[tid], + electron_density[tid], + ) + + +def calc_gamma_linear_stark_cuda( + n_eff_upper, + n_eff_lower, + electron_density, + nthreads=256, + ret_np_ndarray=False, + dtype=float, +): + arg_list = ( + n_eff_upper, + n_eff_lower, + electron_density, + ) + + 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) + + _calc_gamma_linear_stark_cuda[nblocks, nthreads]( + res, + *arg_list, + ) + + return cp.asnumpy(res) if ret_np_ndarray else res + + @numba.njit def calc_gamma_quadratic_stark( ion_number, n_eff_upper, n_eff_lower, electron_density, temperature diff --git a/stardis/radiation_field/opacities/tests/test_broadening.py b/stardis/radiation_field/opacities/tests/test_broadening.py index ffc27185..c05a1db7 100644 --- a/stardis/radiation_field/opacities/tests/test_broadening.py +++ b/stardis/radiation_field/opacities/tests/test_broadening.py @@ -10,6 +10,9 @@ calc_n_effective, _calc_n_effective_cuda, calc_n_effective_cuda, + calc_gamma_linear_stark, + _calc_gamma_linear_stark_cuda, + calc_gamma_linear_stark_cuda, ) GPUs_available = cuda.is_available() @@ -236,3 +239,108 @@ def test_calc_n_effective_cuda_wrapped_sample_cuda_values( calc_n_effective_cuda(*map(cp.asarray, arg_list)), calc_n_effective_cuda_wrapped_sample_cuda_values_expected_result, ) + + +@pytest.mark.parametrize( + "calc_gamma_linear_stark_sample_values_input_n_eff_upper,calc_gamma_linear_stark_sample_values_input_n_eff_lower,calc_gamma_linear_stark_sample_values_input_electron_density, calc_gamma_linear_stark_sample_values_expected_result", + [ + ( + 1, + 0, + (0.51 * 0.642) ** (-3 / 2), + 1.0, + ), + ( + np.array(2 * [1]), + np.array(2 * [0]), + np.array(2 * [(0.51 * 0.642) ** (-3 / 2)]), + np.array(2 * [1.0]), + ), + ], +) +def test_calc_gamma_linear_stark_sample_values( + calc_gamma_linear_stark_sample_values_input_n_eff_upper, + calc_gamma_linear_stark_sample_values_input_n_eff_lower, + calc_gamma_linear_stark_sample_values_input_electron_density, + calc_gamma_linear_stark_sample_values_expected_result, +): + assert np.allclose( + calc_gamma_linear_stark( + calc_gamma_linear_stark_sample_values_input_n_eff_upper, + calc_gamma_linear_stark_sample_values_input_n_eff_lower, + calc_gamma_linear_stark_sample_values_input_electron_density, + ), + calc_gamma_linear_stark_sample_values_expected_result, + ) + + +@pytest.mark.skipif( + not GPUs_available, reason="No GPU is available to test CUDA function" +) +@pytest.mark.parametrize( + "calc_gamma_linear_stark_cuda_unwrapped_sample_values_input_n_eff_upper,calc_gamma_linear_stark_cuda_unwrapped_sample_values_input_n_eff_lower,calc_gamma_linear_stark_cuda_unwrapped_sample_values_input_electron_density,calc_gamma_linear_stark_cuda_unwrapped_sample_values_expected_result", + [ + ( + np.array(2 * [1]), + np.array(2 * [0]), + np.array(2 * [(0.51 * 0.642) ** (-3 / 2)]), + np.array(2 * [1.0]), + ), + ], +) +def test_calc_gamma_linear_stark_cuda_unwrapped_sample_values( + calc_gamma_linear_stark_cuda_unwrapped_sample_values_input_n_eff_upper, + calc_gamma_linear_stark_cuda_unwrapped_sample_values_input_n_eff_lower, + calc_gamma_linear_stark_cuda_unwrapped_sample_values_input_electron_density, + calc_gamma_linear_stark_cuda_unwrapped_sample_values_expected_result, +): + arg_list = ( + calc_gamma_linear_stark_cuda_unwrapped_sample_values_input_n_eff_upper, + calc_gamma_linear_stark_cuda_unwrapped_sample_values_input_n_eff_lower, + calc_gamma_linear_stark_cuda_unwrapped_sample_values_input_electron_density, + ) + + arg_list = tuple(map(cp.array, arg_list)) + result_values = cp.empty_like(arg_list[0], dtype=float) + + nthreads = 256 + length = len(calc_gamma_linear_stark_cuda_unwrapped_sample_values_expected_result) + nblocks = 1 + (length // nthreads) + + _calc_gamma_linear_stark_cuda[nblocks, nthreads](result_values, *arg_list) + + assert np.allclose( + cp.asnumpy(result_values), + calc_gamma_linear_stark_cuda_unwrapped_sample_values_expected_result, + ) + + +@pytest.mark.skipif( + not GPUs_available, reason="No GPU is available to test CUDA function" +) +@pytest.mark.parametrize( + "calc_gamma_linear_stark_cuda_wrapped_sample_values_input_n_eff_upper,calc_gamma_linear_stark_cuda_wrapped_sample_values_input_n_eff_lower,calc_gamma_linear_stark_cuda_wrapped_sample_values_input_electron_density,calc_gamma_linear_stark_cuda_wrapped_sample_values_expected_result", + [ + ( + np.array(2 * [1]), + np.array(2 * [0]), + np.array(2 * [(0.51 * 0.642) ** (-3 / 2)]), + np.array(2 * [1.0]), + ), + ], +) +def test_calc_doppler_width_cuda_wrapped_sample_cuda_values( + calc_gamma_linear_stark_cuda_wrapped_sample_values_input_n_eff_upper, + calc_gamma_linear_stark_cuda_wrapped_sample_values_input_n_eff_lower, + calc_gamma_linear_stark_cuda_wrapped_sample_values_input_electron_density, + calc_gamma_linear_stark_cuda_wrapped_sample_values_expected_result, +): + arg_list = ( + calc_gamma_linear_stark_cuda_wrapped_sample_values_input_n_eff_upper, + calc_gamma_linear_stark_cuda_wrapped_sample_values_input_n_eff_lower, + calc_gamma_linear_stark_cuda_wrapped_sample_values_input_electron_density, + ) + assert np.allclose( + calc_gamma_linear_stark_cuda(*map(cp.asarray, arg_list)), + calc_gamma_linear_stark_cuda_wrapped_sample_values_expected_result, + )