Add wrapped cuda implementation of calc_doppler_width

stardis/opacities/broadening.py

@@ -65,6 +65,37 @@ def _calc_doppler_width_cuda(res, nu_line, temperature, atomic_mass):
         res[tid] = _calc_doppler_width(nu_line[tid], temperature[tid], atomic_mass[tid])
 
 
+def calc_doppler_width_cuda(
+    nu_line,
+    temperature,
+    atomic_mass,
+    nthreads=256,
+    ret_np_ndarray=True,
+    dtype=float,
+):
+    arg_list = (
+        nu_line,
+        temperature,
+        atomic_mass,
+    )
+
+    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_doppler_width_cuda[nblocks, nthreads](
+        res,
+        *arg_list,
+    )
+
+    return cp.asnumpy(res) if ret_np_ndarray else res
+
+
 @numba.njit
 def calc_n_effective(ion_number, ionization_energy, level_energy):
     """

stardis/opacities/tests/test_broadening.py

@@ -3,7 +3,11 @@
 from astropy import constants as const
 from numba import cuda
 
-from stardis.opacities.broadening import calc_doppler_width, _calc_doppler_width_cuda
+from stardis.opacities.broadening import (
+    calc_doppler_width,
+    _calc_doppler_width_cuda,
+    calc_doppler_width_cuda,
+)
 
 GPUs_available = cuda.is_available()
 
@@ -93,3 +97,34 @@ def test_calc_doppler_width_cuda_unwrapped_sample_values(
         cp.asnumpy(result_values),
         calc_doppler_width_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_doppler_width_cuda_sample_values_input_nu_line, calc_doppler_width_cuda_sample_values_input_temperature, calc_doppler_width_cuda_sample_values_input_atomic_mass, calc_doppler_width_cuda_wrapped_sample_cuda_values_expected_result",
+    [
+        (
+            np.array(2 * [SPEED_OF_LIGHT]),
+            np.array(2 * [0.5]),
+            np.array(2 * [BOLTZMANN_CONSTANT]),
+            np.array(2 * [1.0]),
+        ),
+    ],
+)
+def test_calc_doppler_width_cuda_wrapped_sample_cuda_values(
+    calc_doppler_width_cuda_sample_values_input_nu_line,
+    calc_doppler_width_cuda_sample_values_input_temperature,
+    calc_doppler_width_cuda_sample_values_input_atomic_mass,
+    calc_doppler_width_cuda_wrapped_sample_cuda_values_expected_result,
+):
+    arg_list = (
+        calc_doppler_width_cuda_sample_values_input_nu_line,
+        calc_doppler_width_cuda_sample_values_input_temperature,
+        calc_doppler_width_cuda_sample_values_input_atomic_mass,
+    )
+    assert np.allclose(
+        calc_doppler_width_cuda(*map(cp.asarray, arg_list)),
+        calc_doppler_width_cuda_wrapped_sample_cuda_values_expected_result,
+    )