parallel support for line opacity calculation

stardis/base.py

@@ -102,6 +102,7 @@ def run_stardis(config_fname, tracing_lambdas_or_nus):
         stellar_model=stellar_model,
         stellar_radiation_field=stellar_radiation_field,
         opacity_config=config.opacity,
+        parallel_config=config.parallel,
     )
 
     raytrace(

stardis/radiation_field/opacities/opacities_solvers/base.py

@@ -322,10 +322,7 @@ def gaunt_times_departure(tracing_nus, temperatures, gaunt_fpath, departure_fpat
 
 # line opacity
 def calc_alpha_line_at_nu(
-    stellar_plasma,
-    stellar_model,
-    tracing_nus,
-    line_opacity_config,
+    stellar_plasma, stellar_model, tracing_nus, line_opacity_config, parallel_config
 ):
     """
     Calculates line opacity.
@@ -415,22 +412,16 @@ def calc_alpha_line_at_nu(
         line_opacity_config.broadening,
     )  # This can be further improved by only calculating the broadening for the lines that are within the range.
 
-    alpha_line_at_nu = np.zeros((stellar_model.no_of_depth_points, len(tracing_nus)))
+    delta_nus = tracing_nus.value - line_nus[:, np.newaxis]
 
     # If no broadening range, compute the contribution of every line at every frequency.
-    if line_range is None:
-        for i, nu in enumerate(tracing_nus):
-            delta_nus = nu.value - line_nus
-
-            alpha_line_at_nu[:, i] = calc_alan_entries(
-                delta_nus[:, np.newaxis],
-                doppler_widths,
-                gammas,
-                alphas_array,
-            )
+    h_lines_indices = None
+    line_range_value = None
 
     # If there is a broadening range, first make sure the range is in frequency units, and then iterate through each frequency to calculate the contribution of each line within the broadening range.
-    else:  # This if statement block appropriately handles if the broadening range is in frequency or wavelength units.
+    if (
+        line_range is not None
+    ):  # This if statement block appropriately handles if the broadening range is in frequency or wavelength units.
         h_lines_indices = (
             lines_sorted_in_range.atomic_number == 1
         ).to_numpy()  # Hydrogen lines are much broader than other lines, so they need special treatment to ignore the broadening range.
@@ -448,29 +439,116 @@ def calc_alpha_line_at_nu(
                 "Broadening range must be in units of length or frequency."
             )
 
-        # Iterate through each frequency to calculate the contribution of each line within the broadening range.
-        for i, nu in enumerate(tracing_nus):
-            delta_nus = nu.value - line_nus
+    # Iterate through each frequency to calculate the contribution of each line within the broadening range.
+    if parallel_config:
+        alpha_line_at_nu = calc_alan_entries_parallel(
+            stellar_model.no_of_depth_points,
+            tracing_nus.value,
+            delta_nus,
+            doppler_widths,
+            gammas,
+            alphas_array,
+            line_range_value,
+            h_lines_indices,
+        )
+
+    else:
+        alpha_line_at_nu = calc_alan_entries(
+            stellar_model.no_of_depth_points,
+            tracing_nus.value,
+            delta_nus,
+            doppler_widths,
+            gammas,
+            alphas_array,
+            line_range_value,
+            h_lines_indices,
+        )
+
+    return alpha_line_at_nu, gammas, doppler_widths
+
+
+@numba.njit(parallel=True)
+def calc_alan_entries_parallel(
+    no_of_depth_points,
+    tracing_nus_values,
+    delta_nus,
+    doppler_widths,
+    gammas,
+    alphas_array,
+    broadening_range=None,
+    h_lines_indices=None,
+):
+    alpha_line_at_nu = np.zeros((no_of_depth_points, len(tracing_nus_values)))
+
+    if broadening_range is None:
+        for frequency_index in numba.prange(len(tracing_nus_values)):
+            alpha_line_at_nu[:, frequency_index] = _calc_alan_entries(
+                delta_nus[:, frequency_index, np.newaxis],
+                doppler_widths,
+                gammas,
+                alphas_array,
+            )
 
-            broadening_mask = np.abs(delta_nus) < line_range_value[i]
+    else:
+        for frequency_index in numba.prange(len(tracing_nus_values)):
+            broadening_mask = (
+                np.abs(delta_nus[:, frequency_index])
+                < broadening_range[frequency_index]
+            )
             broadening_mask = np.logical_or(broadening_mask, h_lines_indices)
 
-            delta_nus_considered = delta_nus[broadening_mask]
-            gammas_considered = gammas[broadening_mask, :]
-            doppler_widths_considered = doppler_widths[broadening_mask, :]
-            alphas_considered = alphas_array[broadening_mask, :]
-            alpha_line_at_nu[:, i] = calc_alan_entries(
-                delta_nus_considered[:, np.newaxis],
-                doppler_widths_considered,
-                gammas_considered,
-                alphas_considered,
+            alpha_line_at_nu[:, frequency_index] = _calc_alan_entries(
+                delta_nus[:, frequency_index, np.newaxis][broadening_mask],
+                doppler_widths[broadening_mask],
+                gammas[broadening_mask],
+                alphas_array[broadening_mask],
             )
 
-    return alpha_line_at_nu, gammas, doppler_widths
+    return alpha_line_at_nu
 
 
 @numba.njit
 def calc_alan_entries(
+    no_of_depth_points,
+    tracing_nus_values,
+    delta_nus,
+    doppler_widths,
+    gammas,
+    alphas_array,
+    broadening_range=None,
+    h_lines_indices=None,
+):
+    alpha_line_at_nu = np.zeros((no_of_depth_points, len(tracing_nus_values)))
+
+    if broadening_range is None:
+        for frequency_index in range(len(tracing_nus_values)):
+            alpha_line_at_nu[:, frequency_index] = _calc_alan_entries(
+                delta_nus[:, frequency_index, np.newaxis],
+                doppler_widths,
+                gammas,
+                alphas_array,
+            )
+
+    else:
+        for frequency_index in range(len(tracing_nus_values)):
+            broadening_mask = (
+                np.abs(delta_nus[:, frequency_index])
+                < broadening_range[frequency_index]
+            )
+            broadening_mask = np.logical_or(broadening_mask, h_lines_indices)
+
+            alpha_line_at_nu[:, frequency_index] = _calc_alan_entries(
+                delta_nus[:, frequency_index, np.newaxis][broadening_mask],
+                doppler_widths[broadening_mask],
+                gammas[broadening_mask],
+                alphas_array[broadening_mask],
+            )
+
+    return alpha_line_at_nu
+
+
+@numba.njit
+def _calc_alan_entries(
     delta_nus,
     doppler_widths_at_depth_point,
     gammas_at_depth_point,
@@ -497,9 +575,8 @@ def calc_alan_entries(
     float
         Line opacity.
     """
-
     phis = voigt_profile(
-        np.abs(delta_nus), doppler_widths_at_depth_point, gammas_at_depth_point
+        delta_nus, doppler_widths_at_depth_point, gammas_at_depth_point
     )
 
     return np.sum(phis * alphas_at_depth_point, axis=0)
@@ -510,6 +587,7 @@ def calc_alphas(
     stellar_model,
     stellar_radiation_field,
     opacity_config,
+    parallel_config,
 ):
     """
     Calculates each opacity and adds it to the opacity dictionary contained in the radiation field.
@@ -582,6 +660,7 @@ def calc_alphas(
         stellar_model,
         stellar_radiation_field.frequencies,
         opacity_config.line,
+        parallel_config,
     )
     stellar_radiation_field.opacities.opacities_dict[
         "alpha_line_at_nu"