Creating collisional excitation/deexcitation coefficient matrix for NLTE excitation treatment

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

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)