From ef15f1b3cf36da7dd745349ffa371f0557c0c184 Mon Sep 17 00:00:00 2001 From: Sona Chitchyan Date: Wed, 16 Nov 2022 21:14:41 +0400 Subject: [PATCH] Adding nlte_rate_equation_solver.py (#2140) * added nlte_rate_equation_matrix * adding the rate_equation_matrix * moved rate_matrix into a class NLTERateEquationSolver * fixed a bug in rate_matrix_index, added an example for checking if rate matrix works properly * changed the initial electron density for now * added tests, fixed a small bug * added some doctrings * fixed a typo * added some doctrings * added some doctrings part 2 * added some docstrings part 3 * removed unnecessary import * Update tardis/plasma/properties/nlte_rate_equation_solver.py Co-authored-by: Christian Vogl * changed how atomic number is created from rate_matrix_index * changed call of a function * got rid of unnecessary if statement in set_nlte_ion_rate * renamed last_row to charge_conservation_row * switched 0, 1 to atomic_number and ion_number to make it more readable * swtihed from rates to coefficients * changed the matrix set up to only keep necessary row for nlte_ion * ran black * fixing some doctrings * swtiched from using numbers to index names * switched the return statement to NotImplementedError * changed groupby from 0, 1 to atomic number, ion number * fixed an issue in the test * ran black Co-authored-by: Christian Vogl --- tardis/plasma/properties/__init__.py | 1 + .../properties/nlte_rate_equation_solver.py | 423 ++++++++++++++++++ .../plasma/properties/property_collections.py | 4 +- tardis/plasma/properties/rate_matrix_index.py | 25 +- tardis/plasma/tests/test_nlte_solver.py | 73 +++ 5 files changed, 521 insertions(+), 5 deletions(-) create mode 100644 tardis/plasma/properties/nlte_rate_equation_solver.py create mode 100644 tardis/plasma/tests/test_nlte_solver.py diff --git a/tardis/plasma/properties/__init__.py b/tardis/plasma/properties/__init__.py index 7ec3f10e396..002e45bc1fb 100644 --- a/tardis/plasma/properties/__init__.py +++ b/tardis/plasma/properties/__init__.py @@ -11,3 +11,4 @@ from tardis.plasma.properties.transition_probabilities import * from tardis.plasma.properties.helium_nlte import * from tardis.plasma.properties.rate_matrix_index import * +from tardis.plasma.properties.nlte_rate_equation_solver import * diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py new file mode 100644 index 00000000000..3a12e01f860 --- /dev/null +++ b/tardis/plasma/properties/nlte_rate_equation_solver.py @@ -0,0 +1,423 @@ +import pandas as pd +import numpy as np + +from tardis.plasma.properties.base import ProcessingPlasmaProperty + +__all__ = [ + "NLTERateEquationSolver", +] + + +class NLTERateEquationSolver(ProcessingPlasmaProperty): + outputs = ("ion_number_density_nlte", "electron_densities_nlte") + + def calculate( + self, + gamma, + alpha_sp, + alpha_stim, + coll_ion_coeff, + coll_recomb_coeff, + partition_function, + levels, + level_boltzmann_factor, + phi, + rate_matrix_index, + number_density, + ): + """Calculates ion number densities and electron densities using NLTE ionization. + + Parameters + ---------- + gamma : DataFrame + The rate coefficient for radiative ionization. + alpha_sp : DataFrame + The rate coefficient for spontaneous recombination. + alpha_stim : DataFrame + The rate coefficient for stimulated recombination. + coll_ion_coeff : DataFrame + The rate coefficient for collisional ionization in the Seaton + approximation. + coll_recomb_coeff : DataFrame + The rate coefficient for collisional recombination. + partition_function : DataFrame + General partition function. Indexed by atomic number, ion number. + levels : MultiIndex + (atomic_number, ion_number, level_number) + Index of filtered atomic data. + level_boltzmann_factor : DataFrame + General Boltzmann factor. + phi : DataFrame + Saha Factors. + rate_matrix_index : MultiIndex + (atomic_number, ion_number, level/treatment type) + If ion is treated in LTE ionization, 3rd index is "lte_ion", + if treated in NLTE ionization, 3rd index is "nlte_ion". + number_density : DataFrame + Number density in each shell for each species. + + Returns + ------- + ion_number_densities_nlte: DataFrame + Number density with NLTE ionization treatment. + electron_densities_nlte: Series + Electron density with NLTE ionizaion treatment. + """ + + ( + total_photo_ion_coefficients, + total_rad_recomb_coefficients, + total_coll_ion_coefficients, + total_coll_recomb_coefficients, + ) = self.prepare_ion_recomb_coefficients_nlte_ion( + gamma, + alpha_sp, + alpha_stim, + coll_ion_coeff, + coll_recomb_coeff, + partition_function, + levels, + level_boltzmann_factor, + ) + + # >>>TODO:initial electron density should be included in the initial guess, added in a future PR + initial_electron_density = number_density.sum(axis=0) + # <<< + rate_matrix = self.calculate_rate_matrix( + phi[0], + initial_electron_density[0], + rate_matrix_index, + total_photo_ion_coefficients[0], + total_rad_recomb_coefficients[0], + total_coll_ion_coefficients[0], + total_coll_recomb_coefficients[0], + ) + + raise NotImplementedError( + "NLTE ionization hasn't been fully implemented yet!" + ) + + @staticmethod + def calculate_rate_matrix( + phi_shell, + electron_density, + rate_matrix_index, + total_photo_ion_coefficients, + total_rad_recomb_coefficients, + total_coll_ion_coefficients, + total_coll_recomb_coefficients, + ): + """ + + Parameters + ---------- + phi_shell : DataFrame + Saha Factors in the current shell + electron_density : float + Guess for electron density in the current shell + rate_matrix_index : MultiIndex + Index used for constructing the rate matrix + total_photo_ion_coefficients : DataFrame + Photo ionization coefficients + total_rad_recomb_coefficients : DataFrame + Radiative recombination coefficients(should get multiplied by electron density) + total_coll_ion_coefficients : DataFrame + Collisional ionization coefficients(should get multiplied by electron density) + total_coll_recomb_coefficients : DataFrame + Collisional recombination coefficients (should get multiplied by electron density^2) + + Returns + ------- + DataFrame + Rate matrix used for NLTE solver. + """ + rate_matrix = pd.DataFrame( + 0.0, columns=rate_matrix_index, index=rate_matrix_index + ) + total_rad_recomb_coefficients = ( + total_rad_recomb_coefficients * electron_density + ) + total_coll_ion_coefficients = ( + total_coll_ion_coefficients * electron_density + ) + total_coll_recomb_coefficients = ( + total_coll_recomb_coefficients * electron_density**2 + ) + atomic_numbers = ( + rate_matrix_index.get_level_values("atomic_number") + .unique() + .drop("n_e") + ) # dropping the n_e index + for atomic_number in atomic_numbers: + ion_numbers = rate_matrix.loc[atomic_number].index.get_level_values( + "ion_number" + ) + phi_block = phi_shell.loc[atomic_number] + rate_matrix_block = NLTERateEquationSolver.lte_rate_matrix_block( + phi_block, electron_density + ) + + nlte_ion_numbers = ion_numbers[ + rate_matrix.loc[atomic_number].index.get_level_values( + "level_number" + ) + == "nlte_ion" + ] + # >>> lte_ion_numbers is for future use in NLTE excitation treatment + lte_ion_numbers = ion_numbers[ + rate_matrix.loc[atomic_number].index.get_level_values( + "level_number" + ) + == "lte_ion" + ] + # <<< + for ion_number in nlte_ion_numbers: + rate_matrix_block = NLTERateEquationSolver.set_nlte_ion_rate( + rate_matrix_block, + atomic_number, + ion_number, + total_rad_recomb_coefficients.loc[(atomic_number,)], + total_photo_ion_coefficients.loc[(atomic_number,)], + total_coll_ion_coefficients.loc[(atomic_number,)], + total_coll_recomb_coefficients.loc[(atomic_number,)], + ) + rate_matrix.loc[ + (atomic_number, slice(None)), (atomic_number) + ] = rate_matrix_block + + charge_conservation_row = ( + NLTERateEquationSolver.prepare_charge_conservation_row( + atomic_numbers + ) + ) + rate_matrix.loc[("n_e", slice(None))] = charge_conservation_row + return rate_matrix + + @staticmethod + def set_nlte_ion_rate( + rate_matrix_block, + atomic_number, + ion_number, + total_rad_recomb_coefficients, + total_photo_ion_coefficients, + total_coll_ion_coefficients, + total_coll_recomb_coefficients, + ): + """Calculates the row for the species treated in NLTE ionization + + Parameters + ---------- + rate_matrix_block : numpy.array + The diagonal block corresponding to current atomic number. + atomic_number : int + Current atomic number + ion_number : int + Current ion number + total_rad_recomb_coefficients : DataFrame + Rad. recomb. coefficients for current atomic number + total_photo_ion_coefficients : DataFrame + Photo ion. coefficients for current atomic number + total_coll_ion_coefficients : DataFrame + Coll. ion. coefficients for current atomic number + total_coll_recomb_coefficients : DataFrame + Coll. recomb. coefficients for current atomic number + + Returns + ------- + numpy.array + Rate matrix block with a changed row for NLTE ionization treatment + """ + ion_coefficients = ( + total_photo_ion_coefficients + total_coll_ion_coefficients + ) + recomb_coefficients = ( + total_rad_recomb_coefficients + total_coll_recomb_coefficients + ) + if atomic_number != ion_number: + ion_coeff_matrix_ion_row = NLTERateEquationSolver.ion_matrix( + ion_coefficients, atomic_number, ion_number + ) + recomb_coeff_matrix_ion_row = NLTERateEquationSolver.recomb_matrix( + recomb_coefficients, atomic_number, ion_number + ) + rate_matrix_block[ion_number, :] = ( + ion_coeff_matrix_ion_row + recomb_coeff_matrix_ion_row + ) + return rate_matrix_block + + @staticmethod + def lte_rate_matrix_block(phi_block, electron_density): + """Creates the generic LTE block for rate matrix. + + Parameters + ---------- + phi_block : DataFrame + Saha Factors for current atomic number + electron_density : float + Current guess for electron density + + Returns + ------- + numpy.array + LTE block for rate matrix + """ + lte_rate_vector_block = -1.0 * np.hstack([*phi_block.values, -1.0]) + lte_rate_matrix_block = np.diag(lte_rate_vector_block) + n_e_initial = np.ones(len(phi_block)) * electron_density + n_e_matrix = np.diag(n_e_initial, 1) + lte_rate_matrix_block += n_e_matrix + lte_rate_matrix_block[-1, :] = 1.0 + return lte_rate_matrix_block + + @staticmethod + def prepare_phi(phi): + """ + Makes sure that phi does not have any 0 entries. + """ + phi[phi == 0.0] = 1.0e-10 * phi[phi > 0.0].min().min() + return phi + + @staticmethod + def recomb_matrix(recomb_coefficients, atomic_number, ion_number): + """Constructs a recombination rate matrix from the recombination rates. + + Parameters + ---------- + recomb_coefficients : DataFrame + Recombination coefficients. + atomic_number : int64 + Current atomic number. Used for the dimension of a square matrix. + ion_number : int64 + Current ion number. Used for returning the correct row. + + Returns + ------- + numpy.ndarray + """ + offdiag = np.zeros(atomic_number) + index = recomb_coefficients.index + for i in index: + offdiag[i] = recomb_coefficients.loc[i] + diag = np.hstack([np.zeros(1), -offdiag]) + return (np.diag(diag) + np.diag(offdiag, k=1))[ion_number, :] + + @staticmethod + def ion_matrix(ion_coefficients, atomic_number, ion_number): + """Constructs an ionization rate matrix from the ionization rates. + + Parameters + ---------- + ion_coefficients : DataFrame + Recombination coefficients. + atomic_number : int64 + Current atomic number. Used for the dimension of a square matrix. + ion_number : int64 + Current ion number. Used for returning the correct row. + + Returns + ------- + numpy.ndarray + """ + offdiag = np.zeros(atomic_number) + index = ion_coefficients.index + for i in index: + offdiag[i] = ion_coefficients.loc[i] + diag = np.hstack([-offdiag, np.zeros(1)]) + return (np.diag(diag) + np.diag(offdiag, k=-1))[ion_number, :] + + @staticmethod + def prepare_charge_conservation_row(atomic_numbers): + """Prepares the last row of the rate_matrix. This row corresponds to the charge + density equation.""" + charge_conservation_row = [] + for atomic_number in atomic_numbers: + charge_conservation_row.append(np.arange(0.0, atomic_number + 1)) + charge_conservation_row = np.hstack([*charge_conservation_row, -1]) + # TODO needs to be modified for use in nlte_excitation + return charge_conservation_row + + @staticmethod + def prepare_ion_recomb_coefficients_nlte_ion( + gamma, + alpha_sp, + alpha_stim, + coll_ion_coeff, + coll_recomb_coeff, + partition_function, + levels, + level_boltzmann_factor, + ): + """ + Prepares the ionization and recombination coefficients by grouping them for + ion numbers. + + Parameters + ---------- + gamma : DataFrame + The rate coefficient for radiative ionization. + alpha_sp : DataFrame + The rate coefficient for spontaneous recombination. + alpha_stim : DataFrame + The rate coefficient for stimulated recombination. + coll_ion_coeff : DataFrame + The rate coefficient for collisional ionization in the Seaton + approximation. + coll_recomb_coeff : DataFrame + The rate coefficient for collisional recombination. + partition_function : DataFrame + General partition function. Indexed by atomic number, ion number. + levels : MultiIndex + (atomic_number, ion_number, level_number) + Index of filtered atomic data. + level_boltzmann_factor : DataFrame + General Boltzmann factor. + Returns + ------- + total_photo_ion_coefficients: + Photoinization coefficients grouped by atomic number and ion number. + total_rad_recomb_coefficients: + Radiative recombination coefficients grouped by atomic number and ion number. + total_coll_ion_coefficients: + Collisional ionization coefficients grouped by atomic number and ion number. + total_coll_recomb_coefficients: + Collisional recombination coefficients grouped by atomic number and ion number. + """ + indexer = pd.Series( + np.arange(partition_function.shape[0]), + index=partition_function.index, + ) + _ion2level_idx = indexer.loc[levels.droplevel("level_number")].values + partition_function_broadcast = partition_function.values[_ion2level_idx] + level_population_fraction = pd.DataFrame( + level_boltzmann_factor.values / partition_function_broadcast, + index=levels, + ) + total_photo_ion_coefficients = ( + (level_population_fraction.loc[gamma.index] * gamma) + .groupby(level=("atomic_number", "ion_number")) + .sum() + ) + total_rad_recomb_coefficients = ( + (alpha_sp + alpha_stim) + .groupby(level=["atomic_number", "ion_number"]) + .sum() + ) + total_coll_ion_coefficients = ( + ( + level_population_fraction.loc[coll_ion_coeff.index] + * coll_ion_coeff + ) + .groupby(level=("atomic_number", "ion_number")) + .sum() + ) + total_coll_recomb_coefficients = ( + (coll_recomb_coeff) + .groupby(level=("atomic_number", "ion_number")) + .sum() + ) + return ( + total_photo_ion_coefficients, + total_rad_recomb_coefficients, + total_coll_ion_coefficients, + total_coll_recomb_coefficients, + ) diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index d04bd5f5032..006083b5585 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -58,7 +58,9 @@ class PlasmaPropertyCollection(list): BetaSobolev, ] ) -nlte_solver_properties = PlasmaPropertyCollection([NLTEIndexHelper]) +nlte_solver_properties = PlasmaPropertyCollection( + [NLTEIndexHelper, NLTERateEquationSolver] +) helium_nlte_properties = PlasmaPropertyCollection( [ HeliumNLTE, diff --git a/tardis/plasma/properties/rate_matrix_index.py b/tardis/plasma/properties/rate_matrix_index.py index 2e18e3bc487..4be7552c114 100644 --- a/tardis/plasma/properties/rate_matrix_index.py +++ b/tardis/plasma/properties/rate_matrix_index.py @@ -14,12 +14,27 @@ def __init__(self, plasma_parent, nlte_ionization_species=0): self.nlte_ionization_species = nlte_ionization_species def calculate(self, levels, nlte_ionization_species): + """Generates rate_matrix_index using levels and changing the last index(level) to + "lte_ion" if that ion_number is treated in LTE, "nlte_ion" for NLTE ionizatin and + keeps the levels for the rest. + + Parameters + ---------- + levels : MultiIndex + (Atomic number, Ion number, Level) + nlte_ionization_species : list + List of tuples for (atomic number, ion number) which are treated in NLTE ionization. + + Returns + ------- + MultiIndex + """ nlte_excitation_species = [] # not yet implemented rate_matrix_index = pd.MultiIndex.from_tuples( list( self.calculate_rate_matrix_index( levels, - self.nlte_ionization_species, + nlte_ionization_species, nlte_excitation_species, ) ), @@ -27,11 +42,13 @@ def calculate(self, levels, nlte_ionization_species): ).drop_duplicates() return rate_matrix_index - def calculate_rate_matrix_index(self, levels, nlte_excitation_species=[]): + def calculate_rate_matrix_index( + self, levels, nlte_ionization_species, nlte_excitation_species=[] + ): for level in levels: - if level[:2] in self.nlte_ionization_species: + if level[:2] in nlte_ionization_species: yield (*level[:2], "nlte_ion") - elif (level[:2] not in self.nlte_ionization_species) and ( + elif (level[:2] not in nlte_ionization_species) and ( level[:2] not in nlte_excitation_species ): yield (*level[:2], "lte_ion") diff --git a/tardis/plasma/tests/test_nlte_solver.py b/tardis/plasma/tests/test_nlte_solver.py new file mode 100644 index 00000000000..ab5552f694e --- /dev/null +++ b/tardis/plasma/tests/test_nlte_solver.py @@ -0,0 +1,73 @@ +import pytest +import numpy as np +import pandas as pd +from numpy.testing import assert_almost_equal +from tardis.plasma.properties import NLTERateEquationSolver + + +def test_rate_matrix(): + """ + Using a simple case of nlte_ion for HI and HeII, checks if the calculate_rate_matrix generates the correct data. + """ + + simple_index_nlte_ion = pd.MultiIndex.from_tuples( + [(1, 0), (2, 1)], names=("atomic_number", "ion_number") + ) + simple_index_lte_ion = pd.MultiIndex.from_tuples( + [(1, 1), (2, 1), (2, 2)], names=("atomic_number", "ion_number") + ) + simple_rate_matrix_index = pd.MultiIndex.from_tuples( + [ + (1, 0, "nlte_ion"), + (1, 1, "lte_ion"), + (2, 0, "lte_ion"), + (2, 1, "nlte_ion"), + (2, 2, "lte_ion"), + ("n_e", "n_e", "n_e"), + ], + names=("atomic_number", "ion_number", "level_number"), + ) + + simple_photo_ion_rates = [0.03464792, 0.68099508] + simple_rad_recomb_coeff = [0.43303813, 0.66140309] + simple_col_ion_coeff = [0.19351674, 0.69214007] + simple_col_recomb_coeff = [0.06402515, 0.29785023] + simple_phi = [0.18936306, 0.15726292, 0.79851244] + + simple_photo_ion_rates = pd.DataFrame( + simple_photo_ion_rates, index=simple_index_nlte_ion + ) + simple_rad_recomb_coeff = pd.DataFrame( + simple_rad_recomb_coeff, index=simple_index_nlte_ion + ) + simple_col_ion_coeff = pd.DataFrame( + simple_col_ion_coeff, index=simple_index_nlte_ion + ) + simple_col_recomb_coeff = pd.DataFrame( + simple_col_recomb_coeff, index=simple_index_nlte_ion + ) + simple_phi = pd.DataFrame(simple_phi, index=simple_index_lte_ion) + + simple_electron_density = 0.2219604493076 + + actual_rate_matrix = NLTERateEquationSolver.calculate_rate_matrix( + simple_phi, + simple_electron_density, + simple_rate_matrix_index, + simple_photo_ion_rates, + simple_rad_recomb_coeff, + simple_col_ion_coeff, + simple_col_recomb_coeff, + ) + desired_rate_matrix = [ + [-0.077601, 0.099272, 0.000000, 0.000000, 0.000000, 0.0], + [1.000000, 1.000000, 0.000000, 0.000000, 0.000000, 0.0], + [0.000000, 0.000000, -0.157263, 0.221960, 0.000000, 0.0], + [0.000000, 0.000000, 0.000000, -0.834623, 0.161479, 0.0], + [0.000000, 0.000000, 1.000000, 1.000000, 1.000000, 0.0], + [0.000000, 1.000000, 0.000000, 1.000000, 2.000000, -1.0], + ] + + assert_almost_equal( + desired_rate_matrix, np.array(actual_rate_matrix), decimal=6 + )