import logging
import numpy as np
import numpy.typing as npt
import pandas as pd
from astropy import units as u
from tardis import constants as const
from tardis.configuration.sorting_globals import SORTING_ALGORITHM
from tardis.plasma.equilibrium.rates.collision_strengths import (
UpsilonChiantiSolver,
UpsilonCMFGENSolver,
UpsilonRegemorterSolver,
)
logger = logging.getLogger(__name__)
BETA_COLL = (
(const.h**4 / (8 * const.k_B * const.m_e**3 * np.pi**3)) ** 0.5
).cgs
[docs]
class ThermalCollisionalRateSolver:
"""Solve thermal electron-collision bound-bound rate coefficients."""
def __init__(
self,
levels: pd.DataFrame,
radiative_transitions: pd.DataFrame,
thermal_collisional_strengths_temperatures: npt.ArrayLike,
thermal_collisional_strengths: pd.DataFrame,
collision_strengths_type: str,
collisional_strength_approximation: str = "regemorter",
) -> None:
"""Initialize thermal collisional rates for the supplied levels.
Collision strengths whose lower or upper level is absent from
``levels`` are dropped, because their energies and statistical
weights are undefined.
Parameters
----------
levels : pandas.DataFrame
Level energies and statistical weights.
radiative_transitions : pandas.DataFrame
Radiative line data.
thermal_collisional_strengths_temperatures : array-like
Temperatures of the tabulated collision strengths [K].
thermal_collisional_strengths : pandas.DataFrame
Tabulated collision strengths indexed by transition.
collision_strengths_type : str
Collision-strength format, ``"cmfgen"`` or ``"chianti"``.
collisional_strength_approximation : str, optional
Approximation for transitions without collision strengths.
"""
lower_level_index = thermal_collisional_strengths.index.droplevel(
"level_number_upper"
)
upper_level_index = thermal_collisional_strengths.index.droplevel(
"level_number_lower"
)
defined_transition_mask = lower_level_index.isin(
levels.index
) & upper_level_index.isin(levels.index)
if not defined_transition_mask.all():
logger.warning(
"Dropping %d collision strengths whose levels are absent "
"from the level data.",
(~defined_transition_mask).sum(),
)
thermal_collisional_strengths = thermal_collisional_strengths.loc[
defined_transition_mask
]
self.levels = levels
self.collision_strengths_type = collision_strengths_type
if self.collision_strengths_type == "cmfgen":
self.thermal_collision_strength_solver = UpsilonCMFGENSolver(
thermal_collisional_strengths_temperatures,
thermal_collisional_strengths,
)
elif self.collision_strengths_type == "chianti":
self.thermal_collision_strength_solver = UpsilonChiantiSolver(
thermal_collisional_strengths,
)
else:
raise ValueError(
f"collision_strengths_type {collision_strengths_type} not supported"
)
self.radiative_transitions = radiative_transitions
# find the transitions that have radiative rate data but no collisional data
self.missing_collision_strengths_index = (
radiative_transitions.index.difference(
thermal_collisional_strengths.index
)
)
self.all_collisional_strengths_index = (
self.missing_collision_strengths_index.append(
thermal_collisional_strengths.index
).sort_values()
)
self.delta_energies = (
self.levels.loc[
self.all_collisional_strengths_index.droplevel(
"level_number_upper"
)
].energy.values
- self.levels.loc[
self.all_collisional_strengths_index.droplevel(
"level_number_lower"
)
].energy.values
) * u.erg
self.g_upper = self.levels.loc[
self.all_collisional_strengths_index.droplevel("level_number_lower")
].g.values
self.g_lower = self.levels.loc[
self.all_collisional_strengths_index.droplevel("level_number_upper")
].g.values
if collisional_strength_approximation == "regemorter":
self.thermal_collision_strength_approximator = (
UpsilonRegemorterSolver(
radiative_transitions.loc[
self.missing_collision_strengths_index
]
)
)
else:
self.thermal_collision_strength_approximator = None
[docs]
def solve(self, temperatures_electron):
thermal_all_collision_strengths = self.calculate_collision_strengths(
temperatures_electron
)
boltzmann_factor = np.exp(
self.delta_energies[np.newaxis].T
/ (temperatures_electron * const.k_B),
).value
collision_rates_coeff_lu = (
(BETA_COLL / np.sqrt(temperatures_electron) * boltzmann_factor)
.to("cm3 / s")
.value
* thermal_all_collision_strengths
) # see formula A2 in Przybilla, Butler 2004 - Apj 609, 1181
collision_rates_coeff_ul = (
(self.g_lower / self.g_upper)[np.newaxis].T
/ boltzmann_factor
* collision_rates_coeff_lu
)
collision_rates_coeff_ul.index = (
collision_rates_coeff_lu.index.swaplevel(
"level_number_lower", "level_number_upper"
)
)
collision_rates_coeff_df = pd.concat(
[collision_rates_coeff_lu, collision_rates_coeff_ul]
)
collision_rates_coeff_df.index.names = [
"atomic_number",
"ion_number",
"level_number_source",
"level_number_destination",
]
collision_rates_coeff_df = collision_rates_coeff_df.reset_index()
# Add the new columns by duplicating the ion_number column
collision_rates_coeff_df["ion_number_source"] = (
collision_rates_coeff_df["ion_number"]
)
collision_rates_coeff_df["ion_number_destination"] = (
collision_rates_coeff_df["ion_number"]
)
collision_rates_coeff_df = collision_rates_coeff_df.set_index(
[
"atomic_number",
"ion_number",
"ion_number_source",
"ion_number_destination",
"level_number_source",
"level_number_destination",
]
)
return collision_rates_coeff_df
[docs]
def calculate_collision_strengths(self, temperatures_electron):
"""
Calculate collision strengths based on the provided electron temperatures.
Parameters
----------
temperatures_electron : array-like
Array-like of electron temperatures.
Returns
-------
pandas.DataFrame
DataFrame containing the calculated collision strengths.
"""
thermal_collision_strengths = (
self.thermal_collision_strength_solver.solve(temperatures_electron)
)
if self.thermal_collision_strength_approximator is None:
thermal_collision_strength_approximated = pd.DataFrame(
0,
index=self.missing_collision_strengths_index,
columns=np.arange(len(temperatures_electron)),
dtype=np.float64,
)
else:
thermal_collision_strength_approximated = (
self.thermal_collision_strength_approximator.solve(
temperatures_electron
)
)
return pd.concat(
[
thermal_collision_strengths,
thermal_collision_strength_approximated,
]
).sort_index(kind=SORTING_ALGORITHM)