From 487a065b4fc9773764162d5d3cfcf4092cfee67d Mon Sep 17 00:00:00 2001
From: William Black <125844868+smokestacklightnin@users.noreply.github.com>
Date: Wed, 9 Aug 2023 17:03:34 -0700
Subject: [PATCH] Refactor `calc_doppler_width()` and add test for vectorized
 implementation

---
 stardis/opacities/broadening.py            | 12 ++++++++----
 stardis/opacities/tests/test_broadening.py | 20 +++++++++++++++-----
 2 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/stardis/opacities/broadening.py b/stardis/opacities/broadening.py
index b618abcc..598971fe 100644
--- a/stardis/opacities/broadening.py
+++ b/stardis/opacities/broadening.py
@@ -1,5 +1,6 @@
 import numpy as np
 from astropy import constants as const
+import math
 import numba
 
 
@@ -13,7 +14,7 @@
 
 
 @numba.njit
-def calc_doppler_width(nu_line, temperature, atomic_mass):
+def _calc_doppler_width(nu_line, temperature, atomic_mass):
     """
     Calculates doppler width.
     https://ui.adsabs.harvard.edu/abs/2003rtsa.book.....R/
@@ -34,10 +35,15 @@ def calc_doppler_width(nu_line, temperature, atomic_mass):
     return (
         nu_line
         / SPEED_OF_LIGHT
-        * np.sqrt(2 * BOLTZMANN_CONSTANT * temperature / atomic_mass)
+        * math.sqrt(2 * BOLTZMANN_CONSTANT * temperature / atomic_mass)
     )
 
 
+@numba.vectorize(nopython=True)
+def calc_doppler_width(nu_line, temperature, atomic_mass):
+    return _calc_doppler_width(nu_line, temperature, atomic_mass)
+
+
 @numba.njit
 def calc_n_effective(ion_number, ionization_energy, level_energy):
     """
@@ -365,7 +371,6 @@ def calculate_broadening(
     h_mass = atomic_masses[0]
 
     for i in range(len(lines_array)):
-
         atomic_number = int(lines_array[i, line_cols["atomic_number"]])
         atomic_mass = atomic_masses[atomic_number - 1]
         ion_number = int(lines_array[i, line_cols["ion_number"]]) + 1
@@ -378,7 +383,6 @@ def calculate_broadening(
         line_nus[i] = line_nu
 
         for j in range(no_shells):
-
             electron_density = electron_densities[j]
             temperature = temperatures[j]
             h_density = h_densities[j]
diff --git a/stardis/opacities/tests/test_broadening.py b/stardis/opacities/tests/test_broadening.py
index ead1611a..1bb5cd81 100644
--- a/stardis/opacities/tests/test_broadening.py
+++ b/stardis/opacities/tests/test_broadening.py
@@ -1,10 +1,10 @@
 import pytest
-from numpy import allclose, pi as PI
+import numpy as np
 from astropy import constants as const
 
 from stardis.opacities.broadening import calc_doppler_width
 
-
+PI = np.pi
 SPEED_OF_LIGHT = const.c.cgs.value
 BOLTZMANN_CONSTANT = const.k_B.cgs.value
 PLANCK_CONSTANT = const.h.cgs.value
@@ -17,8 +17,18 @@
 @pytest.mark.parametrize(
     "calc_doppler_width_sample_values_input_nu_line,calc_doppler_width_sample_values_input_temperature,calc_doppler_width_sample_values_input_atomic_mass, calc_doppler_width_sample_values_expected_result",
     [
-        (SPEED_OF_LIGHT, 0.5, BOLTZMANN_CONSTANT, 1.0),
-        # (0.0, 1.0 + 0.0j),
+        (
+            SPEED_OF_LIGHT,
+            0.5,
+            BOLTZMANN_CONSTANT,
+            1.0,
+        ),
+        (
+            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_sample_values(
@@ -27,7 +37,7 @@ def test_calc_doppler_width_sample_values(
     calc_doppler_width_sample_values_input_atomic_mass,
     calc_doppler_width_sample_values_expected_result,
 ):
-    assert allclose(
+    assert np.allclose(
         calc_doppler_width(
             calc_doppler_width_sample_values_input_nu_line,
             calc_doppler_width_sample_values_input_temperature,