Creating collisional excitation/deexcitation coefficient matrix for NLTE excitation treatment

tardis/plasma/properties/nlte_rate_equation_solver.py

@@ -848,3 +848,49 @@ 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.
+        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 = np.zeros(size + 1)
+        deexc_coeff = (
+            coll_deexc_coefficient.swaplevel()
+        )  # brings level_number_upper to first index
+        diagonal_deexc = np.zeros(size + 1)
+        for i in range(size):
+            diagonal_exc[i] = coll_exc_coefficient.loc[i].sum()
+            diagonal_deexc[i + 1] = deexc_coeff.loc[i + 1].sum()
+        exc_matrix = np.zeros((size + 1, size + 1))
+        deexc_matrix = np.zeros((size + 1, size + 1))
+        for i in range(size + 1):
+            for j in range(size + 1):
+                if i == j:
+                    exc_matrix[i, j] = -diagonal_exc[i]
+                    deexc_matrix[i, j] = -diagonal_deexc[i]
+                elif i > j:
+                    exc_matrix[i, j] = coll_exc_coefficient.loc[j, i]
+                elif i < j:
+                    deexc_matrix[i, j] = deexc_coeff.loc[j, i]
+        coeff_matrix = exc_matrix + deexc_matrix
+        return coeff_matrix

tardis/plasma/tests/test_nlte_excitation.py

@@ -116,3 +116,24 @@ def test_prepare_bound_bound_rate_matrix(
         np.array(actual_rate_matrix),
         rtol=1e-6,
     )
+
+
+def test_coll_exc_deexc_matrix():
+    """
+    Checks the NLTERateEquationSolver.create_coll_exc_deexc_matrix for simple values of species with 3 levels.
+    """
+    index = pd.MultiIndex.from_tuples(
+        [(0, 1), (0, 2), (1, 2)],
+        names=["level_number_lower", "level_number_upper"],
+    )
+    exc_values = [1, -2, 3]
+    exc_coeff = pd.DataFrame(exc_values, index=index)
+    deexc_values = [4, 9, 10]
+    deexc_coeff = pd.DataFrame(deexc_values, index=index)
+    obtained_coeff_matrix = NLTERateEquationSolver.create_coll_exc_deexc_matrix(
+        exc_coeff, deexc_coeff
+    )
+    desired_coeff_matrix = np.array(
+        [[1.0, 4.0, 9.0], [1.0, -7.0, 10.0], [-2.0, 3.0, -19.0]]
+    )
+    assert_allclose(obtained_coeff_matrix, desired_coeff_matrix)