diff --git a/stardis/opacities/tests/test_broadening.py b/stardis/opacities/tests/test_broadening.py index fd6cf510..f2aead72 100644 --- a/stardis/opacities/tests/test_broadening.py +++ b/stardis/opacities/tests/test_broadening.py @@ -7,6 +7,7 @@ calc_doppler_width, _calc_doppler_width_cuda, calc_doppler_width_cuda, + calc_n_effective, ) GPUs_available = cuda.is_available() @@ -128,3 +129,36 @@ def test_calc_doppler_width_cuda_wrapped_sample_cuda_values( calc_doppler_width_cuda(*map(cp.asarray, arg_list)), calc_doppler_width_cuda_wrapped_sample_cuda_values_expected_result, ) + + +@pytest.mark.parametrize( + "calc_n_effective_sample_values_input_ion_number,calc_n_effective_sample_values_input_ionization_energy,calc_n_effective_sample_values_input_level_energy, calc_n_effective_sample_values_expected_result", + [ + ( + 1.0, + RYDBERG_ENERGY, + 0, + 1.0, + ), + ( + np.array(2 * [1.0]), + np.array(2 * [RYDBERG_ENERGY]), + np.array(2 * [0]), + np.array(2 * [1.0]), + ), + ], +) +def test_calc_n_effective_sample_values( + calc_n_effective_sample_values_input_ion_number, + calc_n_effective_sample_values_input_ionization_energy, + calc_n_effective_sample_values_input_level_energy, + calc_n_effective_sample_values_expected_result, +): + assert np.allclose( + calc_n_effective( + calc_n_effective_sample_values_input_ion_number, + calc_n_effective_sample_values_input_ionization_energy, + calc_n_effective_sample_values_input_level_energy, + ), + calc_n_effective_sample_values_expected_result, + )