Ionization rates (#2897)

* Photoionization rate solver

* Collisional ionization rate solver

* Move to correct location

* Fix imports

* Correctly define photoion coefficients

* Test cells for estimated rate coeff solver

* Add actual rate solver

* Adds arguments and notes to the rate solver

* Some notes about the rate equations

* Fix small docstring error

* Correctly integrate spontaneous recombination rate coeff

* Update test notebook

* Separate analytic and estimated rate solvers

* Add some docstrings

* Renames and unit fixes for collisional ionization

* Fixes class name in init

* Better variable names and a docstring

* Save notebook with output

* Better docstrings

* Add markdown note regarding the source of comparisons

tardis/plasma/electron_energy_distribution/base.py

@@ -21,7 +21,7 @@ class ThermalElectronEnergyDistribution(ElectronEnergyDistribution):
     temperature : Quantity
         Electron temperatures in Kelvin.
     number_density : Quantity
-        Electron number densities in g/cm^3.
+        Electron number densities in cm^-3.
     """
 
     temperature: u.Quantity

tardis/plasma/equilibrium/rates/__init__.py

@@ -2,9 +2,24 @@
     UpsilonCMFGENSolver,
     UpsilonRegemorterSolver,
 )
+from tardis.plasma.equilibrium.rates.collisional_ionization_rates import (
+    CollisionalIonizationRateSolver,
+)
+from tardis.plasma.equilibrium.rates.collisional_ionization_strengths import (
+    CollisionalIonizationSeaton,
+)
 from tardis.plasma.equilibrium.rates.collisional_rates import (
     ThermalCollisionalRateSolver,
 )
+from tardis.plasma.equilibrium.rates.photoionization_rates import (
+    AnalyticPhotoionizationRateSolver,
+    EstimatedPhotoionizationRateSolver,
+)
+from tardis.plasma.equilibrium.rates.photoionization_strengths import (
+    AnalyticPhotoionizationCoeffSolver,
+    EstimatedPhotoionizationCoeffSolver,
+    SpontaneousRecombinationCoeffSolver,
+)
 from tardis.plasma.equilibrium.rates.radiative_rates import (
     RadiativeRatesSolver,
 )

tardis/plasma/equilibrium/rates/collisional_ionization_rates.py

@@ -0,0 +1,50 @@
+from tardis.plasma.equilibrium.rates.collisional_ionization_strengths import (
+    CollisionalIonizationSeaton,
+)
+
+
+class CollisionalIonizationRateSolver:
+    """Solver for collisional ionization and recombination rates."""
+
+    def __init__(self, photoionization_cross_sections):
+        self.photoionization_cross_sections = photoionization_cross_sections
+
+    def solve(self, electron_temperature, saha_factor, approximation="seaton"):
+        """Solve the collisional ionization and recombination rates.
+
+        Parameters
+        ----------
+        electron_temperature : u.Quantity
+            Electron temperatures per cell
+        saha_factor : pandas.DataFrame, dtype float
+            The Saha factor for each cell. Indexed by atom number, ion number, level number.
+        approximation : str, optional
+            The rate approximation to use, by default "seaton"
+
+        Returns
+        -------
+        pd.DataFrame
+            Collisional ionization rates
+        pd.DataFrame
+            Collisional recombination rates
+
+        Raises
+        ------
+        ValueError
+            If an unsupported approximation is requested.
+        """
+        if approximation == "seaton":
+            strength_solver = CollisionalIonizationSeaton(
+                self.photoionization_cross_sections
+            )
+        else:
+            raise ValueError(f"approximation {approximation} not supported")
+
+        collision_ionization_rates = strength_solver.solve(electron_temperature)
+
+        # Inverse of the ionization rate for equilibrium
+        collision_recombination_rates = collision_ionization_rates.multiply(
+            saha_factor.loc[collision_ionization_rates.index]
+        )
+
+        return collision_ionization_rates, collision_recombination_rates

tardis/plasma/equilibrium/rates/collisional_ionization_strengths.py

@@ -0,0 +1,59 @@
+import astropy.units as u
+import numpy as np
+import pandas as pd
+
+from tardis import constants as const
+
+H = const.h.cgs
+K_B = const.k_B.cgs
+
+
+class CollisionalIonizationSeaton:
+    """Solver for collisional ionization rate coefficients in the Seaton approximation."""
+
+    def __init__(self, photoionization_cross_sections):
+        self.photoionization_cross_sections = photoionization_cross_sections
+
+    def solve(self, electron_temperature):
+        """
+        Parameters
+        ----------
+        electron_temperature : u.Quantity
+            The electron temperature in K.
+
+        Returns
+        -------
+        pandas.DataFrame, dtype float
+            The rate coefficient for collisional ionization in the Seaton
+            approximation. Multiply with the electron density and the
+            level number density to obtain the total rate.
+
+        Notes
+        -----
+        The rate coefficient for collisional ionization in the Seaton approximation
+        is calculated according to Eq. 9.60 in [1].
+
+        References
+        ----------
+        .. [1] Hubeny, I. and Mihalas, D., "Theory of Stellar Atmospheres". 2014.
+        """
+        photo_ion_cross_sections_threshold = (
+            self.photoionization_cross_sections.groupby(level=[0, 1, 2]).first()
+        )
+        nu_i = photo_ion_cross_sections_threshold["nu"]
+        u0s = (
+            nu_i.values[np.newaxis].T * u.Hz / electron_temperature * (H / K_B)
+        )
+        factor = np.exp(-u0s) / u0s
+        factor = pd.DataFrame(factor, index=nu_i.index)
+        coll_ion_coeff = 1.55e13 * photo_ion_cross_sections_threshold["x_sect"]
+        coll_ion_coeff = factor.multiply(coll_ion_coeff, axis=0)
+        coll_ion_coeff = coll_ion_coeff.divide(
+            np.sqrt(electron_temperature), axis=1
+        )
+
+        ion_number = coll_ion_coeff.index.get_level_values("ion_number").values
+        coll_ion_coeff[ion_number == 0] *= 0.1
+        coll_ion_coeff[ion_number == 1] *= 0.2
+        coll_ion_coeff[ion_number >= 2] *= 0.3
+        return coll_ion_coeff

tardis/plasma/equilibrium/rates/photoionization_rates.py

@@ -0,0 +1,214 @@
+from tardis.plasma.equilibrium.rates.photoionization_strengths import (
+    AnalyticPhotoionizationCoeffSolver,
+    EstimatedPhotoionizationCoeffSolver,
+    SpontaneousRecombinationCoeffSolver,
+)
+
+
+class AnalyticPhotoionizationRateSolver:
+    """Solve the photoionization and spontaneous recombination rates in the
+    case where the radiation field is computed analytically.
+    """
+
+    def __init__(self, photoionization_cross_sections):
+        self.photoionization_cross_sections = photoionization_cross_sections
+
+        self.spontaneous_recombination_rate_coeff_solver = (
+            SpontaneousRecombinationCoeffSolver(
+                self.photoionization_cross_sections
+            )
+        )
+
+    def compute_rates(
+        self,
+        photoionization_rate_coeff,
+        stimulated_recombination_rate_coeff,
+        spontaneous_recombination_rate_coeff,
+        level_number_density,
+        ion_number_density,
+        electron_number_density,
+        saha_factor,
+    ):
+        """Compute the photoionization and spontaneous recombination rates
+
+        Parameters
+        ----------
+        photoionization_rate_coeff : pd.DataFrame
+            The photoionization rate coefficients for each transition.
+            Columns are cells.
+        stimulated_recombination_rate_coeff : pd.DataFrame
+            The stimulated recombination rate coefficients for each transition.
+            Columns are cells.
+        spontaneous_recombination_rate_coeff : pd.DataFrame
+            The spontaneous recombination rate coefficients for each transition.
+            Columns are cells.
+        level_number_density : pd.DataFrame
+            The electron energy level number density. Columns are cells.
+        ion_number_density : pd.DataFrame
+            The ion number density. Columns are cells.
+        electron_number_density : u.Quantity
+            The free electron number density per cell.
+        saha_factor : pd.DataFrame
+            The LTE population factor. Columns are cells.
+
+        Returns
+        -------
+        pd.DataFrame
+            Photoionization rate for each electron energy level. Columns are cells
+        pd.DataFrame
+            Spontaneous recombination rate for each electron energy level. Columns are cells
+        """
+        photoionization_rate = (
+            photoionization_rate_coeff * level_number_density
+            - saha_factor
+            * stimulated_recombination_rate_coeff
+            * ion_number_density
+            * electron_number_density
+        )
+        spontaneous_recombination_rate = (
+            saha_factor
+            * spontaneous_recombination_rate_coeff
+            * ion_number_density
+            * electron_number_density
+        )
+
+        return photoionization_rate, spontaneous_recombination_rate
+
+    def solve(
+        self,
+        dilute_blackbody_radiationfield_state,
+        electron_energy_distribution,
+        level_number_density,
+        ion_number_density,
+        saha_factor,
+    ):
+        """Solve the photoionization and spontaneous recombination rates in the
+        case where the radiation field is not estimated.
+
+        Parameters
+        ----------
+        dilute_blackbody_radiationfield_state : DiluteBlackBodyRadiationFieldState
+            A dilute black body radiation field state.
+        electron_energy_distribution : ThermalElectronEnergyDistribution
+            Electron properties.
+        level_number_density : pd.DataFrame
+            Electron energy level number density. Columns are cells.
+        ion_number_density : pd.DataFrame
+            Ion number density. Columns are cells.
+        saha_factor : pd.DataFrame
+            Saha factor: the LTE level number density divided by the LTE ion
+            number density and the electron number density.
+
+        Returns
+        -------
+        pd.DataFrame
+            Photoionization rate. Columns are cells.
+        pd.DataFrame
+            Spontaneous recombination rate. Columns are cells.
+        """
+        photoionization_rate_coeff_solver = AnalyticPhotoionizationCoeffSolver(
+            self.photoionization_cross_sections
+        )
+
+        photoionization_rate_coeff, stimulated_recombination_rate_coeff = (
+            photoionization_rate_coeff_solver.solve(
+                dilute_blackbody_radiationfield_state,
+                electron_energy_distribution.temperature,
+            )
+        )
+
+        spontaneous_recombination_rate_coeff = (
+            self.spontaneous_recombination_rate_coeff_solver.solve(
+                electron_energy_distribution.temperature
+            )
+        )
+
+        return self.compute_rates(
+            photoionization_rate_coeff,
+            stimulated_recombination_rate_coeff,
+            spontaneous_recombination_rate_coeff,
+            level_number_density,
+            ion_number_density,
+            electron_energy_distribution.number_density,
+            saha_factor,
+        )
+
+
+class EstimatedPhotoionizationRateSolver(AnalyticPhotoionizationRateSolver):
+    """Solve the photoionization and spontaneous recombination rates in the
+    case where the radiation field is estimated by Monte Carlo processes.
+    """
+
+    def __init__(
+        self, photoionization_cross_sections, level2continuum_edge_idx
+    ):
+        super().__init__(
+            photoionization_cross_sections,
+        )
+        self.level2continuum_edge_idx = level2continuum_edge_idx
+
+    def solve(
+        self,
+        electron_energy_distribution,
+        radfield_mc_estimators,
+        time_simulation,
+        volume,
+        level_number_density,
+        ion_number_density,
+        saha_factor,
+    ):
+        """Solve the photoionization and spontaneous recombination rates in the
+        case where the radiation field is estimated by Monte Carlo processes.
+
+        Parameters
+        ----------
+        electron_energy_distribution : ThermalElectronEnergyDistribution
+            Electron properties.
+        radfield_mc_estimators : RadiationFieldMCEstimators
+            Estimators of the radiation field properties.
+        time_simulation : u.Quantity
+            Time of simulation.
+        volume : u.Quantity
+            Volume per cell.
+        level_number_density : pd.DataFrame
+            Electron energy level number density. Columns are cells.
+        ion_number_density : pd.DataFrame
+            Ion number density. Columns are cells.
+        saha_factor : pd.DataFrame
+            Saha factor: the LTE level number density divided by the LTE ion
+            number density and the electron number density.
+
+        Returns
+        -------
+        pd.DataFrame
+            Photoionization rate. Columns are cells.
+        pd.DataFrame
+            Spontaneous recombination rate. Columns are cells.
+        """
+        photoionization_rate_coeff_solver = EstimatedPhotoionizationCoeffSolver(
+            self.level2continuum_edge_idx
+        )
+
+        photoionization_rate_coeff, stimulated_recombination_rate_coeff = (
+            photoionization_rate_coeff_solver.solve(
+                radfield_mc_estimators,
+                time_simulation,
+                volume,
+            )
+        )
+
+        spontaneous_recombination_rate_coeff = (
+            self.spontaneous_recombination_rate_coeff_solver.solve(
+                electron_energy_distribution.temperature
+            )
+        )
+
+        return self.compute_rates(
+            photoionization_rate_coeff,
+            stimulated_recombination_rate_coeff,
+            spontaneous_recombination_rate_coeff,
+            level_number_density,
+            ion_number_density,
+            electron_energy_distribution.number_density,
+            saha_factor,
+        )