Add CUDA JIT to calc_gamma_linear_stark() (#125)

* Add test for `calc_gamma_linar_stark()`

* Typecast inputs and constants

* Add vectorized version of `calc_gamma_linear_stark()`

* Add `_calc_gamma_linear_stark_cuda()`

* Add wrapped CUDA implementation of `calc_gamma_linear_stark()`

* Remove debugging print statement

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

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,
+    )