Changing construction of the first guess and coefficient matrices in NLTE solver

tardis/plasma/properties/nlte_excitation_data.py

@@ -0,0 +1,41 @@
+import logging
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+class NLTEExcitationData(object):
+    def __init__(self, atom_data, nlte_excitation_species):
+        self.atom_data = atom_data
+        self.lines = atom_data.lines.reset_index()
+        self.nlte_excitation_species = nlte_excitation_species
+
+        if nlte_excitation_species:
+            logger.info("Preparing the NLTE data")
+            self._init_indices()
+
+    def _init_indices(self):
+        self.lines_idx = {}
+        self.lines_level_number_lower = {}
+        self.lines_level_number_upper = {}
+        self.A_uls = {}
+        self.B_uls = {}
+        self.B_lus = {}
+
+        for species in self.nlte_excitation_species:
+            lines_idx = np.where(
+                (self.lines.atomic_number == species[0])
+                & (self.lines.ion_number == species[1])
+            )
+            self.lines_idx[species] = lines_idx
+            self.lines_level_number_lower[
+                species
+            ] = self.lines.level_number_lower.values[lines_idx].astype(int)
+            self.lines_level_number_upper[
+                species
+            ] = self.lines.level_number_upper.values[lines_idx].astype(int)
+
+            self.A_uls[species] = self.atom_data.lines.A_ul.values[lines_idx]
+            self.B_uls[species] = self.atom_data.lines.B_ul.values[lines_idx]
+            self.B_lus[species] = self.atom_data.lines.B_lu.values[lines_idx]

tardis/plasma/properties/nlte_rate_equation_solver.py

@@ -1,8 +1,10 @@
 import pandas as pd
 import numpy as np
 from scipy.optimize import root
+from tardis.io.atom_data.base import NLTEData
 
 from tardis.plasma.properties.base import ProcessingPlasmaProperty
+from tardis.plasma.properties.nlte_excitation_data import NLTEExcitationData
 
 __all__ = [
     "NLTERateEquationSolver",
@@ -25,6 +27,7 @@ def calculate(
         phi,
         rate_matrix_index,
         number_density,
+        nlte_excitation_species,
     ):
         """Calculates ion number densities and electron densities using NLTE ionization.
 
@@ -65,6 +68,7 @@ def calculate(
             Electron density with NLTE ionization treatment.
         """
 
+        # nlte_data = NLTEExcitationData(atomic_data, nlte_excitation_species) - will be used in a future PR
         (
             total_photo_ion_coefficients,
             total_rad_recomb_coefficients,
@@ -79,6 +83,8 @@ def calculate(
             partition_function,
             levels,
             level_boltzmann_factor,
+            nlte_excitation_species,
+            rate_matrix_index,
         )
 
         initial_electron_densities = number_density.sum(axis=0)
@@ -99,6 +105,7 @@ def calculate(
                 number_density[shell]
             )
             first_guess = self.prepare_first_guess(
+                rate_matrix_index,
                 atomic_numbers,
                 number_density[shell],
                 initial_electron_densities[shell],
@@ -369,6 +376,8 @@ def prepare_ion_recomb_coefficients_nlte_ion(
         partition_function,
         levels,
         level_boltzmann_factor,
+        nlte_excitation_species,
+        rate_matrix_index,
     ):
         """
         Prepares the ionization and recombination coefficients by grouping them for
@@ -420,6 +429,7 @@ def prepare_ion_recomb_coefficients_nlte_ion(
             .groupby(level=("atomic_number", "ion_number"))
             .sum()
         )
+
         total_rad_recomb_coefficients = (
             (alpha_sp + alpha_stim)
             .groupby(level=["atomic_number", "ion_number"])
@@ -438,6 +448,61 @@ def prepare_ion_recomb_coefficients_nlte_ion(
             .groupby(level=("atomic_number", "ion_number"))
             .sum()
         )
+        if len(nlte_excitation_species) != 0:
+            total_photo_ion_coefficients_with_levels = (
+                (level_population_fraction.loc[gamma.index] * gamma)
+                .groupby(level=("atomic_number", "ion_number", "level_number"))
+                .sum()
+            )
+            total_rad_recomb_coefficients_with_levels = (
+                (alpha_sp + alpha_stim)
+                .groupby(level=("atomic_number", "ion_number", "level_number"))
+                .sum()
+            )
+
+            total_coll_ion_coefficients_with_levels = (
+                (
+                    level_population_fraction.loc[coll_ion_coeff.index]
+                    * coll_ion_coeff
+                )
+                .groupby(level=("atomic_number", "ion_number", "level_number"))
+                .sum()
+            )
+            total_coll_recomb_coefficients_with_levels = (
+                (coll_recomb_coeff)
+                .groupby(level=("atomic_number", "ion_number", "level_number"))
+                .sum()
+            )
+
+            total_photo_ion_coefficients = (
+                NLTERateEquationSolver.prepare_coefficient_matrices_excitation(
+                    rate_matrix_index,
+                    total_photo_ion_coefficients,
+                    total_photo_ion_coefficients_with_levels,
+                )
+            )
+            total_rad_recomb_coefficients = (
+                NLTERateEquationSolver.prepare_coefficient_matrices_excitation(
+                    rate_matrix_index,
+                    total_rad_recomb_coefficients,
+                    total_rad_recomb_coefficients_with_levels,
+                )
+            )
+            total_coll_ion_coefficients = (
+                NLTERateEquationSolver.prepare_coefficient_matrices_excitation(
+                    rate_matrix_index,
+                    total_coll_ion_coefficients,
+                    total_coll_ion_coefficients_with_levels,
+                )
+            )
+            total_coll_recomb_coefficients = (
+                NLTERateEquationSolver.prepare_coefficient_matrices_excitation(
+                    rate_matrix_index,
+                    total_coll_recomb_coefficients,
+                    total_coll_recomb_coefficients_with_levels,
+                )
+            )
+
         return (
             total_photo_ion_coefficients,
             total_rad_recomb_coefficients,
@@ -555,7 +620,11 @@ def deriv_matrix_block(
         return np.dot(deriv_matrix, current_ion_number_densities)
 
     def prepare_first_guess(
-        self, atomic_numbers, number_density, electron_density
+        self,
+        rate_matrix_index,
+        atomic_numbers,
+        number_density,
+        electron_density,
     ):
         """Constructs a first guess for ion number densities and electron density, where all species are singly ionized.
 
@@ -573,13 +642,12 @@ def prepare_first_guess(
         numpy.array
             Guess for ion number densities and electron density.
         """
-        # TODO needs to be changed for excitation
-        array_size = (number_density.index.values + 1).sum() + 1
-        first_guess = np.zeros(array_size)
-        index = 1
+        first_guess = pd.Series(0.0, index=rate_matrix_index)
         for atomic_number in atomic_numbers:
-            first_guess[index] = number_density.loc[atomic_number]
-            index += atomic_number + 1
+            first_guess.loc[(atomic_number, 1)][0] = number_density.loc[
+                atomic_number
+            ]
+        first_guess = first_guess.values
         first_guess[-1] = electron_density
         return first_guess
 
@@ -699,3 +767,170 @@ def prepare_solution_vector(self, number_density):
             )
         solution_vector = np.hstack(solution_array + [0])
         return solution_vector
+
+    @staticmethod
+    def prepare_coefficient_matrices_excitation(
+        rate_matrix_index,
+        coeff_matrix_without_exc,
+        coeff_matrix_with_exc,
+    ):
+        coeff_array = np.zeros(
+            (rate_matrix_index.size, coeff_matrix_without_exc.shape[1])
+        )
+        size = 0
+        for i in range(rate_matrix_index.size):
+            if rate_matrix_index[i][2] != "lte_ion":
+                size += 1
+        size -= 1
+        coeff_array = np.zeros((size, coeff_matrix_without_exc.shape[1]))
+        index = []
+        row = 0
+        for i in range(rate_matrix_index.size):
+            if rate_matrix_index[i][2] == "nlte_ion":
+                coeff_array[row] = coeff_matrix_without_exc.loc[
+                    rate_matrix_index[i][:-1]
+                ]
+                index.append(rate_matrix_index[i])
+            elif (
+                rate_matrix_index[i][2] == "lte_ion"
+                or rate_matrix_index[i][2] == "n_e"
+            ):
+                continue
+            else:
+                coeff_array[row] = coeff_matrix_with_exc.loc[
+                    rate_matrix_index[i]
+                ]
+                index.append(rate_matrix_index[i])
+            row += 1
+        index = pd.MultiIndex.from_tuples(
+            index, names=["atomic_number", "ion_number", "level_number"]
+        )
+        coeff_matrix = pd.DataFrame(
+            coeff_array, columns=coeff_matrix_without_exc.columns, index=index
+        )
+        return coeff_matrix
+
+    @staticmethod
+    def main_nlte_calculation_bound_bound(
+        atomic_data,
+        t_electrons,
+        j_blues,
+        beta_sobolev,
+        excitation_species,
+        nlte_data,
+    ):
+        """Calculates a matrix with bound-bound rates for NLTE excitation treatment.
+
+        Parameters
+        ----------
+        atomic_data : Object
+            Atomic data from the atomic datafile.
+        t_electrons :  Numpy Array, dtype float
+            Electron temperatures.
+        j_blues : Pandas DataFrame, dtype float
+            J_blue values as calculated in LTE.
+        beta_sobolev : Numpy Array, dtype float
+            Sobolev escape probability
+        excitation_species : tuple
+            Species treated in NLTE excitation.
+        nlte_data : NLTEExcitationData
+            Data relevant to NLTE exciation species.
+
+        Returns
+        -------
+        numpy.array
+            Matrix with excitation-deexcitation rates(should be added to NLTE rate matrix for excitation treatment).
+        """
+
+        number_of_levels = atomic_data.levels.energy.loc[
+            excitation_species
+        ].count()
+        lnl = nlte_data.lines_level_number_lower[excitation_species]
+        lnu = nlte_data.lines_level_number_upper[excitation_species]
+        (lines_index,) = nlte_data.lines_idx[excitation_species]
+
+        try:
+            j_blues_filtered = j_blues.iloc[lines_index]
+        except AttributeError:
+            j_blues_filtered = j_blues
+        try:
+            beta_sobolev_filtered = beta_sobolev.iloc[lines_index]
+        except AttributeError:
+            beta_sobolev_filtered = beta_sobolev
+        A_uls = nlte_data.A_uls[excitation_species]
+        B_uls = nlte_data.B_uls[excitation_species]
+        B_lus = nlte_data.B_lus[excitation_species]
+        r_lu_index = lnu * number_of_levels + lnl
+        r_ul_index = lnl * number_of_levels + lnu
+        r_ul_matrix = np.zeros(
+            (number_of_levels, number_of_levels, len(t_electrons)),
+            dtype=np.float64,
+        )
+        r_ul_matrix_reshaped = r_ul_matrix.reshape(
+            (number_of_levels**2, len(t_electrons))
+        )
+        r_ul_matrix_reshaped[r_ul_index] = (
+            A_uls[np.newaxis].T + B_uls[np.newaxis].T * j_blues_filtered
+        )
+        r_ul_matrix_reshaped[r_ul_index] *= beta_sobolev_filtered
+        r_lu_matrix = np.zeros_like(r_ul_matrix)
+        r_lu_matrix_reshaped = r_lu_matrix.reshape(
+            (number_of_levels**2, len(t_electrons))
+        )
+        r_lu_matrix_reshaped[r_lu_index] = (
+            B_lus[np.newaxis].T * j_blues_filtered * beta_sobolev_filtered
+        )
+
+        rates_matrix_bound_bound = r_lu_matrix + r_ul_matrix
+        for i in range(number_of_levels):
+            rates_matrix_bound_bound[i, i] = -rates_matrix_bound_bound[
+                :, i
+            ].sum(axis=0)
+        rates_matrix_bound_bound[0, :, :] = 1.0
+        return rates_matrix_bound_bound
+
+        # TODO: for tests, do the same thing as for the ionization solver, but only H for excitation treatment
+        # TODO: beta sobolev needs to be recalculated for each iteration, because it depends on numberdensity
+
+    def calculate_coll_exc_coefficients(
+        self,
+        coll_exc_coeff,
+        coll_deexc_coeff,
+        nlte_excitation_species,
+    ):
+        """Calculates collision coefficients used for NLTE excitation treatment.
+
+        Parameters
+        ----------
+        coll_exc_coeff : pandas.DataFrame, dtype float
+            Rate coefficient for collisional excitation.
+        coll_deexc_coeff : pandas.DataFrame, dtype float
+            Rate coefficient for collisional deexcitation.
+        nlte_excitation_species : list
+            List of tuples for (atomic number, ion number) which are treated in NLTE excitation.
+
+        Returns
+        -------
+        pandas.DataFrame, dtype float
+            Returns two dataframes, for collision excitation and deexcitation coefficients.
+        """
+        initial_index = pd.MultiIndex.from_tuples(
+            [] * 4, names=coll_exc_coeff.index.names
+        )
+        filtered_coll_exc_coefficients = pd.DataFrame(
+            columns=coll_exc_coeff.columns, index=initial_index
+        )
+        filtered_coll_deexc_coefficients = pd.DataFrame(
+            columns=coll_exc_coeff.columns, index=initial_index
+        )
+        for species in nlte_excitation_species:
+            subset_exc = coll_exc_coeff.loc[[species[0], species[1]]]
+            subset_deexc = coll_deexc_coeff.loc[[species[0], species[1]]]
+            filtered_coll_exc_coefficients = pd.concat(
+                [filtered_coll_exc_coefficients, subset_exc]
+            )
+            filtered_coll_deexc_coefficients = pd.concat(
+                [filtered_coll_deexc_coefficients, subset_deexc]
+            )
+
+        return filtered_coll_exc_coefficients, filtered_coll_deexc_coefficients

tardis/plasma/tests/test_nlte_solver.py

@@ -296,6 +296,7 @@ def test_critical_case_w0(nlte_raw_plasma_w0):
         nlte_raw_plasma_w0.phi,
         nlte_raw_plasma_w0.rate_matrix_index,
         nlte_raw_plasma_w0.number_density,
+        nlte_raw_plasma_w0.nlte_excitation_species,
     )[0]
     ion_number_density_nlte = ion_number_density_nlte.values
     ion_number_density_nlte[ion_number_density_nlte < 1e-10] = 0.0