Creating collisional excitation/deexcitation coefficient matrix for NLTE excitation treatment

tardis/plasma/properties/nlte_rate_equation_solver.py

@@ -848,3 +848,48 @@ 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
+    ):
+        """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.DataFrame
+            DataFrame of collisional excitation coefficients for current (atomic number, ion_number)
+        coll_deexc_coefficient : pandas.DataFrame
+            DataFrame of collisional deexcitation coefficients for (atomic number, ion_number)
+
+        Returns
+        -------
+        coeff_matrix : np.array (number of levels, number of levels)
+            Square matrix constructed by collisional exc./deexc. coefficients.
+        """
+        size = (
+            coll_exc_coefficient.index.get_level_values("level_number_lower")
+            .unique()
+            .size
+        )
+        diagonal_exc = coll_exc_coefficient.groupby("level_number_lower").sum()
+        diagonal_deexc = coll_deexc_coefficient.groupby(
+            "level_number_upper"
+        ).sum()
+        exc_matrix = np.zeros((size + 1, size + 1))
+        deexc_matrix = np.zeros((size + 1, size + 1))
+        exc_matrix[
+            np.tril_indices(size + 1, k=-1)
+        ] = coll_exc_coefficient.values.reshape(
+            size + 1,
+        )
+        deexc_matrix[
+            np.triu_indices(size + 1, k=1)
+        ] = coll_deexc_coefficient.values.reshape(
+            size + 1,
+        )
+        np.fill_diagonal(exc_matrix, [*diagonal_exc.values * -1, 0])
+        np.fill_diagonal(deexc_matrix, [0, *diagonal_deexc.values * -1])
+        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,45 @@ 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",
+        "desired_coeff_matrix",
+    ],
+    [
+        (
+            [1, -2, 3],
+            [4, 9, 10],
+            [[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],
+            [
+                [-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, 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.DataFrame(coll_exc_coeff_values, index=index)
+    deexc_coeff = pd.DataFrame(coll_deexc_coeff_values, index=index)
+    obtained_coeff_matrix = NLTERateEquationSolver.create_coll_exc_deexc_matrix(
+        exc_coeff, deexc_coeff
+    )
+    assert_allclose(obtained_coeff_matrix, desired_coeff_matrix)