Line opacity parallelization

stardis/base.py

@@ -14,6 +14,7 @@
 from stardis.io.model.marcs import read_marcs_model
 from stardis.io.model.mesa import read_mesa_model
 from stardis.radiation_field.source_functions.blackbody import blackbody_flux_at_nu
+import logging
 
 
 BASE_DIR = Path(__file__).parent
@@ -48,6 +49,7 @@ def run_stardis(config_fname, tracing_lambdas_or_nus):
     adata = AtomData.from_hdf(config.atom_data)
 
     # model
+    logging.info("Reading model")
     if config.model.type == "marcs":
         raw_marcs_model = read_marcs_model(
             Path(config.model.fname), gzipped=config.model.gzipped
@@ -91,19 +93,21 @@ def run_stardis(config_fname, tracing_lambdas_or_nus):
         continuum_interaction_species=[],
     )
     # plasma
+    logging.info("Creating plasma")
     stellar_plasma = create_stellar_plasma(stellar_model, adata, config)
 
     stellar_radiation_field = RadiationField(
         tracing_nus, blackbody_flux_at_nu, stellar_model
     )
-
+    logging.info("Calculating alphas")
     calc_alphas(
         stellar_plasma=stellar_plasma,
         stellar_model=stellar_model,
         stellar_radiation_field=stellar_radiation_field,
         opacity_config=config.opacity,
+        parallel_config=config.parallel,
     )
-
+    logging.info("Raytracing")
     raytrace(
         stellar_model,
         stellar_radiation_field,

stardis/conftest.py

@@ -157,6 +157,7 @@ def example_stellar_radiation_field_parallel(
         stellar_model=example_stellar_model,
         stellar_radiation_field=stellar_radiation_field,
         opacity_config=example_config_parallel.opacity,
+        parallel_config=example_config_parallel.parallel,
     )
 
     raytrace(

stardis/radiation_field/opacities/opacities_solvers/base.py

@@ -326,6 +326,7 @@
     stellar_model,
     tracing_nus,
     line_opacity_config,
+    parallel_config=False,
 ):
     """
     Calculates line opacity.
@@ -415,22 +416,16 @@
         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 +443,116 @@
                 "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
 
-            broadening_mask = np.abs(delta_nus) < line_range_value[i]
+
+@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,
+            )
+
+    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 +579,8 @@
     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 +591,7 @@
     stellar_model,
     stellar_radiation_field,
     opacity_config,
+    parallel_config=False,
 ):
     """
     Calculates each opacity and adds it to the opacity dictionary contained in the radiation field.
@@ -582,6 +664,7 @@
         stellar_model,
         stellar_radiation_field.frequencies,
         opacity_config.line,
+        parallel_config,
     )
     stellar_radiation_field.opacities.opacities_dict[
         "alpha_line_at_nu"