diff --git a/stardis/base.py b/stardis/base.py index d4ae3c88..5d6941e0 100644 --- a/stardis/base.py +++ b/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( diff --git a/stardis/radiation_field/opacities/opacities_solvers/base.py b/stardis/radiation_field/opacities/opacities_solvers/base.py index ec13c69c..dd1eb969 100644 --- a/stardis/radiation_field/opacities/opacities_solvers/base.py +++ b/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"