From e151a7bdd81357f508bef59c1f7f1bf8374883be Mon Sep 17 00:00:00 2001 From: Joshua Shields <54691495+jvshields@users.noreply.github.com> Date: Tue, 8 Oct 2024 11:43:37 -0400 Subject: [PATCH] Molecules (#210) * fix plasma * use aug 25 release * molecules checkpoint * fix logging * spectrum checkpoint - probably incorrect * separate out partition functions * update docstrings * clearer variable names * slow checkpoint * optimization * add comments * add shortlist handling * densities -> density * shortlist handling bugfix * use mapping for ions instead of holding in densities df * give value error if no molecular data provided * docstrings * clearer row iteration * explain alpha * add preprocessing function * make preprocess function return processed dataframe --- stardis/base.py | 8 +- stardis/config_schema.yml | 4 + stardis/io/base.py | 6 +- stardis/plasma/base.py | 30 +- stardis/plasma/molecules.py | 421 ++++++++++++++++++ stardis/radiation_field/base.py | 6 +- .../opacities/opacities_solvers/base.py | 111 +++++ .../opacities/opacities_solvers/broadening.py | 28 ++ 8 files changed, 602 insertions(+), 12 deletions(-) create mode 100644 stardis/plasma/molecules.py diff --git a/stardis/base.py b/stardis/base.py index 7016768c..0754e325 100644 --- a/stardis/base.py +++ b/stardis/base.py @@ -7,6 +7,8 @@ import logging +logger = logging.getLogger(__name__) + def run_stardis( config_fname, tracing_lambdas_or_nus, add_config_keys=None, add_config_vals=None @@ -73,11 +75,11 @@ def set_num_threads(n_threads): """ if n_threads == 1: - logging.info("Running in serial mode") + logger.info("Running in serial mode") elif n_threads == -99: - logging.info("Running with max threads") + logger.info("Running with max threads") elif n_threads > 1: - logging.info(f"Running with {n_threads} threads") + logger.info(f"Running with {n_threads} threads") numba.set_num_threads(n_threads) else: raise ValueError( diff --git a/stardis/config_schema.yml b/stardis/config_schema.yml index dc4d9f58..f20d4b6b 100644 --- a/stardis/config_schema.yml +++ b/stardis/config_schema.yml @@ -126,6 +126,10 @@ properties: type: boolean default: false description: Options for whether to use a vald linelist. Linelist must be included the atom_data and cannot be supplied separately. + include_molecules: + type: boolean + default: false + description: Whether to include molecular lines in the simulation. #required: #- line no_of_thetas: diff --git a/stardis/io/base.py b/stardis/io/base.py index a4ba029f..ba081689 100644 --- a/stardis/io/base.py +++ b/stardis/io/base.py @@ -14,6 +14,8 @@ BASE_DIR = Path(__file__).parent.parent SCHEMA_PATH = BASE_DIR / "config_schema.yml" +logger = logging.getLogger(__name__) + def parse_config_to_model(config_fname, add_config_keys=None, add_config_vals=None): """ @@ -49,7 +51,7 @@ def parse_config_to_model(config_fname, add_config_keys=None, add_config_vals=No ): # If a dictionary was passed, update the config with the dictionary pass else: - print("Updating config with additional keys and values") + logger.info("Updating config with additional keys and values") if isinstance(add_config_keys, str): # Directly set the config item if add_config_keys is a string config.set_config_item(add_config_keys, add_config_vals) @@ -75,7 +77,7 @@ def parse_config_to_model(config_fname, add_config_keys=None, add_config_vals=No adata = AtomData.from_hdf(config.atom_data) # model - logging.info("Reading model") + logger.info("Reading model") if config.model.type == "marcs": raw_marcs_model = read_marcs_model( Path(config.model.fname), gzipped=config.model.gzipped diff --git a/stardis/plasma/base.py b/stardis/plasma/base.py index edc73696..017a216e 100644 --- a/stardis/plasma/base.py +++ b/stardis/plasma/base.py @@ -19,6 +19,13 @@ non_nlte_properties, helium_lte_properties, ) +from stardis.plasma.molecules import ( + MoleculeIonNumberDensity, + AlphaLineValdMolecule, + MoleculePartitionFunction, + AlphaLineShortlistValdMolecule, +) + import tardis.plasma from tardis.opacities.tau_sobolev import TauSobolev @@ -50,6 +57,8 @@ 25200, ] # see directly above +logger = logging.getLogger(__name__) + class HMinusDensity(ProcessingPlasmaProperty): """ @@ -122,7 +131,7 @@ class AlphaLine(ProcessingPlasmaProperty): Attributes ---------- alpha_line : Pandas DataFrame, dtype float - Sobolev optical depth for each line. Indexed by line. + Calculates the alpha (line integrated absorption coefficient in cm^-1) values for each line at each depth point. Indexed by line. Columns as zones. """ @@ -167,6 +176,7 @@ def calculate( class AlphaLineVald(ProcessingPlasmaProperty): """ + Calculates the alpha (line integrated absorption coefficient in cm^-1) values for each line at each depth point. Uses VALD linelists for lines. Indexed by line. Attributes ---------- alpha_line_from_linelist : DataFrame @@ -200,7 +210,6 @@ def calculate( # alphas = ALPHA_COEFFICIENT * n_lower * f_lu * emission_correction ###TODO: handle other broadening parameters - points = len(t_electrons) linelist = atomic_data.linelist_atoms.rename( @@ -301,8 +310,8 @@ def calculate( linelist["level_energy_upper"] = ((linelist["e_up"].values * u.eV).cgs).value # Radiation broadening parameter is approximated as the einstein A coefficient. Vald parameters are in log scale. - linelist["A_ul"] = 10 ** ( - linelist["rad"] + linelist["A_ul"] = 10 ** (linelist["rad"]) / ( + 4 * np.pi ) # see 1995A&AS..112..525P for appropriate units - may be off by a factor of 4pi # Need to remove autoionization lines - can't handle with current broadening treatment because can't calculate effective principal quantum number @@ -313,6 +322,9 @@ def calculate( class AlphaLineShortlistVald(ProcessingPlasmaProperty): """ + Calculates the alpha (line integrated absorption coefficient in cm^-1) values for each line at each depth point. Uses VALD shortform linelists for lines. Indexed by line. + Subtley different from the full list calculation in that it does not require the upper level degeneracy, and pure number density is not directly calculated as a result. + Attributes ---------- alpha_line_from_linelist : DataFrame @@ -496,7 +508,7 @@ def create_stellar_plasma( tardis.plasma.base.BasePlasma """ - logging.info("Creating plasma") + logger.info("Creating plasma") # basic_properties.remove(tardis.plasma.properties.general.NumberDensity) plasma_modules = [] @@ -538,6 +550,14 @@ def create_stellar_plasma( temperature=stellar_model.temperatures, dilution_factor=np.ones_like(stellar_model.temperatures), ) + if config.opacity.line.include_molecules: + plasma_modules.append(MoleculeIonNumberDensity) + plasma_modules.append(MoleculePartitionFunction) + if config.opacity.line.vald_linelist.use_linelist: + if config.opacity.line.vald_linelist.shortlist: + plasma_modules.append(AlphaLineShortlistValdMolecule) + else: + plasma_modules.append(AlphaLineValdMolecule) return BasePlasma( plasma_properties=plasma_modules, diff --git a/stardis/plasma/molecules.py b/stardis/plasma/molecules.py new file mode 100644 index 00000000..f7fb2e77 --- /dev/null +++ b/stardis/plasma/molecules.py @@ -0,0 +1,421 @@ +import numpy as np +import pandas as pd +import logging + +from astropy import constants as const, units as u +from tardis.util.base import element_symbol2atomic_number + +from tardis.plasma.properties.base import ProcessingPlasmaProperty + +ALPHA_COEFFICIENT = (np.pi * const.e.gauss**2) / (const.m_e.cgs * const.c.cgs) + +logger = logging.getLogger(__name__) + + +class MoleculeIonNumberDensity(ProcessingPlasmaProperty): + """ + + Calculates the number density of each molecule at each depth point using Barklem and Collet 2016 equilibrium constants. + Multiplies the number density of each constituent ion and divides by the number density constant. Negative ions are ignored. + Attributes + ---------- + molecule_number_density : DataFrame + A pandas DataFrame with dtype float. This represents the number density of each molecule at each depth point. + molecule_ion_map : DataFrame + A pandas DataFrame with the constituent ions of each molecule. Useful for calculating molecular masses later for doppler broadening. + """ + + # Need to think about negative ions - ignoring for now + # applicable for equilibrium constants given by Barklem and Collet 2016, which are given in SI units + + outputs = ("molecule_number_density", "molecule_ion_map") + + def calculate(self, ion_number_density, t_electrons, atomic_data): + # Preprocessing - split ions into symbol, charge, and number + try: + molecules_df = atomic_data.molecule_data.dissociation_energies.copy() + except: + raise ValueError( + "No molecular dissociation energies found in atomic data. Use Carsus to generate atomic data with the Barklem and Collet 2016 data." + ) + + for ion in [1, 2]: + molecules_df = self.preprocess_ion(molecules_df, ion) + + number_densities_arr = np.zeros( + (len(atomic_data.molecule_data.equilibrium_constants), len(t_electrons)) + ) + # Get the temperatures at which the equilibrium constants are given for interpolation + equilibrium_const_temps = ( + atomic_data.molecule_data.equilibrium_constants.columns.values + ) + included_elements = ion_number_density.index.get_level_values(0).unique() + + for ( + molecule, + molecule_row, + ) in ( + molecules_df.iterrows() + ): # Loop over all molecules, calculate number densities using Barklem and Collet 2016 equilibrium constants - if a component ion does not exist in the plasma or is negative, assume no molecule + if (molecule_row.Ion1_charge == -1) or (molecule_row.Ion2_charge == -1): + logger.warning( + f"Negative ionic molecules not currently supported. Assuming no {molecule}." + ) + continue + + elif molecule_row.Ion1 not in included_elements: + logger.warning( + f"{molecule_row.Ion1} not in included elements. Assuming no {molecule}." + ) + continue + elif molecule_row.Ion2 not in included_elements: + logger.warning( + f"{molecule_row.Ion2} not in included elements. Assuming no {molecule}." + ) + continue + + ion1_number_density = ion_number_density.loc[ + molecule_row.Ion1, molecule_row.Ion1_charge + ] + ion2_number_density = ion_number_density.loc[ + molecule_row.Ion2, molecule_row.Ion2_charge + ] + + # Barklem and Collet 2016 equilibrium constants are pressure constants and are in SI units + pressure_equilibirium_const_at_depth_point = np.interp( + t_electrons, + equilibrium_const_temps, + atomic_data.molecule_data.equilibrium_constants.loc[molecule].values, + ) + # Convert from pressure constants to number density constants using ideal gas law + equilibirium_const_at_depth_point = ( + ( + 10 ** (pressure_equilibirium_const_at_depth_point) + * (u.N / u.m**2) + / (const.k_B * t_electrons * u.K) + ) + .to(u.cm**-3) + .value + ) + + molecule_number_density = (ion1_number_density * ion2_number_density) / ( + equilibirium_const_at_depth_point + ) + + number_densities_arr[ + atomic_data.molecule_data.equilibrium_constants.index.get_loc(molecule) + ] = molecule_number_density + + molecule_densities_df = pd.DataFrame( + number_densities_arr, + index=atomic_data.molecule_data.equilibrium_constants.index, + columns=ion_number_density.columns, + ) + # Keep track of the individual ions - useful to calculate molecular masses later for doppler broadening + molecule_ion_map = pd.DataFrame( + molecules_df[["Ion1", "Ion2"]], + ) + + return molecule_densities_df, molecule_ion_map + + def preprocess_ion(self, molecules_df, ion): + """ + Preprocesses a component ion in the molecule data to add the ion's atomic number and charge to the df. + """ + molecules_df[ + [f"Ion{ion}_symbol", f"Ion{ion}_positive", f"Ion{ion}_negative"] + ] = molecules_df[f"Ion{ion}"].str.extract(r"([A-Z][a-z]?+)(\+*)(\-*)") + molecules_df[f"Ion{ion}"] = molecules_df[f"Ion{ion}_symbol"].apply( + element_symbol2atomic_number + ) + molecules_df[f"Ion{ion}_charge"] = molecules_df[f"Ion{ion}_positive"].apply( + len + ) - molecules_df[f"Ion{ion}_negative"].apply(len) + return molecules_df + + +class MoleculePartitionFunction(ProcessingPlasmaProperty): + """ + Processes the partition function for each molecule at each depth point by interpolating the partition function data. + From Barklem and Collet 2016. + Attributes + ---------- + molecule_partition_function : DataFrame + A pandas DataFrame with dtype float. This represents the partition function + for each molecule at each depth point. + """ + + outputs = ("molecule_partition_function",) + + def calculate(self, t_electrons, atomic_data): + partition_functions = pd.DataFrame( + np.zeros( + (len(atomic_data.molecule_data.partition_functions), len(t_electrons)) + ), + index=atomic_data.molecule_data.partition_functions.index, + ) + + for molecule in atomic_data.molecule_data.partition_functions.index: + partition_functions.loc[molecule] = np.interp( + t_electrons, + atomic_data.molecule_data.partition_functions.columns.values, + atomic_data.molecule_data.partition_functions.loc[molecule].values, + ) + + return partition_functions + + +class AlphaLineValdMolecule(ProcessingPlasmaProperty): + """ + Calculates the alpha (line integrated absorption coefficient in cm^-1) values for each molecular line from Vald at each depth point. This is adapted from the AlphaLineVald calculation. + Attributes + ---------- + molecule_alpha_line_from_linelist : DataFrame + A pandas DataFrame with dtype float. This represents the alpha calculation + for each line from Vald at each depth point. Refer to Rybicki and Lightman + equation 1.80. Voigt profiles are calculated later, and B_12 is substituted + appropriately out for f_lu. This assumes LTE for lower level population. + molecule_lines_from_linelist: DataFrame + A pandas dataframe containing the lines and information about the lines in + the same form and shape as alpha_line_from_linelist. + """ + + outputs = ("molecule_alpha_line_from_linelist", "molecule_lines_from_linelist") + latex_name = (r"\alpha_{\textrm{moleculeline, vald}}",) + latex_formula = ( + r"\dfrac{\pi e^{2} n_{lower} f_{lu}}{m_{e} c}\ + \Big(1-exp(-h \nu / k T) \phi(\nu)\Big)", + ) + + def calculate( + self, + atomic_data, + molecule_number_density, + t_electrons, + molecule_partition_function, + ): + # solve n_lower : n_i = N * g_i / U * e ^ (-E_i/kT) + # get f_lu : loggf -> use g = 2j+1 + # emission_correction = (1-e^(-h*nu / kT)) + # alphas = ALPHA_COEFFICIENT * n_lower * f_lu * emission_correction + + ###TODO: handle other broadening parameters + points = len(t_electrons) + + linelist = atomic_data.linelist_molecules[ + [ + "molecule", + "wavelength", + "log_gf", + "e_low", + "e_up", + "j_lo", + "j_up", + "rad", + "stark", + "waals", + ] + ].copy() + + # Calculate degeneracies + linelist["g_lo"] = linelist.j_lo * 2 + 1 + linelist["g_up"] = linelist.j_up * 2 + 1 + + exponent_by_point = np.exp( + np.outer( + -linelist.e_low.values * u.eV, 1 / (t_electrons * u.K * const.k_B) + ).to(1) + ) + + molecule_densities_div_partition_function = molecule_number_density.copy().div( + molecule_partition_function + ) + molecule_densities_div_partition_function.index.name = "molecule" + + # grab densities for n_lower - need to use linelist as the index and normalize by dividing by the partition function + linelist_with_density_div_partition_function = linelist.merge( + molecule_densities_div_partition_function, + how="left", + on=["molecule"], + ) + + n_lower = ( + ( + exponent_by_point + * linelist_with_density_div_partition_function[np.arange(points)] + ).values.T # arange mask of the dataframe returns the set of densities of the appropriate ion for the line at each point + * linelist_with_density_div_partition_function.g_lo.values + ) + + linelist["f_lu"] = ( + 10**linelist.log_gf / linelist.g_lo + ) # vald log gf is "oscillator strength f times the statistical weight g of the parent level" see 1995A&AS..112..525P, section 2. Structure of VALD + + line_nus = (linelist.wavelength.values * u.AA).to( + u.Hz, equivalencies=u.spectral() + ) + + emission_correction = 1 - np.exp( + ( + -const.h + / const.k_B + * np.outer( + line_nus, + 1 / (t_electrons * u.K), + ) + ).to(1) + ) + + alphas = pd.DataFrame( + ( + ALPHA_COEFFICIENT + * n_lower + * linelist.f_lu.values + * emission_correction.T + ).T + ) + + if np.any(np.isnan(alphas)) or np.any(np.isinf(np.abs(alphas))): + raise ValueError( + "Some alpha_line from vald are nan, inf, -inf " " Something went wrong!" + ) + + alphas["nu"] = line_nus.value + linelist["nu"] = line_nus.value + + # Linelist preparation below is taken from opacities_solvers/base/calc_alpha_line_at_nu + linelist["level_energy_lower"] = ((linelist["e_low"].values * u.eV).cgs).value + linelist["level_energy_upper"] = ((linelist["e_up"].values * u.eV).cgs).value + + # Radiation broadening parameter is approximated as the einstein A coefficient. Vald parameters are in log scale. + linelist["A_ul"] = 10 ** (linelist["rad"]) / ( + 4 * np.pi + ) # see 1995A&AS..112..525P for appropriate units - may be off by a factor of 4pi + + return alphas, linelist + + +class AlphaLineShortlistValdMolecule(ProcessingPlasmaProperty): + """ + Calculates the alpha (line integrated absorption coefficient in cm^-1) values for each molecular line from Vald at each depth point. This is adapted from the AlphaLineShortlistVald calculation. + + Attributes + ---------- + alpha_line_from_linelist : DataFrame + A pandas DataFrame with dtype float. This represents the alpha calculation + for each line from Vald at each depth point. This is adapted from the AlphaLineVald calculation + because shortlists do not contain js or an upper energy level. This works because the degeneracies + cancel between the n_lower recalculation and in calculating f_lu from g*f given by the linelist. + lines_from_linelist: DataFrame + A pandas dataframe containing the lines and information about the lines in + the same form and shape as alpha_line_from_linelist. + """ + + outputs = ("molecule_alpha_line_from_linelist", "molecule_lines_from_linelist") + latex_name = (r"\alpha_{\textrm{moleculeline, vald}}",) + latex_formula = ( + r"\dfrac{\pi e^{2} n_{lower} f_{lu}}{m_{e} c}\ + \Big(1-exp(-h \nu / k T) \phi(\nu)\Big)", + ) + + def calculate( + self, + atomic_data, + molecule_number_density, + t_electrons, + molecule_partition_function, + ): + # solve n_lower : n_i = N * g_i / U * e ^ (-E_i/kT) + # get f_lu : loggf -> use g = 2j+1 + # emission_correction = (1-e^(-h*nu / kT)) + # alphas = ALPHA_COEFFICIENT * n_lower * f_lu * emission_correction + + ###TODO: handle other broadening parameters + points = len(t_electrons) + + linelist = atomic_data.linelist_molecules[ + [ + "molecule", + "wavelength", + "log_gf", + "e_low", + "rad", + "stark", + "waals", + ] + ].copy() + + linelist["e_up"] = ( + ( + linelist.e_low.values * u.eV + + (const.h * const.c / (linelist.wavelength.values * u.AA)) + ) + .to(u.eV) + .value + ) + + exponent_by_point = np.exp( + np.outer( + -linelist.e_low.values * u.eV, 1 / (t_electrons * u.K * const.k_B) + ).to(1) + ) + + molecule_densities_div_partition_function = molecule_number_density.copy().div( + molecule_partition_function + ) + molecule_densities_div_partition_function.index.name = "molecule" + + # grab densities for n_lower - need to use linelist as the index and normalize by dividing by the partition function + linelist_with_density_div_partition_function = linelist.merge( + molecule_densities_div_partition_function, + how="left", + on=["molecule"], + ) + + prefactor = ( + exponent_by_point + * linelist_with_density_div_partition_function[np.arange(points)] + ).values.T # arange mask of the dataframe returns the set of densities of the appropriate ion for the line at each point + + line_nus = (linelist.wavelength.values * u.AA).to( + u.Hz, equivalencies=u.spectral() + ) + + emission_correction = 1 - np.exp( + ( + -const.h + / const.k_B + * np.outer( + line_nus, + 1 / (t_electrons * u.K), + ) + ).to(1) + ) + + alphas = pd.DataFrame( + ( + ALPHA_COEFFICIENT + * prefactor + * 10**linelist.log_gf.values + * emission_correction.T + ).T + ) + + if np.any(np.isnan(alphas)) or np.any(np.isinf(np.abs(alphas))): + raise ValueError( + "Some alpha_line from vald are nan, inf, -inf " " Something went wrong!" + ) + + alphas["nu"] = line_nus.value + linelist["nu"] = line_nus.value + + # Linelist preparation below is taken from opacities_solvers/base/calc_alpha_line_at_nu + linelist["level_energy_lower"] = ((linelist["e_low"].values * u.eV).cgs).value + linelist["level_energy_upper"] = ((linelist["e_up"].values * u.eV).cgs).value + + # Radiation broadening parameter is approximated as the einstein A coefficient. Vald parameters are in log scale. + linelist["A_ul"] = 10 ** (linelist["rad"]) / ( + 4 * np.pi + ) # see 1995A&AS..112..525P for appropriate units - may be off by a factor of 4pi + + return alphas, linelist diff --git a/stardis/radiation_field/base.py b/stardis/radiation_field/base.py index 82a7cb12..82687167 100644 --- a/stardis/radiation_field/base.py +++ b/stardis/radiation_field/base.py @@ -6,6 +6,8 @@ from stardis.radiation_field.source_functions.blackbody import blackbody_flux_at_nu from tardis.io.util import HDFWriterMixin +logger = logging.getLogger(__name__) + class RadiationField(HDFWriterMixin): """ @@ -69,7 +71,7 @@ def create_stellar_radiation_field(tracing_nus, stellar_model, stellar_plasma, c stellar_radiation_field = RadiationField( tracing_nus, blackbody_flux_at_nu, stellar_model ) - logging.info("Calculating alphas") + logger.info("Calculating alphas") calc_alphas( stellar_plasma=stellar_plasma, stellar_model=stellar_model, @@ -77,7 +79,7 @@ def create_stellar_radiation_field(tracing_nus, stellar_model, stellar_plasma, c opacity_config=config.opacity, n_threads=config.n_threads, ) - logging.info("Raytracing") + logger.info("Raytracing") raytrace( stellar_model, stellar_radiation_field, diff --git a/stardis/radiation_field/opacities/opacities_solvers/base.py b/stardis/radiation_field/opacities/opacities_solvers/base.py index 6a78804f..2854c197 100644 --- a/stardis/radiation_field/opacities/opacities_solvers/base.py +++ b/stardis/radiation_field/opacities/opacities_solvers/base.py @@ -7,6 +7,7 @@ from stardis.radiation_field.opacities.opacities_solvers.broadening import ( calculate_broadening, + calculate_molecule_broadening, ) from stardis.radiation_field.opacities.opacities_solvers.voigt import voigt_profile from stardis.radiation_field.opacities.opacities_solvers.util import ( @@ -466,6 +467,92 @@ def calc_alpha_line_at_nu( return alpha_line_at_nu, gammas, doppler_widths +def calc_molecular_alpha_line_at_nu( + stellar_plasma, + stellar_model, + tracing_nus, + line_opacity_config, + n_threads=1, +): + if line_opacity_config.disable: + return 0, 0, 0 + + line_range = line_opacity_config.broadening_range + + lines = stellar_plasma.molecule_lines_from_linelist + lines_sorted = lines.sort_values("nu") + lines_sorted_in_range = lines_sorted[ + lines_sorted.nu.between(tracing_nus.min(), tracing_nus.max()) + ] + line_nus = lines_sorted_in_range.nu.to_numpy() + + alphas_and_nu = stellar_plasma.molecule_alpha_line_from_linelist + + alphas_array = ( + alphas_and_nu[alphas_and_nu.nu.between(tracing_nus.min(), tracing_nus.max())] + .drop(labels="nu", axis=1) + .to_numpy() + ) + + gammas, doppler_widths = calculate_molecule_broadening( + lines_sorted_in_range, + stellar_model, + stellar_plasma, + line_opacity_config.broadening, + ) + + delta_nus = tracing_nus.value - line_nus[:, np.newaxis] + + 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. + 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 = np.full( + len(lines_sorted_in_range), False + ) # This is wonky but necessary for the calc_alan_entries function + if line_range.unit.physical_type == "length": + lambdas = tracing_nus.to(u.AA, equivalencies=u.spectral()) + lambdas_plus_broadening_range = lambdas + line_range.to(u.AA) + nus_plus_broadening_range = lambdas_plus_broadening_range.to( + u.Hz, equivalencies=u.spectral() + ) + line_range_value = (tracing_nus - nus_plus_broadening_range).value + elif line_range.unit.physical_type == "frequency": + line_range_value = line_range.to(u.Hz).value + else: + raise ValueError( + "Broadening range must be in units of length or frequency." + ) + + if n_threads == 1: # Single threaded + 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, + ) + + else: # Parallel threaded + 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, + ) + + return alpha_line_at_nu, gammas, doppler_widths + + @numba.njit(parallel=True) def calc_alan_entries_parallel( no_of_depth_points, @@ -735,6 +822,30 @@ def calc_alphas( stellar_radiation_field.opacities.opacities_dict[ "alpha_line_at_nu_doppler_widths" ] = doppler_widths + + if opacity_config.line.include_molecules: + ( + molecule_alpha_line_at_nu, + molecule_gammas, + molecule_doppler_widths, + ) = calc_molecular_alpha_line_at_nu( + stellar_plasma, + stellar_model, + stellar_radiation_field.frequencies, + opacity_config.line, + n_threads, + ) + + stellar_radiation_field.opacities.opacities_dict[ + "molecule_alpha_line_at_nu" + ] = molecule_alpha_line_at_nu + stellar_radiation_field.opacities.opacities_dict[ + "molecule_alpha_line_at_nu_gammas" + ] = molecule_gammas + stellar_radiation_field.opacities.opacities_dict[ + "molecule_alpha_line_at_nu_doppler_widths" + ] = molecule_doppler_widths + alphas = stellar_radiation_field.opacities.calc_total_alphas() return alphas diff --git a/stardis/radiation_field/opacities/opacities_solvers/broadening.py b/stardis/radiation_field/opacities/opacities_solvers/broadening.py index 6dd2dc77..059badf9 100644 --- a/stardis/radiation_field/opacities/opacities_solvers/broadening.py +++ b/stardis/radiation_field/opacities/opacities_solvers/broadening.py @@ -704,6 +704,34 @@ def calculate_broadening( return gammas, doppler_widths +def calculate_molecule_broadening( + lines, + stellar_model, + stellar_plasma, + broadening_line_opacity_config, +): + if "radiation" in broadening_line_opacity_config: + gammas = lines.A_ul.values[:, np.newaxis] + else: + gammas = np.zeros( + (len(lines), len(stellar_model.geometry.no_of_depth_points)), dtype=float + ) + + ions = stellar_plasma.molecule_ion_map.loc[lines.molecule] + + ion1_masses = stellar_model.composition.nuclide_masses.loc[ions.Ion1].values + ion2_masses = stellar_model.composition.nuclide_masses.loc[ions.Ion2].values + + molecule_masses = (ion1_masses + ion2_masses)[:, np.newaxis] + doppler_widths = calc_doppler_width( + lines.nu.values[:, np.newaxis], + stellar_model.temperatures.value, + molecule_masses, + ) + + return gammas, doppler_widths + + def rotation_broadening( velocity_per_pix, wavelength, flux, v_rot=0 * u.km / u.s, limb_darkening=0.6 ):