Source code for tardis.opacities.macro_atom.macroatom_solver

import numpy as np
import pandas as pd

from tardis.opacities.macro_atom import util
from tardis.opacities.macro_atom.base import (
    get_macro_atom_data,
    initialize_transition_probabilities,
)
from tardis.opacities.macro_atom.macroatom_transitions import (
    line_transition_emission_down,
    line_transition_internal_down,
    line_transition_internal_up,
)
from tardis.opacities.macro_atom.macroatom_state import (
    MacroAtomState,
    LegacyMacroAtomState,
)
from astropy import constants as const


[docs] class LegacyMacroAtomSolver: initialize: bool = True normalize: bool = True def __init__(self, initialize=True, normalize=True): """Solver class for Macro Atom related opacities Parameters ---------- initialize: bool Whether or not to initialize the transition probabilitiy coefficients and block references when solving the first time (default True) normalize: bool Whether or not to normalize the transition probabilities to unity. Default True """ self.initialize = initialize self.normalize = normalize
[docs] def initialize_transition_probabilities(self, atomic_data): """initialize the transition probability coefficients and block references when solving the first time Parameters ---------- atomic_data : tardis.io.atom_data.AtomData Atomic Data """ coef_and_block_ref = initialize_transition_probabilities(atomic_data) self.transition_probability_coef = coef_and_block_ref[ "transition_probability_coef" ] self.block_references = coef_and_block_ref["block_references"] self.initialize = False
[docs] def solve_transition_probabilities( self, atomic_data, mean_intensities_lines_blue_wing, tau_sobolev, beta_sobolev, stimulated_emission_factor, ): """Solve the basic transition probabilities for the macroatom Parameters ---------- atomic_data : tardis.io.atom_data.AtomData Atomic Data mean_intensities_lines_blue_wing : pd.DataFrame Mean intensity of the radiation field of each line in the blue wing for each shell For more detail see Lucy 2003, https://doi.org/10.1051/0004-6361:20030357 tau_sobolev : pd.DataFrame Expansion Optical Depths beta_sobolev : pd.DataFrame Modified expansion Optical Depths stimulated_emission_factor : np.ndarray Returns ------- pd.DataFrame Transition Probabilities """ if self.initialize: self.initialize_transition_probabilities(atomic_data) # Referenced in https://github.com/tardis-sn/tardis/issues/3009 if len(mean_intensities_lines_blue_wing) == 0: return None macro_atom_data = get_macro_atom_data(atomic_data) transition_probabilities = np.empty( (self.transition_probability_coef.shape[0], beta_sobolev.shape[1]) ) transition_type = macro_atom_data.transition_type.values lines_idx = macro_atom_data.lines_idx.values tpos = macro_atom_data.transition_probability.values # This function modifies transition_probabilities inplace util.fast_calculate_transition_probabilities( tpos, beta_sobolev.values, mean_intensities_lines_blue_wing.values, stimulated_emission_factor, transition_type, lines_idx, self.block_references, transition_probabilities, self.normalize, ) transition_probabilities_df = pd.DataFrame( transition_probabilities, index=macro_atom_data.transition_line_id, columns=tau_sobolev.columns, ) return transition_probabilities_df
[docs] def solve( self, mean_intensities_lines_blue_wing, atomic_data, tau_sobolev, stimulated_emission_factor, beta_sobolev=None, ): """Solved the Macro Atom State Parameters ---------- mean_intensities_lines_blue_wing : pd.DataFrame Mean intensity of the radiation field of each line in the blue wing for each shell atomic_data : tardis.io.atom_data.AtomData Atomic Data tau_sobolev : pd.DataFrame Expansion Optical Depths stimulated_emission_factor : pd.DataFrame beta_sobolev : pd.DataFrame Returns ------- tardis.opacities.macroatom_state.MacroAtomState State of the macro atom ready to be placed into the OpacityState """ transition_probabilities = self.solve_transition_probabilities( atomic_data, mean_intensities_lines_blue_wing, tau_sobolev, beta_sobolev, stimulated_emission_factor, ) macro_block_references = atomic_data.macro_atom_references[ "block_references" ] macro_atom_info = atomic_data.macro_atom_data return LegacyMacroAtomState( transition_probabilities, macro_atom_info["transition_type"], macro_atom_info["destination_level_idx"], macro_atom_info["lines_idx"], macro_block_references, atomic_data.lines_upper2macro_reference_idx, )
[docs] class BoundBoundMacroAtomSolver: levels: pd.DataFrame lines: pd.DataFrame def __init__(self, levels, lines): self.levels = levels self.lines = lines
[docs] def solve( self, mean_intensities_blue_wing, beta_sobolevs, stimulated_emission_factors, ): """ Solves the transition probabilities and returns a DataFrame with the probabilities and a DataFrame with the macro atom transition metadata. Referenced as $p_i$ in Lucy 2003, https://doi.org/10.1051/0004-6361:20030357 Parameters ---------- mean_intensities_blue_wing : pd.DataFrame Mean intensity of the radiation field of each line in the blue wing for each shell. For more detail see Lucy 2003, https://doi.org/10.1051/0004-6361:20030357 Referenced as 'J^b_{lu}' internally, or 'J^b_{ji}' in the original paper. beta_sobolevs : pd.DataFrame Escape probabilites for the Sobolev approximation. stimulated_emission_factors : pd.DataFrame Stimulated emission factors for the lines. Returns ------- MacroAtomState A MacroAtomState object containing the transition probabilities, transition metadata, and a mapping from line IDs to macro atom level upper indices. """ oscillator_strength_ul = self.lines.f_ul.values.reshape(-1, 1) oscillator_strength_lu = self.lines.f_lu.values.reshape(-1, 1) nus = self.lines.nu.values.reshape(-1, 1) line_ids = self.lines.line_id.values energies_upper = ( self.levels[["energy"]] .rename(columns={"energy": "level_number_upper"}) .reindex(self.lines.index.droplevel("level_number_lower")) .values ) energies_lower = ( self.levels[["energy"]] .rename(columns={"energy": "level_number_lower"}) .reindex(self.lines.index.droplevel("level_number_upper")) .values ) transition_a_i_l_u_array = self.lines.reset_index()[ [ "atomic_number", "ion_number", "level_number_lower", "level_number_upper", ] ].values # This is a helper array to make the source and destination columns. The letters stand for atomic_number, ion_number, lower level, upper level. lines_level_upper = self.lines.index.droplevel("level_number_lower") p_emission_down, emission_down_metadata = line_transition_emission_down( oscillator_strength_ul, nus, energies_upper, energies_lower, beta_sobolevs, transition_a_i_l_u_array, line_ids, ) p_internal_down, internal_down_metadata = line_transition_internal_down( oscillator_strength_ul, nus, energies_lower, beta_sobolevs, transition_a_i_l_u_array, line_ids, ) p_internal_up, internal_up_metadata = line_transition_internal_up( oscillator_strength_lu, nus, energies_lower, mean_intensities_blue_wing, beta_sobolevs, stimulated_emission_factors, transition_a_i_l_u_array, line_ids, ) probabilities_df = pd.concat( [p_emission_down, p_internal_down, p_internal_up] ) macro_atom_transition_metadata = pd.concat( [ emission_down_metadata, internal_down_metadata, internal_up_metadata, ] ) # Normalize the probabilities by source. This used to be optional but is never not done in TARDIS. probabilities_df = probabilities_df.div( probabilities_df.groupby("source").transform("sum"), ) probabilities_df.replace(np.nan, 0, inplace=True) # fill value for nans where the transition probabilites are all 0, which happens for ground levels that should never be accessed in the active macroatom. probabilities_df.drop(columns=["source"], inplace=True) probabilities_df = probabilities_df.reset_index( drop=True ) # Reset to create a unique macro_atom_transition_id. probabilities_df.index.rename("macro_atom_transition_id", inplace=True) macro_atom_transition_metadata = ( macro_atom_transition_metadata.reset_index() ) macro_atom_transition_metadata.index.rename( "macro_atom_transition_id", inplace=True ) macro_atom_transition_metadata["source_level"] = ( macro_atom_transition_metadata.source.apply(lambda x: x[2]) ) macro_atom_transition_metadata = ( macro_atom_transition_metadata.sort_values( ["atomic_number", "ion_number", "source_level"] ) ) # This is how carsus sorted the macro atom transitions. probabilities_df = probabilities_df.loc[ macro_atom_transition_metadata.index ] # Reorder to match the metadata, which was sorted to match carsus. # We have to create the line2macro object after sorting. unique_source_index = pd.MultiIndex.from_tuples( macro_atom_transition_metadata.source.unique(), names=["atomic_number", "ion_number", "level_number"], ) unique_source_series = pd.Series( index=unique_source_index, data=range(len(macro_atom_transition_metadata.source.unique())), ) line2macro_level_upper = unique_source_series.loc[lines_level_upper] macro_atom_transition_metadata.drop( columns=[ "atomic_number", "ion_number", "level_number_lower", "level_number_upper", "source_level", ], inplace=True, ) return MacroAtomState( probabilities_df, macro_atom_transition_metadata, line2macro_level_upper, )