import numpy as np
import pandas as pd
from tardis.configuration.sorting_globals import SORTING_ALGORITHM
from tardis.io.atom_data import AtomData
from tardis.opacities.macro_atom import util
from tardis.opacities.macro_atom.absorbing_markov_chain import (
create_absorbing_probs,
)
from tardis.opacities.macro_atom.base import (
get_macro_atom_data,
initialize_transition_probabilities,
)
from tardis.opacities.macro_atom.macroatom_continuum_transitions import (
collisional_transition_deexc_to_k_packet,
collisional_transition_excitation_internal,
collisional_transition_internal_down,
collisional_transition_ionization_internal,
collisional_transition_recombination_internal,
collisional_transition_recombination_to_k_packet,
continuum_transition_photoionization_internal,
continuum_transition_recombination_emission,
continuum_transition_recombination_internal,
create_coll_excitation_cooling_metadata,
create_coll_ionization_cooling_metadata,
create_free_bound_cooling_metadata,
create_free_free_cooling_metadata,
probability_collision_deexc_to_k_packet,
probability_collision_exc_internal,
probability_collision_internal_down,
probability_collision_ionization_internal,
probability_collision_recombination_internal,
probability_collision_recombination_to_k_packet,
probability_photoionization_internal,
probability_recombination_emission,
probability_recombination_internal,
)
from tardis.opacities.macro_atom.macroatom_line_transitions import (
line_transition_emission_down,
line_transition_internal_down,
line_transition_internal_up,
probability_emission_down,
probability_internal_down,
probability_internal_up,
)
from tardis.opacities.macro_atom.macroatom_state import (
LegacyMacroAtomState,
MacroAtomState,
)
from tardis.plasma.properties.continuum_processes.rates import (
get_ground_state_multi_index,
)
from tardis.transport.montecarlo.macro_atom import MacroAtomTransitionType
[docs]
class LegacyMacroAtomSolver:
initialize: bool = True
normalize: bool = True
def __init__(self, initialize: bool = True, normalize: bool = True) -> None:
"""
Initialize the LegacyMacroAtomSolver.
Parameters
----------
initialize : bool, optional
Whether or not to initialize the transition probability coefficients and block references when solving the first time. Default is True.
normalize : bool, optional
Whether or not to normalize the transition probabilities to unity. Default is True.
"""
self.initialize = initialize
self.normalize = normalize
[docs]
def initialize_transition_probabilities(
self, atomic_data: AtomData
) -> None:
"""
Initialize the transition probability coefficients and block references.
This method should be called when solving for the first time to set up
the necessary coefficients and block references.
Parameters
----------
atomic_data : AtomData
Atomic data containing the necessary information for initialization.
"""
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: AtomData,
mean_intensities_lines_blue_wing: pd.DataFrame,
tau_sobolev: pd.DataFrame,
beta_sobolev: pd.DataFrame | None,
stimulated_emission_factor: pd.DataFrame | np.ndarray,
) -> pd.DataFrame | None:
"""
Solve the basic transition probabilities for the macroatom.
Parameters
----------
atomic_data : AtomData
Atomic data containing macro atom information.
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 | None
Modified expansion optical depths.
stimulated_emission_factor : pd.DataFrame | np.ndarray
Stimulated emission factors.
Returns
-------
pd.DataFrame | None
Transition probabilities. Returns None if mean_intensities_lines_blue_wing is empty.
"""
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: pd.DataFrame,
atomic_data: AtomData,
tau_sobolev: pd.DataFrame,
stimulated_emission_factor: pd.DataFrame,
beta_sobolev: pd.DataFrame | None = None,
) -> LegacyMacroAtomState:
"""
Solve 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 : AtomData
Atomic data containing macro atom information.
tau_sobolev : pd.DataFrame
Expansion optical depths.
stimulated_emission_factor : pd.DataFrame
Stimulated emission factors.
beta_sobolev : pd.DataFrame | None, optional
Modified expansion optical depths. Default is None.
Returns
-------
LegacyMacroAtomState
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: pd.DataFrame,
lines: pd.DataFrame,
line_interaction_type: str = "macroatom",
) -> None:
"""
Initialize the BoundBoundMacroAtomSolver.
Parameters
----------
levels : pd.DataFrame
DataFrame containing atomic level information.
lines : pd.DataFrame
DataFrame containing spectral line information.
line_interaction_type : str, optional
Type of line interaction to use. Default is "macroatom".
"""
self.levels = levels
self.lines = lines
self.line_interaction_type = line_interaction_type
self._precompute_static_data()
def _precompute_static_data(self) -> None:
"""
Precompute static atomic data arrays.
This method extracts and reshapes line and level data that will be used
repeatedly during solve iterations. Precomputing these avoids redundant
operations on each iteration.
"""
self._oscillator_strength_ul = self.lines.f_ul.to_numpy().reshape(-1, 1)
self._oscillator_strength_lu = self.lines.f_lu.to_numpy().reshape(-1, 1)
self._nus = self.lines.nu.to_numpy().reshape(-1, 1)
self._energies_upper = (
self.levels[["energy"]]
.reindex(self.lines.index.droplevel("level_number_lower"))
.to_numpy()
)
self._energies_lower = (
self.levels[["energy"]]
.reindex(self.lines.index.droplevel("level_number_upper"))
.to_numpy()
)
self._transition_a_i_l_u_array = (
self.lines.reset_index()[
[
"atomic_number",
"ion_number",
"level_number_lower",
"level_number_upper",
]
]
.convert_dtypes(int)
.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.
self._lines_level_upper = self.lines.index.droplevel(
"level_number_lower"
)
[docs]
def solve(
self,
mean_intensities_blue_wing: pd.DataFrame,
beta_sobolevs: pd.DataFrame,
stimulated_emission_factors: np.ndarray,
) -> MacroAtomState:
"""
Solve the transition probabilities for the macroatom.
This method calculates transition probabilities and returns a MacroAtomState object
with the probabilities and 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 probabilities for the Sobolev approximation.
stimulated_emission_factors : np.ndarray
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.
"""
is_first_iteration = not hasattr(self, "computed_metadata")
if is_first_iteration:
(
normalized_probabilities,
macro_atom_transition_metadata,
line2macro_level_upper,
macro_block_references,
references_index,
) = self._solve_first_macroatom_iteration(
mean_intensities_blue_wing,
beta_sobolevs,
stimulated_emission_factors,
)
else:
normalized_probabilities = self._solve_next_macroatom_iteration(
mean_intensities_blue_wing,
beta_sobolevs,
stimulated_emission_factors,
)
(
macro_atom_transition_metadata,
line2macro_level_upper,
macro_block_references,
references_index,
) = self.computed_metadata
return MacroAtomState(
normalized_probabilities,
macro_atom_transition_metadata,
line2macro_level_upper,
macro_block_references,
references_index,
)
def _solve_first_macroatom_iteration(
self,
mean_intensities_blue_wing: pd.DataFrame,
beta_sobolevs: pd.DataFrame,
stimulated_emission_factors: np.ndarray,
) -> tuple[pd.DataFrame, pd.DataFrame, pd.Series, pd.Series, pd.Series]:
"""
Handle the first iteration of the solve method.
Fully computes all metadata for the macroatom and adds it to the class with
the computed_metadata attribute. This method performs the complete calculation
including transition probability computation, normalization, sorting, and
metadata preparation.
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 probabilities for the Sobolev approximation. These probabilities
represent the fraction of photons that escape the line formation region
without being reabsorbed.
stimulated_emission_factors : np.ndarray
Factors accounting for stimulated emission in the transitions. These
modify the transition probabilities based on the radiation field strength.
Returns
-------
normalized_probabilities : pd.DataFrame
DataFrame containing normalized transition probabilities where each source
group sums to 1.0.
macro_atom_transition_metadata : pd.DataFrame
DataFrame containing metadata for transitions including source and
destination levels, transition types, and line indices.
line2macro_level_upper : pd.Series
Series mapping line transitions to macro atom level indices for upper levels.
macro_block_references : pd.Series
Series with unique source levels as index and their first occurrence
index in the metadata as values.
"""
if self.line_interaction_type in ["downbranch", "macroatom"]:
p_emission_down, emission_down_metadata = (
line_transition_emission_down(
self._oscillator_strength_ul,
self._nus,
self._energies_upper,
self._energies_lower,
beta_sobolevs,
self._transition_a_i_l_u_array,
self.lines.line_id.to_numpy(),
)
)
else:
raise ValueError(
f"Unknown line interaction type: {self.line_interaction_type}"
)
if self.line_interaction_type == "downbranch":
probabilities_df = p_emission_down
macro_atom_transition_metadata = emission_down_metadata
elif self.line_interaction_type == "macroatom":
p_internal_down, internal_down_metadata = (
line_transition_internal_down(
self._oscillator_strength_ul,
self._nus,
self._energies_lower,
beta_sobolevs,
self._transition_a_i_l_u_array,
self.lines.line_id.to_numpy(),
)
)
p_internal_up, internal_up_metadata = line_transition_internal_up(
self._oscillator_strength_lu,
self._nus,
self._energies_lower,
mean_intensities_blue_wing,
beta_sobolevs,
stimulated_emission_factors,
self._transition_a_i_l_u_array,
self.lines.line_id.to_numpy(),
)
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. This also removes the source column from the probabilities DataFrame.
normalized_probabilities = self.normalize_transition_probabilities(
probabilities_df, macro_atom_transition_metadata
)
normalized_probabilities, macro_atom_transition_metadata = (
self.reindex_sort_and_clean_probabilities_and_metadata(
normalized_probabilities, macro_atom_transition_metadata
)
)
# We have to create the line2macro object after sorting.
line2macro_level_upper, reference_index = (
self.create_line2macro_level_upper_and_reference_idx(
macro_atom_transition_metadata, self._lines_level_upper
)
)
macro_atom_transition_metadata.drop(
columns=[
"atomic_number",
"ion_number",
"level_number_lower",
"level_number_upper",
"source_level",
],
inplace=True,
)
self.create_source_and_destination_idx_columns(
macro_atom_transition_metadata
)
macro_block_references = self.create_macro_block_references(
macro_atom_transition_metadata
)
self.computed_metadata = (
macro_atom_transition_metadata,
line2macro_level_upper,
macro_block_references,
reference_index,
)
return (
normalized_probabilities,
macro_atom_transition_metadata,
line2macro_level_upper,
macro_block_references,
reference_index,
)
def _solve_next_macroatom_iteration(
self,
mean_intensities_blue_wing: pd.DataFrame,
beta_sobolevs: pd.DataFrame,
stimulated_emission_factors: np.ndarray,
) -> pd.DataFrame:
"""
Handle subsequent iterations of the solve method.
Uses precomputed metadata and only recalculates the probabilities. This method
is optimized for speed by reusing the transition metadata, block references,
and line mappings computed in the first iteration.
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.
This parameter may have updated values compared to the first iteration.
beta_sobolevs : pd.DataFrame
Escape probabilities for the Sobolev approximation. These probabilities
represent the fraction of photons that escape the line formation region
without being reabsorbed. Values may be updated from the first iteration.
stimulated_emission_factors : np.ndarray
Factors accounting for stimulated emission in the transitions. These
modify the transition probabilities based on the radiation field strength.
May contain updated values from the radiation field calculation.
Returns
-------
pd.DataFrame
DataFrame containing normalized transition probabilities where each source
group sums to 1.0. The structure matches the first iteration output but
with updated probability values.
"""
(
macro_atom_transition_metadata,
line2macro_level_upper,
macro_block_references,
reference_index,
) = self.computed_metadata
line_trans_internal_up_ids = macro_atom_transition_metadata[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.INTERNAL_UP
].transition_line_idx.to_numpy()
line_trans_internal_down_ids = macro_atom_transition_metadata[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.INTERNAL_DOWN
].transition_line_idx.to_numpy()
line_trans_emission_down_ids = macro_atom_transition_metadata[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.BB_EMISSION
].transition_line_idx.to_numpy()
probabilities_df = pd.DataFrame(
np.zeros(
(
macro_atom_transition_metadata.shape[0],
beta_sobolevs.shape[1],
)
),
index=macro_atom_transition_metadata.index,
columns=beta_sobolevs.columns,
)
probabilities_df[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.BB_EMISSION
] = probability_emission_down(
beta_sobolevs.iloc[line_trans_emission_down_ids],
self._nus[line_trans_emission_down_ids],
self._oscillator_strength_ul[line_trans_emission_down_ids],
self._energies_upper[line_trans_emission_down_ids],
self._energies_lower[line_trans_emission_down_ids],
).to_numpy()
probabilities_df[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.INTERNAL_DOWN
] = probability_internal_down(
beta_sobolevs.iloc[line_trans_internal_down_ids],
self._nus[line_trans_internal_down_ids],
self._oscillator_strength_ul[line_trans_internal_down_ids],
self._energies_lower[line_trans_internal_down_ids],
).to_numpy()
probabilities_df[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.INTERNAL_UP
] = probability_internal_up(
beta_sobolevs.iloc[line_trans_internal_up_ids],
self._nus[line_trans_internal_up_ids],
self._oscillator_strength_lu[line_trans_internal_up_ids],
stimulated_emission_factors[line_trans_internal_up_ids],
mean_intensities_blue_wing.iloc[line_trans_internal_up_ids],
self._energies_lower[line_trans_internal_up_ids],
).to_numpy()
normalized_probabilities = self.normalize_transition_probabilities(
probabilities_df, macro_atom_transition_metadata
)
return normalized_probabilities
[docs]
def create_source_and_destination_idx_columns(
self,
macro_atom_transition_metadata: pd.DataFrame,
) -> None:
"""
Create numerical indices for source and destination levels
by mapping unique source levels to sequential integers. The destination
indices use -99 for destinations that are not sources (emission-only levels).
Parameters
----------
macro_atom_transition_metadata : pd.DataFrame
DataFrame containing macro atom transition metadata with 'source' and
'destination' columns.
"""
source_to_index = {
source: idx
for idx, source in enumerate(
macro_atom_transition_metadata.source.unique()
)
}
# -99 should never be used downstream. The presence of it means the destination is not a source,
# which means that the destination is only referenced from emission
# (or macroatom deactivation) for the given macroatom configuration.
macro_atom_transition_metadata["destination_level_idx"] = (
(macro_atom_transition_metadata.destination.map(source_to_index))
.fillna(-99)
.astype(np.int64)
)
macro_atom_transition_metadata["source_level_idx"] = (
macro_atom_transition_metadata.source.map(source_to_index)
).astype(np.int64)
[docs]
def create_macro_block_references(
self, macro_atom_transition_metadata: pd.DataFrame
) -> pd.Series:
"""
Create macro block references from the macro atom transition metadata.
This method creates a mapping from unique source levels to their first occurrence index in the metadata.
Parameters
----------
macro_atom_transition_metadata
DataFrame containing metadata for macro atom transitions.
Returns
-------
pd.Series
Series with unique source levels as index and their first occurrence index in the metadata as values.
"""
unique_source_multi_index = pd.MultiIndex.from_tuples(
macro_atom_transition_metadata.source.unique(),
names=["atomic_number", "ion_number", "level_number"],
)
macro_data = (
macro_atom_transition_metadata.reset_index()
.groupby("source")
.apply(lambda x: x.index[0])
)
# Append a dummy index so that the interactions can access a "block end" if a packet activates the macroatom highest level of the heaviest element in the montecarlo.
# Without this the kernel will crash trying to access an index that doesn't exist.
macro_data = np.append(
macro_data.values, len(macro_atom_transition_metadata)
)
unique_source_multi_index = unique_source_multi_index.append(
pd.MultiIndex.from_tuples(
[(-99, -99, -99)],
names=["atomic_number", "ion_number", "level_number"],
)
)
macro_block_references = pd.Series(
data=macro_data,
index=unique_source_multi_index,
name="macro_block_references",
)
return macro_block_references
[docs]
def create_line2macro_level_upper_and_reference_idx(
self,
macro_atom_transition_metadata: pd.DataFrame,
lines_level_upper: pd.MultiIndex,
) -> tuple[pd.Series, pd.Series]:
"""
Create a mapping from line transitions to macro atom level indices for upper levels.
This method creates a mapping that connects line transition upper levels to their
corresponding macro atom level indices. It first extracts unique source levels
from the macro atom transition metadata and assigns sequential indices to them,
then maps the line upper levels to these indices.
Parameters
----------
macro_atom_transition_metadata
DataFrame containing macro atom transition metadata.
lines_level_upper
MultiIndex containing line upper level information.
Returns
-------
tuple[pd.Series, pd.Series]
Line to macro level upper mapping and unique source indices reference.
"""
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]
return line2macro_level_upper, unique_source_series
[docs]
@staticmethod
def normalize_transition_probabilities(
probabilities_df: pd.DataFrame,
macro_atom_transition_metadata: pd.DataFrame,
) -> pd.DataFrame:
"""
Normalize transition probabilities by their source levels.
Parameters
----------
probabilities_df :
DataFrame containing transition probabilities.
macro_atom_transition_metadata:
Dataframe containing all metadata for the transitions.
Needed to obtain sources to normalize by.
Returns
-------
pd.DataFrame
Normalized probabilities where each source group sums to 1.0.
NaN values are replaced with 0.0 for cases where all transition
probabilities are zero (typically ground levels in macroatom).
"""
# Normalize the probabilities by source. This used to be optional but is never not done in TARDIS.
probabilities_df["source"] = macro_atom_transition_metadata["source"]
normalized_probabilities = probabilities_df.div(
probabilities_df.groupby("source").transform("sum"),
)
normalized_probabilities.replace(np.nan, 0.0, inplace=True)
return normalized_probabilities.drop(columns=["source"])
[docs]
class ContinuumMacroAtomSolver(BoundBoundMacroAtomSolver):
levels: pd.DataFrame
lines: pd.DataFrame
line_interaction_type: str
photoionization_data: pd.DataFrame
ionization_energies: pd.Series
selected_continuum_transitions: np.ndarray
def __init__(
self,
levels: pd.DataFrame,
lines: pd.DataFrame,
photoionization_data: pd.DataFrame,
ionization_energies: pd.Series,
selected_continuum_transitions: np.ndarray = np.array([]),
line_interaction_type: str = "macroatom",
) -> None:
"""
Initialize the ContinuumMacroAtomSolver.
Parameters
----------
levels
DataFrame containing atomic level energy information.
lines
DataFrame containing spectral line information.
photoionization_data
DataFrame containing photoionization cross-section information.
ionization_energies
Series containing ionization energies for each level.
selected_continuum_transitions
Array of selected continuum transitions (atomic_number, ion_number) pairs.
If empty, all photoionization transitions are included.
line_interaction_type
Type of line interaction to use. Default is "macroatom".
"""
super().__init__(
lines=lines,
levels=levels,
line_interaction_type=line_interaction_type,
)
if line_interaction_type != "macroatom":
raise NotImplementedError(
"ContinuumMacroAtomSolver only supports line_interaction_type='macroatom' currently."
)
if selected_continuum_transitions.size > 0:
included_species = photoionization_data.index.droplevel(
"level_number"
).isin(selected_continuum_transitions)
self.photoionization_data = photoionization_data[included_species]
else:
self.photoionization_data = photoionization_data
# selected_continuum_transitions = [
# (1, 0),
# (1, 1),
# ] # Temporary hack to test the continuum macro atom implementation.
# included_species = photoionization_data.index.droplevel(
# "level_number"
# ).isin(selected_continuum_transitions)
# self.photoionization_data = photoionization_data[included_species]
# Here we probably want to check and throw an error if the photoionization data contains atoms not in the lines and levels dataframes.
self.photoionization_data_level_energies = levels.loc[
self.photoionization_data.index.unique()
].energy
# self.ionization_frequency_thresholds = (
# self.photoionization_data.groupby(level=[0, 1, 2]).first().nu
# )
self.ionization_energies = ionization_energies
[docs]
def solve(
self,
mean_intensities_blue_wing: pd.DataFrame,
beta_sobolevs: pd.DataFrame,
stimulated_emission_factors: np.ndarray,
stim_recomb_corrected_photoionization_rate_coeff: pd.DataFrame,
spontaneous_recombination_coeff: pd.DataFrame,
coll_deexc_coeff: pd.DataFrame,
coll_exc_coeff: pd.DataFrame,
coll_ion_coeff: pd.DataFrame,
coll_recomb_coeff: pd.DataFrame,
electron_densities: pd.Series,
delta_E_yg: pd.Series,
coll_exc_cool_rate: pd.Series,
coll_exc_cool_arr: np.ndarray,
coll_ion_cool_rate: pd.Series,
coll_ion_cool_arr: np.ndarray,
fb_cool_rate: pd.Series,
fb_cool_probs_arr: np.ndarray,
ff_cool_rate: pd.Series,
) -> MacroAtomState:
"""
Solve the Macro Atom State including continuum transitions.
This method calculates transition probabilities for both bound-bound (line) and continuum
transitions and returns a MacroAtomState object with the probabilities and macro atom transition metadata.
Referenced as $p_i$ in Lucy 2003, https://doi.org/10.1051/0004-6361:20030357
Parameters
----------
mean_intensities_blue_wing
Mean intensity of the radiation field of each line in the blue wing for each shell.
Referenced as 'J^b_{lu}' internally, or 'J^b_{ji}' in the original paper.
beta_sobolevs
Escape probabilities for the Sobolev approximation.
stimulated_emission_factors
Stimulated emission factors for the lines.
stim_recomb_corrected_photoionization_rate_coeff
Corrected photoionization rate coefficients for continuum transitions.
spontaneous_recombination_coeff
Spontaneous recombination coefficients for continuum transitions.
coll_deexc_coeff
Collisional de-excitation coefficients.
coll_exc_coeff
Collisional excitation coefficients.
coll_ion_coeff
Collisional ionization coefficients.
coll_recomb_coeff
Collisional recombination coefficients.
electron_densities
Electron number densities for each cell.
delta_E_yg
Energy differences for transitions.
coll_exc_cool_rate
Collisional excitation cooling rates per cell.
coll_exc_cool_arr
Array of collisional excitation cooling rates by transition.
coll_ion_cool_rate
Collisional ionization cooling rates per cell.
coll_ion_cool_arr
Array of collisional ionization cooling rates by transition.
fb_cool_rate
Free-bound cooling rates per cell.
fb_cool_probs_arr
Array of free-bound cooling probabilities by bound level.
ff_cool_rate
Free-free (bremsstrahlung) cooling rates per cell.
Returns
-------
MacroAtomState
State of the macro atom including continuum transitions, ready to be placed into the OpacityState.
"""
is_first_iteration = not hasattr(self, "computed_metadata")
if is_first_iteration:
self.set_static_properties(
delta_E_yg, coll_ion_coeff, coll_exc_coeff
)
(
normalized_probabilities,
macro_atom_transition_metadata,
line2macro_level_upper,
macro_block_references,
references_index,
normalized_deactivating_probs,
deactivating_metadata,
absorbing_probability_matrix,
) = self._solve_first_macroatom_iteration(
mean_intensities_blue_wing,
beta_sobolevs,
stimulated_emission_factors,
stim_recomb_corrected_photoionization_rate_coeff,
spontaneous_recombination_coeff,
coll_deexc_coeff,
coll_exc_coeff,
coll_ion_coeff,
coll_recomb_coeff,
electron_densities,
coll_exc_cool_rate,
coll_exc_cool_arr,
coll_ion_cool_rate,
coll_ion_cool_arr,
fb_cool_rate,
fb_cool_probs_arr,
ff_cool_rate,
)
else:
(
normalized_probabilities,
deactivating_metadata,
normalized_deactivating_probs,
absorbing_probability_matrix,
) = self._solve_next_macroatom_iteration(
mean_intensities_blue_wing,
beta_sobolevs,
stimulated_emission_factors,
stim_recomb_corrected_photoionization_rate_coeff,
spontaneous_recombination_coeff,
coll_deexc_coeff,
coll_exc_coeff,
coll_ion_coeff,
coll_recomb_coeff,
electron_densities,
coll_exc_cool_rate,
coll_exc_cool_arr,
coll_ion_cool_rate,
coll_ion_cool_arr,
fb_cool_rate,
fb_cool_probs_arr,
ff_cool_rate,
)
(
macro_atom_transition_metadata,
line2macro_level_upper,
macro_block_references,
references_index,
) = self.computed_metadata
return MacroAtomState(
normalized_probabilities,
macro_atom_transition_metadata,
line2macro_level_upper,
macro_block_references,
references_index,
normalized_deactivating_probs,
deactivating_metadata,
absorbing_probability_matrix,
)
[docs]
def set_static_properties(
self,
delta_E_yg: pd.Series,
coll_ion_coeff: pd.DataFrame,
coll_exc_coeff: pd.DataFrame,
) -> None:
"""
Set static properties that are computed from plasma data.
This method stores energy differences and collision coefficient-derived
properties that are needed for the solve iterations but cannot be
computed in __init__ because they depend on plasma properties.
Parameters
----------
delta_E_yg
Energy differences for collisional transitions.
coll_ion_coeff
Collisional ionization coefficients.
coll_exc_coeff
Collisional excitation coefficients.
"""
self._delta_E_yg = delta_E_yg # This should be moved to the init, but can't yet because delta_E_yg comes from the plasma which isn't given to macroatom init
self._delta_E_yg_ionization = (
self.ionization_energies[
get_ground_state_multi_index(coll_ion_coeff.index)
]
- self.levels.energy.loc[coll_ion_coeff.index].values
)
self._coll_energies_lower = self.levels.energy.loc[
coll_exc_coeff.index.droplevel("level_number_upper")
]
self._coll_indices_lower = coll_exc_coeff.index.droplevel(
"level_number_upper"
)
self._coll_ion_energies = self.levels.energy.loc[coll_ion_coeff.index]
def _solve_first_macroatom_iteration(
self,
mean_intensities_blue_wing: pd.DataFrame,
beta_sobolevs: pd.DataFrame,
stimulated_emission_factors: np.ndarray,
stim_recomb_corrected_photoionization_rate_coeff: pd.DataFrame,
spontaneous_recombination_coeff: pd.DataFrame,
coll_deexc_coeff: pd.DataFrame,
coll_exc_coeff: pd.DataFrame,
coll_ion_coeff: pd.DataFrame,
coll_recomb_coeff: pd.DataFrame,
electron_densities: pd.Series,
coll_exc_cool_rate: pd.Series,
coll_exc_cool_arr: np.ndarray,
coll_ion_cool_rate: pd.Series,
coll_ion_cool_arr: np.ndarray,
fb_cool_rate: pd.Series,
fb_cool_probs_arr: np.ndarray,
ff_cool_rate: pd.Series,
) -> tuple[
pd.DataFrame,
pd.DataFrame,
pd.Series,
pd.Series,
pd.Series,
pd.DataFrame,
pd.DataFrame,
np.ndarray,
]:
"""
Handle the first iteration of the solve method for continuum macro atom.
Fully computes all metadata for the macroatom including continuum transitions and adds it to the class
with the computed_metadata attribute. This method performs the complete calculation including transition
probability computation, normalization, sorting, and metadata preparation for both bound-bound and bound-free transitions.
Parameters
----------
mean_intensities_blue_wing
Mean intensity of the radiation field of each line in the blue wing for each shell.
beta_sobolevs
Escape probabilities for the Sobolev approximation.
stimulated_emission_factors
Stimulated emission factors for the lines.
stim_recomb_corrected_photoionization_rate_coeff
Corrected photoionization rate coefficients for continuum transitions.
spontaneous_recombination_coeff
Spontaneous recombination coefficients for continuum transitions.
coll_deexc_coeff
Collisional de-excitation coefficients.
coll_exc_coeff
Collisional excitation coefficients.
coll_ion_coeff
Collisional ionization coefficients.
coll_recomb_coeff
Collisional recombination coefficients.
electron_densities
Electron number densities for each cell.
coll_exc_cool_rate
Collisional excitation cooling rates per cell.
coll_exc_cool_arr
Array of collisional excitation cooling rates by transition.
coll_ion_cool_rate
Collisional ionization cooling rates per cell.
coll_ion_cool_arr
Array of collisional ionization cooling rates by transition.
fb_cool_rate
Free-bound cooling rates per cell.
fb_cool_probs_arr
Array of free-bound cooling probabilities by bound level.
ff_cool_rate
Free-free (bremsstrahlung) cooling rates per cell.
Returns
-------
normalized_probabilities
DataFrame containing normalized transition probabilities where each source group sums to 1.0.
macro_atom_transition_metadata
DataFrame containing metadata for transitions including source and destination levels, transition types, and line indices.
line2macro_level_upper
Series mapping line transitions to macro atom level indices for upper levels.
macro_block_references
Series with unique source levels as index and their first occurrence index in the metadata as values.
references_index
Series with unique source levels as index and their assigned indices as values.
normalized_deactivating_probs
Dataframe containing emission probabilities from a chosen absorbing state
deactivating_metadata
Dataframe containing metadata for deactivation channels from a chosen absorbing state.
absorbing_probability_matrix
Ndarray describing a single jump from an interaction handler activation to a state to deactivate from.
"""
# Assemble bound-bound transitions first.
p_internal_up, internal_up_metadata = line_transition_internal_up(
self._oscillator_strength_lu,
self._nus,
self._energies_lower,
mean_intensities_blue_wing,
beta_sobolevs,
stimulated_emission_factors,
self._transition_a_i_l_u_array,
self.lines.line_id.to_numpy(),
)
internal_up_metadata[
"photoionization_key_idx"
] = -99 # Bound-bound transitions don't have continuum ids
internal_up_metadata[
"collision_key_idx"
] = -99 # Bound-bound transitions don't have collision ids
p_internal_down, internal_down_metadata = line_transition_internal_down(
self._oscillator_strength_ul,
self._nus,
self._energies_lower,
beta_sobolevs,
self._transition_a_i_l_u_array,
self.lines.line_id.to_numpy(),
)
internal_down_metadata[
"photoionization_key_idx"
] = -99 # Bound-bound transitions don't have continuum ids
internal_down_metadata[
"collision_key_idx"
] = -99 # Bound-bound transitions don't have collision ids
p_emission_down, emission_down_metadata = line_transition_emission_down(
self._oscillator_strength_ul,
self._nus,
self._energies_upper,
self._energies_lower,
beta_sobolevs,
self._transition_a_i_l_u_array,
self.lines.line_id.to_numpy(),
)
emission_down_metadata[
"photoionization_key_idx"
] = -99 # Bound-bound transitions don't have continuum ids
emission_down_metadata[
"collision_key_idx"
] = -99 # Bound-bound transitions don't have collision ids
# Then assemble photoionization transitions
p_photoionization_internal, photoionization_internal_metadata = (
continuum_transition_photoionization_internal(
stim_recomb_corrected_photoionization_rate_coeff,
self.photoionization_data_level_energies,
)
)
p_recombination_internal, recombination_internal_metadata = (
continuum_transition_recombination_internal(
spontaneous_recombination_coeff,
self.photoionization_data_level_energies,
)
)
p_recombination_emission, recombination_emission_metadata = (
continuum_transition_recombination_emission(
spontaneous_recombination_coeff, self._delta_E_yg_ionization
)
)
# Then assemble the collisional transitions
p_coll_down_to_k_packet, coll_down_to_k_packet_metadata = (
collisional_transition_deexc_to_k_packet(
coll_deexc_coeff,
electron_densities,
self._delta_E_yg,
)
)
p_coll_internal_down, coll_internal_down_metadata = (
collisional_transition_internal_down(
coll_deexc_coeff, electron_densities, self._coll_energies_lower
)
)
p_coll_internal_up, coll_internal_up_metadata = (
collisional_transition_excitation_internal(
coll_exc_coeff,
electron_densities,
self._coll_energies_lower,
)
)
p_coll_ionization_internal, coll_ionization_internal_metadata = (
collisional_transition_ionization_internal(
coll_ion_coeff, electron_densities, self._coll_ion_energies
)
)
p_coll_recomb_internal, coll_recomb_internal_metadata = (
collisional_transition_recombination_internal(
coll_recomb_coeff, electron_densities, self._coll_ion_energies
)
)
p_coll_recomb_to_k_packet, coll_recomb_to_k_packet_metadata = (
collisional_transition_recombination_to_k_packet(
coll_recomb_coeff,
electron_densities,
self._delta_E_yg_ionization,
)
)
probabilities_df = pd.concat(
[
p_emission_down,
p_internal_down,
p_internal_up,
p_photoionization_internal,
p_recombination_emission,
p_recombination_internal,
p_coll_down_to_k_packet,
p_coll_internal_down,
p_coll_internal_up,
p_coll_ionization_internal,
p_coll_recomb_internal,
p_coll_recomb_to_k_packet,
],
ignore_index=True,
)
macro_atom_transition_metadata = pd.concat(
[
emission_down_metadata,
internal_down_metadata,
internal_up_metadata,
photoionization_internal_metadata,
recombination_emission_metadata,
recombination_internal_metadata,
coll_down_to_k_packet_metadata,
coll_internal_down_metadata,
coll_internal_up_metadata,
coll_ionization_internal_metadata,
coll_recomb_internal_metadata,
coll_recomb_to_k_packet_metadata,
],
ignore_index=True,
)
probabilities_df, macro_atom_transition_metadata = (
self.reindex_sort_and_clean_probabilities_and_metadata(
probabilities_df, macro_atom_transition_metadata
)
)
probabilities_df, macro_atom_transition_metadata = (
self.normalize_cooling_block_by_rate_fractions(
probabilities_df,
macro_atom_transition_metadata,
coll_exc_cool_rate,
coll_exc_cool_arr,
coll_ion_cool_rate,
coll_ion_cool_arr,
fb_cool_rate,
fb_cool_probs_arr,
ff_cool_rate,
)
)
self.create_source_and_destination_idx_columns(
macro_atom_transition_metadata
)
normalized_probabilities = self.normalize_transition_probabilities(
probabilities_df, macro_atom_transition_metadata
)
# normalized_probabilities = probabilities_df # Useful for debugging
(
absorbing_probability_matrix,
deactivating_probs,
deactivating_metadata,
) = create_absorbing_probs(
normalized_probabilities, macro_atom_transition_metadata
)
normalized_deactivating_probs = self.normalize_transition_probabilities(
deactivating_probs, deactivating_metadata
)
line2macro_level_upper, reference_index = (
self.create_line2macro_level_upper_and_reference_idx(
deactivating_metadata, self._lines_level_upper
)
)
macro_block_references = self.create_macro_block_references(
deactivating_metadata
)
# getting all sorts of downstream problems because we trim the (1,0,0) because
# it does not have deactivation transitions so gets trimmed by the internal drop
# but trimming it screws up indexing
line2macro_level_upper += 1
macro_block_references = np.hstack([[0], macro_block_references])
reference_index = reference_index + 1
reference_index.loc[1, 0, 0] = 0
reference_index.sort_index(inplace=True)
self.computed_metadata = (
macro_atom_transition_metadata,
line2macro_level_upper,
macro_block_references,
reference_index,
)
return (
normalized_probabilities,
macro_atom_transition_metadata,
line2macro_level_upper,
macro_block_references,
reference_index,
normalized_deactivating_probs,
deactivating_metadata,
absorbing_probability_matrix,
)
def _solve_next_macroatom_iteration(
self,
mean_intensities_blue_wing: pd.DataFrame,
beta_sobolevs: pd.DataFrame,
stimulated_emission_factors: np.ndarray,
stim_recomb_corrected_photoionization_rate_coeff: pd.DataFrame,
spontaneous_recombination_coeff: pd.DataFrame,
coll_deexc_coeff: pd.DataFrame,
coll_exc_coeff: pd.DataFrame,
coll_ion_coeff: pd.DataFrame,
coll_recomb_coeff: pd.DataFrame,
electron_densities: pd.Series,
coll_exc_cool_rate: np.ndarray,
coll_exc_cool_arr: np.ndarray,
coll_ion_cool_rate: np.ndarray,
coll_ion_cool_arr: np.ndarray,
fb_cool_rate: np.ndarray,
fb_cool_probs_arr: np.ndarray,
ff_cool_rate: np.ndarray,
) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, np.ndarray]:
"""
Handle subsequent iterations of the solve method for continuum macro atom.
Uses precomputed metadata and only recalculates the probabilities. This method
is optimized for speed by reusing the transition metadata, block references,
and line mappings computed in the first iteration. Updates transition probabilities
for both bound-bound and continuum transitions.
Parameters
----------
mean_intensities_blue_wing
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.
May contain updated values compared to the first iteration.
beta_sobolevs
Escape probabilities for the Sobolev approximation. Values may be updated from
the first iteration.
stimulated_emission_factors
Factors accounting for stimulated emission in the transitions.
May contain updated values from the radiation field calculation.
stim_recomb_corrected_photoionization_rate_coeff
Corrected photoionization rate coefficients for continuum transitions.
spontaneous_recombination_coeff
Spontaneous recombination coefficients for continuum transitions.
coll_deexc_coeff
Collisional de-excitation coefficients.
coll_exc_coeff
Collisional excitation coefficients.
coll_ion_coeff
Collisional ionization coefficients.
coll_recomb_coeff
Collisional recombination coefficients.
electron_densities
Electron number densities for each cell.
coll_exc_cool_rate
Collisional excitation cooling rates per cell.
coll_exc_cool_arr
Array of collisional excitation cooling rates by transition and cell.
coll_ion_cool_rate
Collisional ionization cooling rates per cell.
coll_ion_cool_arr
Array of collisional ionization cooling rates by transition and cell.
fb_cool_rate
Free-bound cooling rates per cell.
fb_cool_probs_arr
Array of free-bound cooling probabilities by free bound transition and cell.
ff_cool_rate
Free-free (bremsstrahlung) cooling rates per cell.
Returns
-------
DataFrame
DataFrame containing normalized transition probabilities where each source
group sums to 1.0. The structure matches the first iteration output but
with updated probability values.
"""
(
macro_atom_transition_metadata,
line2macro_level_upper,
macro_block_references,
reference_index,
) = self.computed_metadata
line_trans_internal_up_ids = macro_atom_transition_metadata[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.INTERNAL_UP
].transition_line_idx.to_numpy()
line_trans_internal_down_ids = macro_atom_transition_metadata[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.INTERNAL_DOWN
].transition_line_idx.to_numpy()
line_trans_emission_down_ids = macro_atom_transition_metadata[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.BB_EMISSION
].transition_line_idx.to_numpy()
# Photoionization and recombination indices
continuum_photoion_internal_idxs = macro_atom_transition_metadata[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.PHOTOIONIZATION_INTERNAL
].photoionization_key_idx.to_numpy()
# collisional indices
collisional_down_k_idxs = macro_atom_transition_metadata[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.COLL_DOWN_TO_K_PACKET
].collision_key_idx.to_numpy()
collisional_down_internal_idxs = macro_atom_transition_metadata[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.COLL_DOWN_INTERNAL
].collision_key_idx.to_numpy()
collisional_up_internal_idxs = macro_atom_transition_metadata[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.COLL_UP_INTERNAL
].collision_key_idx.to_numpy()
collisional_ionization_internal_idxs = macro_atom_transition_metadata[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.COLL_ION_INTERNAL
].collision_key_idx.to_numpy()
collisional_recomb_internal_idxs = macro_atom_transition_metadata[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.COLL_RECOMB_INTERNAL
].collision_key_idx.to_numpy()
collisional_recomb_emission_idxs = macro_atom_transition_metadata[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.COLL_RECOMB_TO_K_PACKET
].collision_key_idx.to_numpy()
probabilities_df = pd.DataFrame(
np.zeros(
(
macro_atom_transition_metadata.shape[0],
beta_sobolevs.shape[1],
)
),
index=macro_atom_transition_metadata.index,
columns=beta_sobolevs.columns,
)
# Update the line transitions first.
probabilities_df[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.BB_EMISSION
] = probability_emission_down(
beta_sobolevs.iloc[line_trans_emission_down_ids],
self._nus[line_trans_emission_down_ids],
self._oscillator_strength_ul[line_trans_emission_down_ids],
self._energies_upper[line_trans_emission_down_ids],
self._energies_lower[line_trans_emission_down_ids],
).to_numpy()
probabilities_df[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.INTERNAL_DOWN
] = probability_internal_down(
beta_sobolevs.iloc[line_trans_internal_down_ids],
self._nus[line_trans_internal_down_ids],
self._oscillator_strength_ul[line_trans_internal_down_ids],
self._energies_lower[line_trans_internal_down_ids],
).to_numpy()
probabilities_df[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.INTERNAL_UP
] = probability_internal_up(
beta_sobolevs.iloc[line_trans_internal_up_ids],
self._nus[line_trans_internal_up_ids],
self._oscillator_strength_lu[line_trans_internal_up_ids],
stimulated_emission_factors[line_trans_internal_up_ids],
mean_intensities_blue_wing.iloc[line_trans_internal_up_ids],
self._energies_lower[line_trans_internal_up_ids],
).to_numpy()
# Then update the continuum transitions.
photoionization_internal_sources = pd.MultiIndex.from_tuples(
macro_atom_transition_metadata[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.PHOTOIONIZATION_INTERNAL
].source
)
probabilities_df[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.PHOTOIONIZATION_INTERNAL
] = probability_photoionization_internal(
stim_recomb_corrected_photoionization_rate_coeff.loc[
photoionization_internal_sources
],
self.photoionization_data_level_energies.iloc[
continuum_photoion_internal_idxs
],
).to_numpy()
probabilities_df[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.PHOTO_RECOMB_EMISSION
] = probability_recombination_emission(
spontaneous_recombination_coeff, self._delta_E_yg_ionization
).to_numpy() # WARNING - This assumes that the recombination probabilities do not get reordered within a recomb_emission block,
# because they are matched again against the photoionization nus.
# There's no reason for them to get reordered, but they are not explicitly tracked.
# Also, you'd have to add more metadata to track them explicitly, which could be done but isn't free.
recomb_internal_destinations = pd.MultiIndex.from_tuples(
macro_atom_transition_metadata[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.PHOTO_RECOMB_INTERNAL
].destination
)
probabilities_df[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.PHOTO_RECOMB_INTERNAL
] = probability_recombination_internal(
spontaneous_recombination_coeff.loc[recomb_internal_destinations],
self.photoionization_data_level_energies.loc[
recomb_internal_destinations
],
).to_numpy()
coll_down_to_k_probs = probability_collision_deexc_to_k_packet(
coll_deexc_coeff,
electron_densities,
self._delta_E_yg,
)
probabilities_df[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.COLL_DOWN_TO_K_PACKET
] = coll_down_to_k_probs.iloc[collisional_down_k_idxs].to_numpy()
probabilities_df[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.COLL_DOWN_INTERNAL
] = probability_collision_internal_down(
coll_deexc_coeff.iloc[collisional_down_internal_idxs],
electron_densities,
self._coll_energies_lower.iloc[collisional_down_internal_idxs],
).to_numpy()
probabilities_df[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.COLL_UP_INTERNAL
] = probability_collision_exc_internal(
coll_exc_coeff.iloc[collisional_up_internal_idxs],
electron_densities,
self._coll_energies_lower.iloc[collisional_up_internal_idxs],
).to_numpy()
probabilities_df[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.COLL_ION_INTERNAL
] = probability_collision_ionization_internal(
coll_ion_coeff.iloc[collisional_ionization_internal_idxs],
electron_densities,
self._coll_ion_energies.iloc[collisional_ionization_internal_idxs],
).to_numpy()
probabilities_df[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.COLL_RECOMB_INTERNAL
] = probability_collision_recombination_internal(
coll_recomb_coeff.iloc[collisional_recomb_internal_idxs],
electron_densities,
self._coll_ion_energies.iloc[collisional_recomb_internal_idxs],
).to_numpy()
probabilities_df[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.COLL_RECOMB_TO_K_PACKET
] = probability_collision_recombination_to_k_packet(
coll_recomb_coeff.iloc[collisional_recomb_emission_idxs],
electron_densities,
self._delta_E_yg_ionization.iloc[collisional_recomb_emission_idxs],
).to_numpy()
probabilities_df, macro_atom_transition_metadata = (
self.normalize_cooling_block_by_rate_fractions(
probabilities_df,
macro_atom_transition_metadata,
coll_exc_cool_rate,
coll_exc_cool_arr,
coll_ion_cool_rate,
coll_ion_cool_arr,
fb_cool_rate,
fb_cool_probs_arr,
ff_cool_rate,
)
)
# normalized_probabilities = probabilities_df # Useful for debugging
normalized_probabilities = self.normalize_transition_probabilities(
probabilities_df, macro_atom_transition_metadata
)
(
absorbing_probability_matrix,
deactivating_probs,
deactivating_metadata,
) = create_absorbing_probs(
normalized_probabilities, macro_atom_transition_metadata
)
normalized_deactivating_probs = self.normalize_transition_probabilities(
deactivating_probs, deactivating_metadata
)
return (
normalized_probabilities,
deactivating_metadata,
normalized_deactivating_probs,
absorbing_probability_matrix,
)
[docs]
def normalize_cooling_block_by_rate_fractions(
self,
probabilities: pd.DataFrame,
macro_atom_transition_metadata: pd.DataFrame,
coll_exc_cool_rate: np.ndarray,
coll_exc_cool_arr: np.ndarray,
coll_ion_cool_rate: np.ndarray,
coll_ion_cool_arr: np.ndarray,
fb_cool_rate: np.ndarray,
fb_cool_probs_arr: np.ndarray,
ff_cool_rate: np.ndarray,
) -> tuple[pd.DataFrame, pd.DataFrame]:
"""
Create and normalize k-packet source block with cooling transitions.
This method handles cooling transitions (free-free, free-bound, collisional) which are not
regular macroatom transitions. They must be added to and normalized separately because
we cannot complete the normalization until all cooling pathways exist.
The method incorporates cooling rates to establish the proper relative probabilities
for deactivation pathways.
Parameters
----------
probabilities
DataFrame containing transition probabilities.
macro_atom_transition_metadata
DataFrame containing all metadata for the transitions.
coll_exc_cool_rate
Collisional excitation cooling rates per cell.
coll_exc_cool_arr
Array of collisional excitation cooling rates by transition.
coll_ion_cool_rate
Collisional ionization cooling rates per cell.
coll_ion_cool_arr
Array of collisional ionization cooling rates by transition.
fb_cool_rate
Free-bound cooling rates per cell.
fb_cool_probs_arr
Array of free-bound cooling probabilities by bound level.
ff_cool_rate
Free-free (bremsstrahlung) cooling rates per cell.
Returns
-------
probabilities
Probabilities with the cooling block normalized by rate, and free-free deactivation
added if it did not exist yet.
macro_atom_transition_metadata
Updated metadata with cooling transitions appended if they did not exist.
"""
# Check if cooling block exists
if ~(
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.FB_COOLING
).any():
metadata_size = len(macro_atom_transition_metadata)
ff_cool_metadata = create_free_free_cooling_metadata(metadata_size)
fb_cool_metadata = create_free_bound_cooling_metadata(
metadata_size + len(ff_cool_metadata), fb_cool_probs_arr
)
coll_exc_cool_metadata = create_coll_excitation_cooling_metadata(
metadata_size + len(ff_cool_metadata) + len(fb_cool_metadata),
coll_exc_cool_arr,
macro_atom_transition_metadata[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.INTERNAL_UP
].destination,
)
coll_ion_cool_metadata = create_coll_ionization_cooling_metadata(
metadata_size
+ len(ff_cool_metadata)
+ len(fb_cool_metadata)
+ len(coll_exc_cool_metadata),
coll_ion_cool_arr,
)
macro_atom_transition_metadata = pd.concat(
[
macro_atom_transition_metadata,
ff_cool_metadata,
fb_cool_metadata,
coll_exc_cool_metadata,
coll_ion_cool_metadata,
]
)
ff_df = pd.DataFrame(ff_cool_rate).T
ff_df.index = [metadata_size]
fb_df = pd.DataFrame(fb_cool_rate * fb_cool_probs_arr.T)
fb_df.index = fb_df.index + metadata_size + len(ff_df)
coll_exc_df = pd.DataFrame(coll_exc_cool_rate * coll_exc_cool_arr.T)
coll_exc_df.index = (
coll_exc_df.index + metadata_size + len(ff_df) + len(fb_df)
)
coll_ion_df = pd.DataFrame(coll_ion_cool_rate * coll_ion_cool_arr.T)
coll_ion_df.index = (
coll_ion_df.index
+ metadata_size
+ len(ff_df)
+ len(fb_df)
+ len(coll_exc_df)
)
probabilities = pd.concat(
[probabilities, ff_df, fb_df, coll_exc_df, coll_ion_df]
)
else:
probabilities[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.FF_COOLING
] = ff_cool_rate
probabilities[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.FB_COOLING
] = fb_cool_rate * fb_cool_probs_arr.T
probabilities[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.COLL_EXC_COOL
] = coll_exc_cool_rate * coll_exc_cool_arr.T
probabilities[
macro_atom_transition_metadata.transition_type
== MacroAtomTransitionType.COLL_ION_COOL
] = coll_ion_cool_rate * coll_ion_cool_arr.T
return probabilities, macro_atom_transition_metadata