diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py index aee30448433..d76817b9a77 100644 --- a/tardis/plasma/properties/nlte_rate_equation_solver.py +++ b/tardis/plasma/properties/nlte_rate_equation_solver.py @@ -857,3 +857,70 @@ def prepare_r_uls_r_lus( r_lu_matrix, ) # TODO: beta sobolev needs to be recalculated for each iteration, because it depends on number density + + @staticmethod + def create_coll_exc_deexc_matrix( + coll_exc_coefficient, + coll_deexc_coefficient, + number_of_levels, + ): + """Generates a coefficient matrix from collisional excitation/deexcitation coefficients. + + Needs to be multiplied by electron density when added to the overall rate_matrix. + Parameters + ---------- + coll_exc_coefficient : pandas.Series + Series of collisional excitation coefficients for current (atomic number, ion_number) + in the current shell. + coll_deexc_coefficient : pandas.Series + Series of collisional deexcitation coefficients for (atomic number, ion_number) + in the current shell. + number_of_levels : int + Number of levels for the current atomic number, ion number. + + Returns + ------- + coeff_matrix : np.array (number of levels, number of levels) + Square matrix constructed by collisional exc./deexc. coefficients. + """ + diagonal_exc = np.zeros(number_of_levels) + diagonal_deexc = np.zeros(number_of_levels) + col_exc_coefficient_sum_lower = coll_exc_coefficient.groupby( + "level_number_lower" + ).sum() + col_deexc_coefficient_sum_upper = coll_deexc_coefficient.groupby( + "level_number_upper" + ).sum() + + diagonal_exc[col_exc_coefficient_sum_lower.index] = ( + -1 * col_exc_coefficient_sum_lower.values + ) + diagonal_deexc[col_deexc_coefficient_sum_upper.index] = ( + -1 * col_deexc_coefficient_sum_upper.values + ) + exc_matrix = np.zeros((number_of_levels, number_of_levels)) + deexc_matrix = np.zeros((number_of_levels, number_of_levels)) + exc_matrix[ + ( + coll_exc_coefficient.index.get_level_values( + "level_number_upper" + ), + coll_exc_coefficient.index.get_level_values( + "level_number_lower" + ), + ) + ] = coll_exc_coefficient.values + deexc_matrix[ + ( + coll_exc_coefficient.index.get_level_values( + "level_number_lower" + ), + coll_exc_coefficient.index.get_level_values( + "level_number_upper" + ), + ) + ] = coll_deexc_coefficient.values + np.fill_diagonal(exc_matrix, diagonal_exc) + np.fill_diagonal(deexc_matrix, diagonal_deexc) + coeff_matrix = exc_matrix + deexc_matrix + return coeff_matrix diff --git a/tardis/plasma/tests/test_nlte_excitation.py b/tardis/plasma/tests/test_nlte_excitation.py index fc909970777..657d49711c0 100644 --- a/tardis/plasma/tests/test_nlte_excitation.py +++ b/tardis/plasma/tests/test_nlte_excitation.py @@ -1,5 +1,6 @@ import pandas as pd import numpy as np +import pytest from numpy.testing import assert_allclose @@ -116,3 +117,51 @@ def test_prepare_bound_bound_rate_matrix( np.array(actual_rate_matrix), rtol=1e-6, ) + + +@pytest.mark.parametrize( + [ + "coll_exc_coeff_values", + "coll_deexc_coeff_values", + "number_of_levels", + "desired_coeff_matrix", + ], + [ + ( + [1, -2, 3], + [4, 9, 10], + 3, + [[1.0, 4.0, 9.0], [1.0, -7.0, 10.0], [-2.0, 3.0, -19.0]], + ), + ( + [0.21, 0.045, 0.1234], + [0.7865, 0.987, 0.00123], + 3, + [ + [-0.255, 0.7865, 0.987], + [0.21, -0.9099, 0.00123], + [0.045, 0.1234, -0.98823], + ], + ), + ], +) +def test_coll_exc_deexc_matrix( + coll_exc_coeff_values, + coll_deexc_coeff_values, + number_of_levels, + desired_coeff_matrix, +): + """ + Checks the NLTERateEquationSolver.create_coll_exc_deexc_matrix for simple values of species with 3 levels. + NOTE: Values used for testing are not physical. + """ + index = pd.MultiIndex.from_tuples( + [(0, 1), (0, 2), (1, 2)], + names=["level_number_lower", "level_number_upper"], + ) + exc_coeff = pd.Series(coll_exc_coeff_values, index=index) + deexc_coeff = pd.Series(coll_deexc_coeff_values, index=index) + obtained_coeff_matrix = NLTERateEquationSolver.create_coll_exc_deexc_matrix( + exc_coeff, deexc_coeff, number_of_levels + ) + assert_allclose(obtained_coeff_matrix, desired_coeff_matrix)