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.base import (
get_macro_atom_data,
initialize_transition_probabilities,
)
from tardis.opacities.macro_atom.macroatom_state import (
LegacyMacroAtomState,
MacroAtomState,
)
from tardis.opacities.macro_atom.macroatom_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.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):
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",
]
].to_numpy() # 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,
self._lines_level_upper,
)
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,
lines_level_upper: pd.MultiIndex,
) -> 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.
lines_level_upper : pd.MultiIndex
MultiIndex containing the upper level information for each line transition,
used for creating the line-to-macro-atom level mapping.
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 = normalize_transition_probabilities(
probabilities_df
)
normalized_probabilities, macro_atom_transition_metadata = (
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 = (
create_line2macro_level_upper_and_reference_idx(
macro_atom_transition_metadata, lines_level_upper
)
)
macro_atom_transition_metadata.drop(
columns=[
"atomic_number",
"ion_number",
"level_number_lower",
"level_number_upper",
"source_level",
],
inplace=True,
)
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)
macro_block_references = 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()
probabilities_df["source"] = (
macro_atom_transition_metadata.source.values
)
normalized_probabilities = normalize_transition_probabilities(
probabilities_df
)
return normalized_probabilities
[docs]
def create_macro_block_references(macro_atom_transition_metadata):
"""
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 : pandas.DataFrame
DataFrame containing metadata for macro atom transitions.
Returns
-------
pandas.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(
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 : pd.DataFrame
DataFrame containing macro atom transition metadata
lines_level_upper : pd.MultiIndex
MultiIndex containing line upper level information
Returns
-------
pd.Series
Series mapping line transitions to macro atom level indices
pd.Series
Series with unique source levels as index and their assigned indices as values
"""
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]
def normalize_transition_probabilities(
probabilities_df: pd.DataFrame,
) -> pd.DataFrame:
"""
Normalize transition probabilities by their source levels.
Parameters
----------
probabilities_df : pd.DataFrame
DataFrame containing transition probabilities with a 'source' column
for grouping.
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.
normalized_probabilities = probabilities_df.div(
probabilities_df.groupby("source").transform("sum"),
)
normalized_probabilities.replace(np.nan, 0, inplace=True)
return normalized_probabilities.drop(columns=["source"])