From 4cb0c5aecc02095bb7996cd26b8b860835765de3 Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Mon, 24 Oct 2022 10:26:49 -0400
Subject: [PATCH 01/27] added nlte_rate_equation_matrix

---
 tardis/plasma/properties/nlte_rate_equation_solver.py | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 create mode 100644 tardis/plasma/properties/nlte_rate_equation_solver.py

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..e69de29bb2d

From 2df7398efd6218a67d8a84da5e6886703b7876f1 Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Mon, 24 Oct 2022 10:28:57 -0400
Subject: [PATCH 02/27] adding the rate_equation_matrix

---
 .../properties/nlte_rate_equation_solver.py   | 166 ++++++++++++++++++
 1 file changed, 166 insertions(+)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index e69de29bb2d..cf55c175491 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -0,0 +1,166 @@
+import pandas as pd
+import numpy as np
+
+from tardis.plasma.properties.base import ProcessingPlasmaProperty
+
+__all__ = [
+    "NLTERateEquationMatrix",
+]
+
+
+class NLTERateEquationMatrix(ProcessingPlasmaProperty):
+    outputs = ("rate_equation_matrix",)
+
+    def calculate(
+        self,
+        phi,
+        electron_density,
+        number_density,
+        rate_matrix_index,
+        last_row,
+        rad_recomb_rate_coeff,
+        photo_ion_rates,
+        coll_ion_coefficient,
+        coll_recomb_coefficient,
+    ):
+        rate_matrix = pd.DataFrame(
+            0.0, columns=rate_matrix_index, index=rate_matrix_index
+        )
+        rad_recomb_rates = rad_recomb_rate_coeff * electron_density
+        coll_ion_rates = coll_ion_coefficient * electron_density
+        coll_recomb_rates= coll_recomb_coefficient * electron_density**2
+        atomic_numbers = number_density.index
+        for atomic_number in atomic_numbers:
+            ion_numbers = rate_matrix.loc[atomic_number].index.get_level_values(0)
+            phi_block = phi.loc[atomic_number]
+            rate_matrix_block = self.lte_rate_matrix_block(phi_block, electron_density)
+
+            nlte_ion_numbers = ion_numbers[
+                rate_matrix.loc[atomic_number].index.get_level_values(1) == "nlte_ion"
+            ]
+            lte_ion_numbers = ion_numbers[
+                rate_matrix.loc[atomic_number].index.get_level_values(1) == "lte_ion"
+            ]
+            for ion_number in nlte_ion_numbers:
+                rate_matrix_block = self.set_nlte_ion_rate(
+                    rate_matrix_block,
+                    atomic_number,
+                    ion_number,
+                    rad_recomb_rates.loc[(atomic_number,)],
+                    photo_ion_rates.loc[(atomic_number,)],
+                    coll_ion_rates.loc[(atomic_number,)],
+                    coll_recomb_rates.loc[(atomic_number,)],
+                )
+            rate_matrix.loc[
+                (atomic_number, slice(None)), (atomic_number)
+            ] = rate_matrix_block
+
+        rate_matrix.loc[("n_e", slice(None))] = last_row
+        return rate_matrix
+
+
+    def set_nlte_ion_rate(
+        self,
+        rate_matrix_block,
+        atomic_number,
+        ion_number,
+        radiative_recombination_rate,
+        photo_ion_rates,
+        coll_ion_rate,
+        coll_recomb_rate,
+    ):
+        ion_rates = photo_ion_rates + coll_ion_rate
+        recomb_rate = radiative_recombination_rate + coll_recomb_rate
+        if atomic_number == ion_number:
+            rate_matrix_block[ion_number, :] = 1.0
+        else:
+            ion_rate_matrix = self.ion_matrix(ion_rates, atomic_number)
+            recomb_rate_matrix = self.recomb_matrix(recomb_rate, atomic_number)
+            rate_matrix_block[ion_number, :] = (ion_rate_matrix + recomb_rate_matrix)[
+                ion_number, :
+            ]
+        return rate_matrix_block
+
+
+    def lte_rate_matrix_block(self, phi_block, electron_density):
+        lte_rate_vector_block = np.hstack([-phi_block, 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
+
+
+    def prepare_phi(self, phi):
+        phi[phi == 0.0] = 1.0e-10 * phi[phi > 0.0].min().min()
+        return phi
+
+
+    def recomb_matrix(self, recomb_rate, atomic_number):
+        offdiag = np.zeros(atomic_number)
+        index = recomb_rate.index
+        for i in index:
+            offdiag[i] = recomb_rate[i]
+        diag = np.hstack([np.zeros(1), -offdiag])
+        return np.diag(diag) + np.diag(offdiag, k=1)
+
+
+    def ion_matrix(self, ion_rate, atomic_number):
+        offdiag = np.zeros(atomic_number)
+        index = ion_rate.index
+        for i in index:
+            offdiag[i] = ion_rate[i]
+        diag = np.hstack([-offdiag, np.zeros(1)])
+        return np.diag(diag) + np.diag(offdiag, k=-1)
+
+
+    @staticmethod
+    def prepare_last_row(atomic_numbers):
+        last_row = []
+        for atomic_number in atomic_numbers:
+            last_row.append(np.arange(0.0, atomic_number + 1))
+        last_row = np.hstack([*last_row, -1])
+        return last_row
+
+
+    def prepare_ion_recomb_rates_nlte_ion(
+        self,
+        level_population_fraction,
+        gamma,
+        alpha_sp,
+        alpha_stim,
+        coll_ion_coeff,
+        coll_recomb_coeff,
+        partition_function,
+        levels,
+        level_boltzmann_factor,
+    ):
+        indexer = pd.Series(
+            np.arange(partition_function.shape[0]),
+            index=partition_function.index,
+        )
+        _ion2level_idx = indexer.loc[levels.droplevel(2)].values
+        partition_function_broadcast = partition_function.values[_ion2level_idx]
+        level_population_fraction = pd.DataFrame(
+            level_boltzmann_factor.values / partition_function_broadcast,
+            index=levels,
+        )
+        photo_ion_rate = (
+            (level_population_fraction.loc[gamma.index] * gamma).groupby(level=(0, 1)).sum()
+        )
+        rad_recomb_rate_coeff = (
+            alpha_sp.groupby(level=[0, 1]).sum() + alpha_stim.groupby(level=[0, 1]).sum()
+        )
+        coll_ion_coefficient = (
+            (level_population_fraction.loc[coll_ion_coeff.index] * coll_ion_coeff)
+            .groupby(level=(0, 1))
+            .sum()
+        )
+        coll_recomb_coefficient = (coll_recomb_coeff).groupby(level=(0, 1)).sum()
+        return (
+            photo_ion_rate,
+            rad_recomb_rate_coeff,
+            coll_ion_coefficient,
+            coll_recomb_coefficient,
+        )

From cdc42f2b3c530568a802340a1b74f51a6eaaa2aa Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Tue, 25 Oct 2022 13:51:03 -0400
Subject: [PATCH 03/27] moved rate_matrix into a class NLTERateEquationSolver

---
 tardis/plasma/properties/__init__.py                 |  1 +
 .../plasma/properties/nlte_rate_equation_solver.py   | 12 ++++++++----
 tardis/plasma/properties/property_collections.py     |  2 +-
 tardis/plasma/properties/rate_matrix_index.py        |  1 -
 4 files changed, 10 insertions(+), 6 deletions(-)

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
index cf55c175491..6b61a964150 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -4,14 +4,17 @@
 from tardis.plasma.properties.base import ProcessingPlasmaProperty
 
 __all__ = [
-    "NLTERateEquationMatrix",
+    "NLTERateEquationSolver",
 ]
 
 
-class NLTERateEquationMatrix(ProcessingPlasmaProperty):
-    outputs = ("rate_equation_matrix",)
+class NLTERateEquationSolver(ProcessingPlasmaProperty):
+    outputs = ("ion_number_density_nlte", "n_e")
+    def calculate(self, *args, **kwargs):
+        return super().calculate(*args, **kwargs)
 
-    def calculate(
+
+    def calculate_rate_matrix(
         self,
         phi,
         electron_density,
@@ -23,6 +26,7 @@ def calculate(
         coll_ion_coefficient,
         coll_recomb_coefficient,
     ):
+        1/0
         rate_matrix = pd.DataFrame(
             0.0, columns=rate_matrix_index, index=rate_matrix_index
         )
diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py
index d04bd5f5032..ffc65912080 100644
--- a/tardis/plasma/properties/property_collections.py
+++ b/tardis/plasma/properties/property_collections.py
@@ -58,7 +58,7 @@ 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..e072dc74e59 100644
--- a/tardis/plasma/properties/rate_matrix_index.py
+++ b/tardis/plasma/properties/rate_matrix_index.py
@@ -19,7 +19,6 @@ def calculate(self, levels, nlte_ionization_species):
             list(
                 self.calculate_rate_matrix_index(
                     levels,
-                    self.nlte_ionization_species,
                     nlte_excitation_species,
                 )
             ),

From ebd513a82b1fda20ad00d9b43004333e29285b79 Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Wed, 26 Oct 2022 18:20:52 -0400
Subject: [PATCH 04/27] fixed a bug in rate_matrix_index, added an example for
 checking if rate matrix works properly

---
 .../properties/nlte_rate_equation_solver.py   | 47 ++++++++++++-------
 tardis/plasma/properties/rate_matrix_index.py |  7 +--
 2 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index 6b61a964150..e25fce3a3d6 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -9,34 +9,49 @@
 
 
 class NLTERateEquationSolver(ProcessingPlasmaProperty):
-    outputs = ("ion_number_density_nlte", "n_e")
-    def calculate(self, *args, **kwargs):
-        return super().calculate(*args, **kwargs)
+    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,
+        ):
+
+        photo_ion_rate, rad_recomb_rate_coeff, coll_ion_coefficient, coll_recomb_coefficient = self.prepare_ion_recomb_rates_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 = np.zeros(len(phi.columns))
+        #<<<
+        rate_matrix = self.calculate_rate_matrix(phi[0], initial_electron_density[0], rate_matrix_index, photo_ion_rate[0], rad_recomb_rate_coeff[0], coll_ion_coefficient[0], coll_recomb_coefficient[0])
+        return None
 
 
     def calculate_rate_matrix(
         self,
-        phi,
+        phi_shell,
         electron_density,
-        number_density,
         rate_matrix_index,
-        last_row,
+        photo_ion_rate,
         rad_recomb_rate_coeff,
-        photo_ion_rates,
         coll_ion_coefficient,
         coll_recomb_coefficient,
     ):
-        1/0
         rate_matrix = pd.DataFrame(
             0.0, columns=rate_matrix_index, index=rate_matrix_index
         )
         rad_recomb_rates = rad_recomb_rate_coeff * electron_density
         coll_ion_rates = coll_ion_coefficient * electron_density
         coll_recomb_rates= coll_recomb_coefficient * electron_density**2
-        atomic_numbers = number_density.index
+        atomic_numbers = rate_matrix_index.get_level_values(0).unique()[:-1].values #dropping the n_e index
         for atomic_number in atomic_numbers:
             ion_numbers = rate_matrix.loc[atomic_number].index.get_level_values(0)
-            phi_block = phi.loc[atomic_number]
+            phi_block = phi_shell.loc[atomic_number]
             rate_matrix_block = self.lte_rate_matrix_block(phi_block, electron_density)
 
             nlte_ion_numbers = ion_numbers[
@@ -51,14 +66,15 @@ def calculate_rate_matrix(
                     atomic_number,
                     ion_number,
                     rad_recomb_rates.loc[(atomic_number,)],
-                    photo_ion_rates.loc[(atomic_number,)],
+                    photo_ion_rate.loc[(atomic_number,)],
                     coll_ion_rates.loc[(atomic_number,)],
                     coll_recomb_rates.loc[(atomic_number,)],
                 )
             rate_matrix.loc[
                 (atomic_number, slice(None)), (atomic_number)
             ] = rate_matrix_block
-
+        
+        last_row = self.prepare_last_row(atomic_numbers)
         rate_matrix.loc[("n_e", slice(None))] = last_row
         return rate_matrix
 
@@ -87,7 +103,7 @@ def set_nlte_ion_rate(
 
 
     def lte_rate_matrix_block(self, phi_block, electron_density):
-        lte_rate_vector_block = np.hstack([-phi_block, 1.0])
+        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)
@@ -125,12 +141,11 @@ def prepare_last_row(atomic_numbers):
         for atomic_number in atomic_numbers:
             last_row.append(np.arange(0.0, atomic_number + 1))
         last_row = np.hstack([*last_row, -1])
+        #TODO needs to be modified for use in nlte_excitation
         return last_row
 
-
+    @staticmethod
     def prepare_ion_recomb_rates_nlte_ion(
-        self,
-        level_population_fraction,
         gamma,
         alpha_sp,
         alpha_stim,
diff --git a/tardis/plasma/properties/rate_matrix_index.py b/tardis/plasma/properties/rate_matrix_index.py
index e072dc74e59..871a5f2c7c0 100644
--- a/tardis/plasma/properties/rate_matrix_index.py
+++ b/tardis/plasma/properties/rate_matrix_index.py
@@ -19,6 +19,7 @@ def calculate(self, levels, nlte_ionization_species):
             list(
                 self.calculate_rate_matrix_index(
                     levels,
+                    nlte_ionization_species,
                     nlte_excitation_species,
                 )
             ),
@@ -26,11 +27,11 @@ 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")

From 2dabff3cc2230591b743a6cff4cebe6fb8dfbc9f Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Thu, 27 Oct 2022 10:55:51 -0400
Subject: [PATCH 05/27] changed the initial electron density for now

---
 tardis/plasma/properties/nlte_rate_equation_solver.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index e25fce3a3d6..6bb0a1d93b3 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -21,15 +21,16 @@ def calculate(self,
         level_boltzmann_factor,
         phi,
         rate_matrix_index,
+        number_density,
         ):
 
         photo_ion_rate, rad_recomb_rate_coeff, coll_ion_coefficient, coll_recomb_coefficient = self.prepare_ion_recomb_rates_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 = np.zeros(len(phi.columns))
+        initial_electron_density = number_density.sum(axis=0)
         #<<<
         rate_matrix = self.calculate_rate_matrix(phi[0], initial_electron_density[0], rate_matrix_index, photo_ion_rate[0], rad_recomb_rate_coeff[0], coll_ion_coefficient[0], coll_recomb_coefficient[0])
-        return None
+        return -1, -1
 
 
     def calculate_rate_matrix(

From d664bce0fdf1b52d6e21f8324f6ce0ee7f8a6568 Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Mon, 31 Oct 2022 15:57:27 -0400
Subject: [PATCH 06/27] added tests, fixed a small bug

---
 .../properties/nlte_rate_equation_solver.py   | 39 ++++++++++---------
 tardis/plasma/tests/test_nlte_solver.py       | 32 +++++++++++++++
 2 files changed, 52 insertions(+), 19 deletions(-)
 create mode 100644 tardis/plasma/tests/test_nlte_solver.py

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index 6bb0a1d93b3..b32d7477f05 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -1,3 +1,4 @@
+import stat
 import pandas as pd
 import numpy as np
 
@@ -24,17 +25,16 @@ def calculate(self,
         number_density,
         ):
 
-        photo_ion_rate, rad_recomb_rate_coeff, coll_ion_coefficient, coll_recomb_coefficient = self.prepare_ion_recomb_rates_nlte_ion(gamma, alpha_sp,     alpha_stim, coll_ion_coeff, coll_recomb_coeff, partition_function, levels, level_boltzmann_factor)
+        photo_ion_rate, rad_recomb_rate_coeff, coll_ion_coefficient, coll_recomb_coefficient = NLTERateEquationSolver.prepare_ion_recomb_rates_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, photo_ion_rate[0], rad_recomb_rate_coeff[0], coll_ion_coefficient[0], coll_recomb_coefficient[0])
+        rate_matrix = NLTERateEquationSolver.calculate_rate_matrix(phi[0], initial_electron_density[0], rate_matrix_index, photo_ion_rate[0], rad_recomb_rate_coeff[0], coll_ion_coefficient[0], coll_recomb_coefficient[0])
         return -1, -1
 
-
+    @staticmethod
     def calculate_rate_matrix(
-        self,
         phi_shell,
         electron_density,
         rate_matrix_index,
@@ -53,7 +53,7 @@ def calculate_rate_matrix(
         for atomic_number in atomic_numbers:
             ion_numbers = rate_matrix.loc[atomic_number].index.get_level_values(0)
             phi_block = phi_shell.loc[atomic_number]
-            rate_matrix_block = self.lte_rate_matrix_block(phi_block, electron_density)
+            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(1) == "nlte_ion"
@@ -62,7 +62,7 @@ def calculate_rate_matrix(
                 rate_matrix.loc[atomic_number].index.get_level_values(1) == "lte_ion"
             ]
             for ion_number in nlte_ion_numbers:
-                rate_matrix_block = self.set_nlte_ion_rate(
+                rate_matrix_block = NLTERateEquationSolver.set_nlte_ion_rate(
                     rate_matrix_block,
                     atomic_number,
                     ion_number,
@@ -75,13 +75,12 @@ def calculate_rate_matrix(
                 (atomic_number, slice(None)), (atomic_number)
             ] = rate_matrix_block
         
-        last_row = self.prepare_last_row(atomic_numbers)
+        last_row = NLTERateEquationSolver.prepare_last_row(atomic_numbers)
         rate_matrix.loc[("n_e", slice(None))] = last_row
         return rate_matrix
 
-
+    @staticmethod
     def set_nlte_ion_rate(
-        self,
         rate_matrix_block,
         atomic_number,
         ion_number,
@@ -95,15 +94,16 @@ def set_nlte_ion_rate(
         if atomic_number == ion_number:
             rate_matrix_block[ion_number, :] = 1.0
         else:
-            ion_rate_matrix = self.ion_matrix(ion_rates, atomic_number)
-            recomb_rate_matrix = self.recomb_matrix(recomb_rate, atomic_number)
+            ion_rate_matrix = NLTERateEquationSolver.ion_matrix(ion_rates, atomic_number)
+            recomb_rate_matrix = NLTERateEquationSolver.recomb_matrix(recomb_rate, atomic_number)
             rate_matrix_block[ion_number, :] = (ion_rate_matrix + recomb_rate_matrix)[
                 ion_number, :
             ]
         return rate_matrix_block
 
 
-    def lte_rate_matrix_block(self, phi_block, electron_density):
+    @staticmethod
+    def lte_rate_matrix_block(phi_block, electron_density):
         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
@@ -113,25 +113,26 @@ def lte_rate_matrix_block(self, phi_block, electron_density):
         return lte_rate_matrix_block
 
 
-    def prepare_phi(self, phi):
+    @staticmethod
+    def prepare_phi(phi):
         phi[phi == 0.0] = 1.0e-10 * phi[phi > 0.0].min().min()
         return phi
 
-
-    def recomb_matrix(self, recomb_rate, atomic_number):
+    @staticmethod
+    def recomb_matrix(recomb_rate, atomic_number):
         offdiag = np.zeros(atomic_number)
         index = recomb_rate.index
         for i in index:
-            offdiag[i] = recomb_rate[i]
+            offdiag[i] = recomb_rate.loc[i]
         diag = np.hstack([np.zeros(1), -offdiag])
         return np.diag(diag) + np.diag(offdiag, k=1)
 
-
-    def ion_matrix(self, ion_rate, atomic_number):
+    @staticmethod
+    def ion_matrix(ion_rate, atomic_number):
         offdiag = np.zeros(atomic_number)
         index = ion_rate.index
         for i in index:
-            offdiag[i] = ion_rate[i]
+            offdiag[i] = ion_rate.loc[i]
         diag = np.hstack([-offdiag, np.zeros(1)])
         return np.diag(diag) + np.diag(offdiag, k=-1)
 
diff --git a/tardis/plasma/tests/test_nlte_solver.py b/tardis/plasma/tests/test_nlte_solver.py
new file mode 100644
index 00000000000..811c142b1f8
--- /dev/null
+++ b/tardis/plasma/tests/test_nlte_solver.py
@@ -0,0 +1,32 @@
+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():
+
+    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')])
+
+    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)
+    
+    

From b15be1b6a35fc7f9fd80cf0723272854d9f5f2d5 Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Wed, 2 Nov 2022 11:31:57 -0400
Subject: [PATCH 07/27] added some doctrings

---
 .../properties/nlte_rate_equation_solver.py   | 157 +++++++++++++++---
 .../plasma/properties/property_collections.py |   4 +-
 tardis/plasma/properties/rate_matrix_index.py |   4 +-
 tardis/plasma/tests/test_nlte_solver.py       |  61 +++++--
 4 files changed, 185 insertions(+), 41 deletions(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index b32d7477f05..91240760f2a 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -11,7 +11,9 @@
 
 class NLTERateEquationSolver(ProcessingPlasmaProperty):
     outputs = ("ion_number_density_nlte", "electron_densities_nlte")
-    def calculate(self,
+
+    def calculate(
+        self,
         gamma,
         alpha_sp,
         alpha_stim,
@@ -23,14 +25,36 @@ def calculate(self,
         phi,
         rate_matrix_index,
         number_density,
-        ):
+    ):
+
+        (
+            photo_ion_rate,
+            rad_recomb_rate_coeff,
+            coll_ion_coefficient,
+            coll_recomb_coefficient,
+        ) = NLTERateEquationSolver.prepare_ion_recomb_rates_nlte_ion(
+            gamma,
+            alpha_sp,
+            alpha_stim,
+            coll_ion_coeff,
+            coll_recomb_coeff,
+            partition_function,
+            levels,
+            level_boltzmann_factor,
+        )
 
-        photo_ion_rate, rad_recomb_rate_coeff, coll_ion_coefficient, coll_recomb_coefficient = NLTERateEquationSolver.prepare_ion_recomb_rates_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
+        # >>>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 = NLTERateEquationSolver.calculate_rate_matrix(phi[0], initial_electron_density[0], rate_matrix_index, photo_ion_rate[0], rad_recomb_rate_coeff[0], coll_ion_coefficient[0], coll_recomb_coefficient[0])
+        # <<<
+        rate_matrix = NLTERateEquationSolver.calculate_rate_matrix(
+            phi[0],
+            initial_electron_density[0],
+            rate_matrix_index,
+            photo_ion_rate[0],
+            rad_recomb_rate_coeff[0],
+            coll_ion_coefficient[0],
+            coll_recomb_coefficient[0],
+        )
         return -1, -1
 
     @staticmethod
@@ -43,23 +67,55 @@ def calculate_rate_matrix(
         coll_ion_coefficient,
         coll_recomb_coefficient,
     ):
+        """_summary_
+
+        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
+        photo_ion_rate : DataFrame
+            Photo ionization rates
+        rad_recomb_rate_coeff : DataFrame
+            Radiative recombination coefficients(should get multiplied by electron density)
+        coll_ion_coefficient : DataFrame
+            Collisionional ionization coefficients(should get multiplied by electron density)
+        coll_recomb_coefficient : DataFrame
+            Collisionional 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
         )
         rad_recomb_rates = rad_recomb_rate_coeff * electron_density
         coll_ion_rates = coll_ion_coefficient * electron_density
-        coll_recomb_rates= coll_recomb_coefficient * electron_density**2
-        atomic_numbers = rate_matrix_index.get_level_values(0).unique()[:-1].values #dropping the n_e index
+        coll_recomb_rates = coll_recomb_coefficient * electron_density**2
+        atomic_numbers = (
+            rate_matrix_index.get_level_values(0).unique()[:-1].values
+        )  # dropping the n_e index
         for atomic_number in atomic_numbers:
-            ion_numbers = rate_matrix.loc[atomic_number].index.get_level_values(0)
+            ion_numbers = rate_matrix.loc[atomic_number].index.get_level_values(
+                0
+            )
             phi_block = phi_shell.loc[atomic_number]
-            rate_matrix_block = NLTERateEquationSolver.lte_rate_matrix_block(phi_block, electron_density)
+            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(1) == "nlte_ion"
+                rate_matrix.loc[atomic_number].index.get_level_values(1)
+                == "nlte_ion"
             ]
             lte_ion_numbers = ion_numbers[
-                rate_matrix.loc[atomic_number].index.get_level_values(1) == "lte_ion"
+                rate_matrix.loc[atomic_number].index.get_level_values(1)
+                == "lte_ion"
             ]
             for ion_number in nlte_ion_numbers:
                 rate_matrix_block = NLTERateEquationSolver.set_nlte_ion_rate(
@@ -74,7 +130,7 @@ def calculate_rate_matrix(
             rate_matrix.loc[
                 (atomic_number, slice(None)), (atomic_number)
             ] = rate_matrix_block
-        
+
         last_row = NLTERateEquationSolver.prepare_last_row(atomic_numbers)
         rate_matrix.loc[("n_e", slice(None))] = last_row
         return rate_matrix
@@ -89,21 +145,62 @@ def set_nlte_ion_rate(
         coll_ion_rate,
         coll_recomb_rate,
     ):
+        """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
+        radiative_recombination_rate : DataFrame
+            Rad. recomb. rate for current atomic number
+        photo_ion_rates : DataFrame
+            Photo ion. rate for current atomic number
+        coll_ion_rate : _type_
+            Coll. ion. rate for current atomic number
+        coll_recomb_rate : _type_
+            Coll. recomb. rate for current atomic number
+
+        Returns
+        -------
+        numpy.array
+            Rate matrix block with a changed row for NLTE ionization treatment
+        """
         ion_rates = photo_ion_rates + coll_ion_rate
         recomb_rate = radiative_recombination_rate + coll_recomb_rate
         if atomic_number == ion_number:
             rate_matrix_block[ion_number, :] = 1.0
         else:
-            ion_rate_matrix = NLTERateEquationSolver.ion_matrix(ion_rates, atomic_number)
-            recomb_rate_matrix = NLTERateEquationSolver.recomb_matrix(recomb_rate, atomic_number)
-            rate_matrix_block[ion_number, :] = (ion_rate_matrix + recomb_rate_matrix)[
-                ion_number, :
-            ]
+            ion_rate_matrix = NLTERateEquationSolver.ion_matrix(
+                ion_rates, atomic_number
+            )
+            recomb_rate_matrix = NLTERateEquationSolver.recomb_matrix(
+                recomb_rate, atomic_number
+            )
+            rate_matrix_block[ion_number, :] = (
+                ion_rate_matrix + recomb_rate_matrix
+            )[ion_number, :]
         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
@@ -112,7 +209,6 @@ def lte_rate_matrix_block(phi_block, electron_density):
         lte_rate_matrix_block[-1, :] = 1.0
         return lte_rate_matrix_block
 
-
     @staticmethod
     def prepare_phi(phi):
         phi[phi == 0.0] = 1.0e-10 * phi[phi > 0.0].min().min()
@@ -136,14 +232,13 @@ def ion_matrix(ion_rate, atomic_number):
         diag = np.hstack([-offdiag, np.zeros(1)])
         return np.diag(diag) + np.diag(offdiag, k=-1)
 
-
     @staticmethod
     def prepare_last_row(atomic_numbers):
         last_row = []
         for atomic_number in atomic_numbers:
             last_row.append(np.arange(0.0, atomic_number + 1))
         last_row = np.hstack([*last_row, -1])
-        #TODO needs to be modified for use in nlte_excitation
+        # TODO needs to be modified for use in nlte_excitation
         return last_row
 
     @staticmethod
@@ -168,17 +263,25 @@ def prepare_ion_recomb_rates_nlte_ion(
             index=levels,
         )
         photo_ion_rate = (
-            (level_population_fraction.loc[gamma.index] * gamma).groupby(level=(0, 1)).sum()
+            (level_population_fraction.loc[gamma.index] * gamma)
+            .groupby(level=(0, 1))
+            .sum()
         )
         rad_recomb_rate_coeff = (
-            alpha_sp.groupby(level=[0, 1]).sum() + alpha_stim.groupby(level=[0, 1]).sum()
+            alpha_sp.groupby(level=[0, 1]).sum()
+            + alpha_stim.groupby(level=[0, 1]).sum()
         )
         coll_ion_coefficient = (
-            (level_population_fraction.loc[coll_ion_coeff.index] * coll_ion_coeff)
+            (
+                level_population_fraction.loc[coll_ion_coeff.index]
+                * coll_ion_coeff
+            )
             .groupby(level=(0, 1))
             .sum()
         )
-        coll_recomb_coefficient = (coll_recomb_coeff).groupby(level=(0, 1)).sum()
+        coll_recomb_coefficient = (
+            (coll_recomb_coeff).groupby(level=(0, 1)).sum()
+        )
         return (
             photo_ion_rate,
             rad_recomb_rate_coeff,
diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py
index ffc65912080..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, NLTERateEquationSolver])
+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 871a5f2c7c0..17147ff49f0 100644
--- a/tardis/plasma/properties/rate_matrix_index.py
+++ b/tardis/plasma/properties/rate_matrix_index.py
@@ -27,7 +27,9 @@ def calculate(self, levels, nlte_ionization_species):
         ).drop_duplicates()
         return rate_matrix_index
 
-    def calculate_rate_matrix_index(self, levels, nlte_ionization_species, nlte_excitation_species=[]):
+    def calculate_rate_matrix_index(
+        self, levels, nlte_ionization_species, nlte_excitation_species=[]
+    ):
         for level in levels:
             if level[:2] in nlte_ionization_species:
                 yield (*level[:2], "nlte_ion")
diff --git a/tardis/plasma/tests/test_nlte_solver.py b/tardis/plasma/tests/test_nlte_solver.py
index 811c142b1f8..40d824dd8d1 100644
--- a/tardis/plasma/tests/test_nlte_solver.py
+++ b/tardis/plasma/tests/test_nlte_solver.py
@@ -4,11 +4,25 @@
 from numpy.testing import assert_almost_equal
 from tardis.plasma.properties import NLTERateEquationSolver
 
+
 def test_rate_matrix():
 
-    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')])
+    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"),
+        ]
+    )
 
     simple_photo_ion_rates = [0.03464792, 0.68099508]
     simple_rad_recomb_coeff = [0.43303813, 0.66140309]
@@ -16,17 +30,40 @@ def test_rate_matrix():
     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_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]]
+    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)
-    
-    
+    assert_almost_equal(
+        desired_rate_matrix, np.array(actual_rate_matrix), decimal=6
+    )

From c3cf82b05fa36cbf760497998286f3050a4c335f Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Thu, 3 Nov 2022 15:20:46 -0400
Subject: [PATCH 08/27] fixed a typo

---
 tardis/plasma/properties/nlte_rate_equation_solver.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index 91240760f2a..9076e501414 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -67,7 +67,7 @@ def calculate_rate_matrix(
         coll_ion_coefficient,
         coll_recomb_coefficient,
     ):
-        """_summary_
+        """
 
         Parameters
         ----------

From dadbb0a695004cee3f96063627d7a92903104334 Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Mon, 7 Nov 2022 10:33:24 -0500
Subject: [PATCH 09/27] added some doctrings

---
 .../properties/nlte_rate_equation_solver.py   | 57 ++++++++++++++++++-
 1 file changed, 56 insertions(+), 1 deletion(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index 9076e501414..6ae29ee308b 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -26,6 +26,27 @@ def calculate(
         rate_matrix_index,
         number_density,
     ):
+        """Calculates ion number densities and electron densities using NLTE ionization.
+
+        Parameters
+        ----------
+        gamma : DataFrame
+        alpha_sp : DataFrame
+        alpha_stim : DataFrame
+        coll_ion_coeff : DataFrame
+        coll_recomb_coeff : DataFrame
+        partition_function : DataFrame
+        levels : MultiIndex
+        level_boltzmann_factor : DataFrame
+        phi : DataFrame
+        rate_matrix_index : MultiIndex
+        number_density : DataFrame
+
+        Returns
+        -------
+        ion_number_densities_nlte: DataFrame
+        electron_densities_nlte: Series
+        """
 
         (
             photo_ion_rate,
@@ -55,7 +76,7 @@ def calculate(
             coll_ion_coefficient[0],
             coll_recomb_coefficient[0],
         )
-        return -1, -1
+        return -1, -1 #function still empty, that's why return statement is arbitrary at this point
 
     @staticmethod
     def calculate_rate_matrix(
@@ -211,11 +232,27 @@ def lte_rate_matrix_block(phi_block, electron_density):
 
     @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_rate, atomic_number):
+        """Constructs a recombination rate matrix from the recombination rates.
+
+        Parameters
+        ----------
+        recomb_rate : DataFrame
+            Recombination rates.
+        atomic_number : int64
+            Current atomic number. Used for the dimension of a square matrix.
+
+        Returns
+        -------
+        numpy.ndarray
+        """
         offdiag = np.zeros(atomic_number)
         index = recomb_rate.index
         for i in index:
@@ -225,6 +262,19 @@ def recomb_matrix(recomb_rate, atomic_number):
 
     @staticmethod
     def ion_matrix(ion_rate, atomic_number):
+        """Constructs an ionization rate matrix from the ionization rates.
+
+        Parameters
+        ----------
+        recomb_rate : DataFrame
+            Recombination rates.
+        atomic_number : int64
+            Current atomic number. Used for the dimension of a square matrix.
+
+        Returns
+        -------
+        numpy.ndarray
+        """
         offdiag = np.zeros(atomic_number)
         index = ion_rate.index
         for i in index:
@@ -234,6 +284,8 @@ def ion_matrix(ion_rate, atomic_number):
 
     @staticmethod
     def prepare_last_row(atomic_numbers):
+        """Prepares the last row of the rate_matrix. This row corresponds to the charge density equation.
+        """
         last_row = []
         for atomic_number in atomic_numbers:
             last_row.append(np.arange(0.0, atomic_number + 1))
@@ -252,6 +304,9 @@ def prepare_ion_recomb_rates_nlte_ion(
         levels,
         level_boltzmann_factor,
     ):
+        """
+        Prepares the ionization and recombination rates/coefficients by grouping them for ion numbers.
+        """
         indexer = pd.Series(
             np.arange(partition_function.shape[0]),
             index=partition_function.index,

From 6131cb722e27aa4379f36e327f28f843900f96df Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Mon, 7 Nov 2022 10:35:16 -0500
Subject: [PATCH 10/27] added some doctrings part 2

---
 tardis/plasma/tests/test_nlte_solver.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tardis/plasma/tests/test_nlte_solver.py b/tardis/plasma/tests/test_nlte_solver.py
index 40d824dd8d1..1095617d9d1 100644
--- a/tardis/plasma/tests/test_nlte_solver.py
+++ b/tardis/plasma/tests/test_nlte_solver.py
@@ -6,6 +6,9 @@
 
 
 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")

From b6d0769355a9b3f9a3703436074a15c737cef38b Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Mon, 7 Nov 2022 10:38:45 -0500
Subject: [PATCH 11/27] added some docstrings part 3

---
 .../plasma/properties/nlte_rate_equation_solver.py  |  8 +++++---
 tardis/plasma/properties/rate_matrix_index.py       | 13 +++++++++++++
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index 6ae29ee308b..ddb78f36692 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -76,7 +76,10 @@ def calculate(
             coll_ion_coefficient[0],
             coll_recomb_coefficient[0],
         )
-        return -1, -1 #function still empty, that's why return statement is arbitrary at this point
+        return (
+            -1,
+            -1,
+        )  # function still empty, that's why return statement is arbitrary at this point
 
     @staticmethod
     def calculate_rate_matrix(
@@ -284,8 +287,7 @@ def ion_matrix(ion_rate, atomic_number):
 
     @staticmethod
     def prepare_last_row(atomic_numbers):
-        """Prepares the last row of the rate_matrix. This row corresponds to the charge density equation.
-        """
+        """Prepares the last row of the rate_matrix. This row corresponds to the charge density equation."""
         last_row = []
         for atomic_number in atomic_numbers:
             last_row.append(np.arange(0.0, atomic_number + 1))
diff --git a/tardis/plasma/properties/rate_matrix_index.py b/tardis/plasma/properties/rate_matrix_index.py
index 17147ff49f0..8836765290f 100644
--- a/tardis/plasma/properties/rate_matrix_index.py
+++ b/tardis/plasma/properties/rate_matrix_index.py
@@ -14,6 +14,19 @@ 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(

From 69778e1cae6523136bbecbdee6ab5da67ba88d7c Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Mon, 7 Nov 2022 10:40:41 -0500
Subject: [PATCH 12/27] removed unnecessary import

---
 tardis/plasma/properties/nlte_rate_equation_solver.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index ddb78f36692..aa712a13550 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -1,4 +1,3 @@
-import stat
 import pandas as pd
 import numpy as np
 

From db4896d1b95165321e5598fc409dc51b501c1ef4 Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Mon, 7 Nov 2022 10:42:40 -0500
Subject: [PATCH 13/27] Update
 tardis/plasma/properties/nlte_rate_equation_solver.py

Co-authored-by: Christian Vogl <cvogl@mpa-garching.mpg.de>
---
 tardis/plasma/properties/nlte_rate_equation_solver.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index aa712a13550..6b65a0a099e 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -107,7 +107,7 @@ def calculate_rate_matrix(
         coll_ion_coefficient : DataFrame
             Collisionional ionization coefficients(should get multiplied by electron density)
         coll_recomb_coefficient : DataFrame
-            Collisionional recombination coefficients(should get multiplied by electron density^2)
+            Collisional recombination coefficients (should get multiplied by electron density^2)
 
         Returns
         -------

From 66c67b787c97f30d60df7b4153cbde9e888f068e Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Mon, 7 Nov 2022 10:47:18 -0500
Subject: [PATCH 14/27] changed how atomic number is created from
 rate_matrix_index

---
 tardis/plasma/properties/nlte_rate_equation_solver.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index aa712a13550..92efd6d1705 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -121,7 +121,7 @@ def calculate_rate_matrix(
         coll_ion_rates = coll_ion_coefficient * electron_density
         coll_recomb_rates = coll_recomb_coefficient * electron_density**2
         atomic_numbers = (
-            rate_matrix_index.get_level_values(0).unique()[:-1].values
+            rate_matrix_index.get_level_values(0).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(

From 023b5cfd1e08dc9079a4136f0953c538021b0ffb Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Mon, 7 Nov 2022 10:55:21 -0500
Subject: [PATCH 15/27] changed call of a function

---
 tardis/plasma/properties/nlte_rate_equation_solver.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index b319ac475f1..ebedab5edf9 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -66,7 +66,7 @@ def calculate(
         # >>>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 = NLTERateEquationSolver.calculate_rate_matrix(
+        rate_matrix = self.calculate_rate_matrix(
             phi[0],
             initial_electron_density[0],
             rate_matrix_index,
@@ -182,9 +182,9 @@ def set_nlte_ion_rate(
             Rad. recomb. rate for current atomic number
         photo_ion_rates : DataFrame
             Photo ion. rate for current atomic number
-        coll_ion_rate : _type_
+        coll_ion_rate : DataFrame
             Coll. ion. rate for current atomic number
-        coll_recomb_rate : _type_
+        coll_recomb_rate : DataFrame
             Coll. recomb. rate for current atomic number
 
         Returns

From 313edba8665061f9116d7d84a29df63a889f361a Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Mon, 7 Nov 2022 14:20:30 -0500
Subject: [PATCH 16/27] got rid of unnecessary if statement in
 set_nlte_ion_rate

---
 tardis/plasma/properties/nlte_rate_equation_solver.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index ebedab5edf9..94f124a083e 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -194,9 +194,7 @@ def set_nlte_ion_rate(
         """
         ion_rates = photo_ion_rates + coll_ion_rate
         recomb_rate = radiative_recombination_rate + coll_recomb_rate
-        if atomic_number == ion_number:
-            rate_matrix_block[ion_number, :] = 1.0
-        else:
+        if atomic_number != ion_number:
             ion_rate_matrix = NLTERateEquationSolver.ion_matrix(
                 ion_rates, atomic_number
             )

From 9686a32dd9a851eaacfda4287e21fe5f302c4a97 Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Mon, 7 Nov 2022 14:22:31 -0500
Subject: [PATCH 17/27] renamed last_row to charge_conservation_row

---
 .../plasma/properties/nlte_rate_equation_solver.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index 94f124a083e..71496a0bac2 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -154,8 +154,8 @@ def calculate_rate_matrix(
                 (atomic_number, slice(None)), (atomic_number)
             ] = rate_matrix_block
 
-        last_row = NLTERateEquationSolver.prepare_last_row(atomic_numbers)
-        rate_matrix.loc[("n_e", slice(None))] = last_row
+        charge_conservation_row = NLTERateEquationSolver.prepare_charge_conservation_row(atomic_numbers)
+        rate_matrix.loc[("n_e", slice(None))] = charge_conservation_row
         return rate_matrix
 
     @staticmethod
@@ -283,14 +283,14 @@ def ion_matrix(ion_rate, atomic_number):
         return np.diag(diag) + np.diag(offdiag, k=-1)
 
     @staticmethod
-    def prepare_last_row(atomic_numbers):
+    def prepare_charge_conservation_row(atomic_numbers):
         """Prepares the last row of the rate_matrix. This row corresponds to the charge density equation."""
-        last_row = []
+        charge_conservation_row = []
         for atomic_number in atomic_numbers:
-            last_row.append(np.arange(0.0, atomic_number + 1))
-        last_row = np.hstack([*last_row, -1])
+            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 last_row
+        return charge_conservation_row
 
     @staticmethod
     def prepare_ion_recomb_rates_nlte_ion(

From 06ffa346cd51dbf5524c9dfa2a9b1e3e4ff594dd Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Mon, 7 Nov 2022 14:28:25 -0500
Subject: [PATCH 18/27] switched 0, 1 to atomic_number and ion_number to make
 it more readable

---
 tardis/plasma/properties/nlte_rate_equation_solver.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index 71496a0bac2..b470bb7fbce 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -310,7 +310,7 @@ def prepare_ion_recomb_rates_nlte_ion(
             np.arange(partition_function.shape[0]),
             index=partition_function.index,
         )
-        _ion2level_idx = indexer.loc[levels.droplevel(2)].values
+        _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,
@@ -318,12 +318,11 @@ def prepare_ion_recomb_rates_nlte_ion(
         )
         photo_ion_rate = (
             (level_population_fraction.loc[gamma.index] * gamma)
-            .groupby(level=(0, 1))
+            .groupby(level=("atomic_number", "ion_number"))
             .sum()
         )
         rad_recomb_rate_coeff = (
-            alpha_sp.groupby(level=[0, 1]).sum()
-            + alpha_stim.groupby(level=[0, 1]).sum()
+            (alpha_sp + alpha_stim).groupby(level=["atomic_number", "ion_number"]).sum()
         )
         coll_ion_coefficient = (
             (
@@ -334,7 +333,7 @@ def prepare_ion_recomb_rates_nlte_ion(
             .sum()
         )
         coll_recomb_coefficient = (
-            (coll_recomb_coeff).groupby(level=(0, 1)).sum()
+            (coll_recomb_coeff).groupby(level=("atomic_number", "ion_number")).sum()
         )
         return (
             photo_ion_rate,

From 95fa1fef5f1b9a01b4180def4bb254faabf163e2 Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Mon, 7 Nov 2022 15:24:56 -0500
Subject: [PATCH 19/27] swtihed from rates to coefficients

---
 .../properties/nlte_rate_equation_solver.py   | 130 +++++++++---------
 1 file changed, 66 insertions(+), 64 deletions(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index b470bb7fbce..d6ce8653a8d 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -48,11 +48,11 @@ def calculate(
         """
 
         (
-            photo_ion_rate,
-            rad_recomb_rate_coeff,
-            coll_ion_coefficient,
-            coll_recomb_coefficient,
-        ) = NLTERateEquationSolver.prepare_ion_recomb_rates_nlte_ion(
+            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,
@@ -70,10 +70,10 @@ def calculate(
             phi[0],
             initial_electron_density[0],
             rate_matrix_index,
-            photo_ion_rate[0],
-            rad_recomb_rate_coeff[0],
-            coll_ion_coefficient[0],
-            coll_recomb_coefficient[0],
+            total_photo_ion_coefficients[0],
+            total_rad_recomb_coefficients[0],
+            total_coll_ion_coefficients[0],
+            total_coll_recomb_coefficients[0],
         )
         return (
             -1,
@@ -85,10 +85,10 @@ def calculate_rate_matrix(
         phi_shell,
         electron_density,
         rate_matrix_index,
-        photo_ion_rate,
-        rad_recomb_rate_coeff,
-        coll_ion_coefficient,
-        coll_recomb_coefficient,
+        total_photo_ion_coefficients,
+        total_rad_recomb_coefficients,
+        total_coll_ion_coefficients,
+        total_coll_recomb_coefficients,
     ):
         """
 
@@ -100,13 +100,13 @@ def calculate_rate_matrix(
             Guess for electron density in the current shell
         rate_matrix_index : MultiIndex
             Index used for constructing the rate matrix
-        photo_ion_rate : DataFrame
-            Photo ionization rates
-        rad_recomb_rate_coeff : DataFrame
+        total_photo_ion_coefficients : DataFrame
+            Photo ionization coefficients
+        total_rad_recomb_coefficients : DataFrame
             Radiative recombination coefficients(should get multiplied by electron density)
-        coll_ion_coefficient : DataFrame
+        total_coll_ion_coefficients : DataFrame
             Collisionional ionization coefficients(should get multiplied by electron density)
-        coll_recomb_coefficient : DataFrame
+        total_coll_recomb_coefficients : DataFrame
             Collisional recombination coefficients (should get multiplied by electron density^2)
 
         Returns
@@ -117,9 +117,9 @@ def calculate_rate_matrix(
         rate_matrix = pd.DataFrame(
             0.0, columns=rate_matrix_index, index=rate_matrix_index
         )
-        rad_recomb_rates = rad_recomb_rate_coeff * electron_density
-        coll_ion_rates = coll_ion_coefficient * electron_density
-        coll_recomb_rates = coll_recomb_coefficient * electron_density**2
+        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(0).unique().drop("n_e")
         )  # dropping the n_e index
@@ -136,19 +136,21 @@ def calculate_rate_matrix(
                 rate_matrix.loc[atomic_number].index.get_level_values(1)
                 == "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(1)
                 == "lte_ion"
             ]
+            #<<<
             for ion_number in nlte_ion_numbers:
                 rate_matrix_block = NLTERateEquationSolver.set_nlte_ion_rate(
                     rate_matrix_block,
                     atomic_number,
                     ion_number,
-                    rad_recomb_rates.loc[(atomic_number,)],
-                    photo_ion_rate.loc[(atomic_number,)],
-                    coll_ion_rates.loc[(atomic_number,)],
-                    coll_recomb_rates.loc[(atomic_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)
@@ -163,10 +165,10 @@ def set_nlte_ion_rate(
         rate_matrix_block,
         atomic_number,
         ion_number,
-        radiative_recombination_rate,
-        photo_ion_rates,
-        coll_ion_rate,
-        coll_recomb_rate,
+        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
 
@@ -178,31 +180,31 @@ def set_nlte_ion_rate(
             Current atomic number
         ion_number : int
             Current ion number
-        radiative_recombination_rate : DataFrame
-            Rad. recomb. rate for current atomic number
-        photo_ion_rates : DataFrame
-            Photo ion. rate for current atomic number
-        coll_ion_rate : DataFrame
-            Coll. ion. rate for current atomic number
-        coll_recomb_rate : DataFrame
-            Coll. recomb. rate for current atomic 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_rates = photo_ion_rates + coll_ion_rate
-        recomb_rate = radiative_recombination_rate + coll_recomb_rate
+        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_rate_matrix = NLTERateEquationSolver.ion_matrix(
-                ion_rates, atomic_number
+            ion_coeff_matrix = NLTERateEquationSolver.ion_matrix(
+                ion_coefficients, atomic_number
             )
-            recomb_rate_matrix = NLTERateEquationSolver.recomb_matrix(
-                recomb_rate, atomic_number
+            recomb_coeff_matrix = NLTERateEquationSolver.recomb_matrix(
+                recomb_coefficients, atomic_number
             )
             rate_matrix_block[ion_number, :] = (
-                ion_rate_matrix + recomb_rate_matrix
+                ion_coeff_matrix + recomb_coeff_matrix
             )[ion_number, :]
         return rate_matrix_block
 
@@ -239,13 +241,13 @@ def prepare_phi(phi):
         return phi
 
     @staticmethod
-    def recomb_matrix(recomb_rate, atomic_number):
+    def recomb_matrix(recomb_coefficients, atomic_number):
         """Constructs a recombination rate matrix from the recombination rates.
 
         Parameters
         ----------
-        recomb_rate : DataFrame
-            Recombination rates.
+        recomb_coefficients : DataFrame
+            Recombination coefficients.
         atomic_number : int64
             Current atomic number. Used for the dimension of a square matrix.
 
@@ -254,20 +256,20 @@ def recomb_matrix(recomb_rate, atomic_number):
         numpy.ndarray
         """
         offdiag = np.zeros(atomic_number)
-        index = recomb_rate.index
+        index = recomb_coefficients.index
         for i in index:
-            offdiag[i] = recomb_rate.loc[i]
+            offdiag[i] = recomb_coefficients.loc[i]
         diag = np.hstack([np.zeros(1), -offdiag])
         return np.diag(diag) + np.diag(offdiag, k=1)
 
     @staticmethod
-    def ion_matrix(ion_rate, atomic_number):
+    def ion_matrix(ion_coefficients, atomic_number):
         """Constructs an ionization rate matrix from the ionization rates.
 
         Parameters
         ----------
-        recomb_rate : DataFrame
-            Recombination rates.
+        ion_coefficients : DataFrame
+            Recombination coefficients.
         atomic_number : int64
             Current atomic number. Used for the dimension of a square matrix.
 
@@ -276,9 +278,9 @@ def ion_matrix(ion_rate, atomic_number):
         numpy.ndarray
         """
         offdiag = np.zeros(atomic_number)
-        index = ion_rate.index
+        index = ion_coefficients.index
         for i in index:
-            offdiag[i] = ion_rate.loc[i]
+            offdiag[i] = ion_coefficients.loc[i]
         diag = np.hstack([-offdiag, np.zeros(1)])
         return np.diag(diag) + np.diag(offdiag, k=-1)
 
@@ -293,7 +295,7 @@ def prepare_charge_conservation_row(atomic_numbers):
         return charge_conservation_row
 
     @staticmethod
-    def prepare_ion_recomb_rates_nlte_ion(
+    def prepare_ion_recomb_coefficients_nlte_ion(
         gamma,
         alpha_sp,
         alpha_stim,
@@ -304,7 +306,7 @@ def prepare_ion_recomb_rates_nlte_ion(
         level_boltzmann_factor,
     ):
         """
-        Prepares the ionization and recombination rates/coefficients by grouping them for ion numbers.
+        Prepares the ionization and recombination coefficients by grouping them for ion numbers.
         """
         indexer = pd.Series(
             np.arange(partition_function.shape[0]),
@@ -316,15 +318,15 @@ def prepare_ion_recomb_rates_nlte_ion(
             level_boltzmann_factor.values / partition_function_broadcast,
             index=levels,
         )
-        photo_ion_rate = (
+        total_photo_ion_coefficients = (
             (level_population_fraction.loc[gamma.index] * gamma)
             .groupby(level=("atomic_number", "ion_number"))
             .sum()
         )
-        rad_recomb_rate_coeff = (
+        total_rad_recomb_coefficients = (
             (alpha_sp + alpha_stim).groupby(level=["atomic_number", "ion_number"]).sum()
         )
-        coll_ion_coefficient = (
+        total_coll_ion_coefficients = (
             (
                 level_population_fraction.loc[coll_ion_coeff.index]
                 * coll_ion_coeff
@@ -332,12 +334,12 @@ def prepare_ion_recomb_rates_nlte_ion(
             .groupby(level=(0, 1))
             .sum()
         )
-        coll_recomb_coefficient = (
+        total_coll_recomb_coefficients = (
             (coll_recomb_coeff).groupby(level=("atomic_number", "ion_number")).sum()
         )
         return (
-            photo_ion_rate,
-            rad_recomb_rate_coeff,
-            coll_ion_coefficient,
-            coll_recomb_coefficient,
+            total_photo_ion_coefficients,
+            total_rad_recomb_coefficients,
+            total_coll_ion_coefficients,
+            total_coll_recomb_coefficients,
         )

From 0b8e71dd0f5d1856a978a94f8e64a880a17f39c6 Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Mon, 7 Nov 2022 15:30:47 -0500
Subject: [PATCH 20/27] changed the matrix set up to only keep necessary row
 for nlte_ion

---
 .../properties/nlte_rate_equation_solver.py   | 24 +++++++++++--------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index d6ce8653a8d..6634f3ec3c2 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -197,15 +197,15 @@ def set_nlte_ion_rate(
         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 = NLTERateEquationSolver.ion_matrix(
-                ion_coefficients, atomic_number
+            ion_coeff_matrix_ion_row = NLTERateEquationSolver.ion_matrix(
+                ion_coefficients, atomic_number, ion_number
             )
-            recomb_coeff_matrix = NLTERateEquationSolver.recomb_matrix(
-                recomb_coefficients, atomic_number
+            recomb_coeff_matrix_ion_row = NLTERateEquationSolver.recomb_matrix(
+                recomb_coefficients, atomic_number, ion_number
             )
             rate_matrix_block[ion_number, :] = (
-                ion_coeff_matrix + recomb_coeff_matrix
-            )[ion_number, :]
+                ion_coeff_matrix_ion_row + recomb_coeff_matrix_ion_row
+            )
         return rate_matrix_block
 
     @staticmethod
@@ -241,7 +241,7 @@ def prepare_phi(phi):
         return phi
 
     @staticmethod
-    def recomb_matrix(recomb_coefficients, atomic_number):
+    def recomb_matrix(recomb_coefficients, atomic_number, ion_number):
         """Constructs a recombination rate matrix from the recombination rates.
 
         Parameters
@@ -250,6 +250,8 @@ def recomb_matrix(recomb_coefficients, atomic_number):
             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
         -------
@@ -260,10 +262,10 @@ def recomb_matrix(recomb_coefficients, atomic_number):
         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)
+        return (np.diag(diag) + np.diag(offdiag, k=1))[ion_number, :]
 
     @staticmethod
-    def ion_matrix(ion_coefficients, atomic_number):
+    def ion_matrix(ion_coefficients, atomic_number, ion_number):
         """Constructs an ionization rate matrix from the ionization rates.
 
         Parameters
@@ -272,6 +274,8 @@ def ion_matrix(ion_coefficients, atomic_number):
             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
         -------
@@ -282,7 +286,7 @@ def ion_matrix(ion_coefficients, atomic_number):
         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)
+        return (np.diag(diag) + np.diag(offdiag, k=-1))[ion_number, :]
 
     @staticmethod
     def prepare_charge_conservation_row(atomic_numbers):

From 574f1f359b382cf3002610743d67d8e1587a822b Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Mon, 7 Nov 2022 15:31:27 -0500
Subject: [PATCH 21/27] ran black

---
 .../properties/nlte_rate_equation_solver.py   | 36 ++++++++++++++-----
 1 file changed, 27 insertions(+), 9 deletions(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index 6634f3ec3c2..48f3dc952ba 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -117,9 +117,15 @@ def calculate_rate_matrix(
         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
+        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(0).unique().drop("n_e")
         )  # dropping the n_e index
@@ -141,7 +147,7 @@ def calculate_rate_matrix(
                 rate_matrix.loc[atomic_number].index.get_level_values(1)
                 == "lte_ion"
             ]
-            #<<<
+            # <<<
             for ion_number in nlte_ion_numbers:
                 rate_matrix_block = NLTERateEquationSolver.set_nlte_ion_rate(
                     rate_matrix_block,
@@ -156,7 +162,11 @@ def calculate_rate_matrix(
                 (atomic_number, slice(None)), (atomic_number)
             ] = rate_matrix_block
 
-        charge_conservation_row = NLTERateEquationSolver.prepare_charge_conservation_row(atomic_numbers)
+        charge_conservation_row = (
+            NLTERateEquationSolver.prepare_charge_conservation_row(
+                atomic_numbers
+            )
+        )
         rate_matrix.loc[("n_e", slice(None))] = charge_conservation_row
         return rate_matrix
 
@@ -194,8 +204,12 @@ def set_nlte_ion_rate(
         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
+        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
@@ -328,7 +342,9 @@ def prepare_ion_recomb_coefficients_nlte_ion(
             .sum()
         )
         total_rad_recomb_coefficients = (
-            (alpha_sp + alpha_stim).groupby(level=["atomic_number", "ion_number"]).sum()
+            (alpha_sp + alpha_stim)
+            .groupby(level=["atomic_number", "ion_number"])
+            .sum()
         )
         total_coll_ion_coefficients = (
             (
@@ -339,7 +355,9 @@ def prepare_ion_recomb_coefficients_nlte_ion(
             .sum()
         )
         total_coll_recomb_coefficients = (
-            (coll_recomb_coeff).groupby(level=("atomic_number", "ion_number")).sum()
+            (coll_recomb_coeff)
+            .groupby(level=("atomic_number", "ion_number"))
+            .sum()
         )
         return (
             total_photo_ion_coefficients,

From 31c821a7162c0f1f2f42cb1538bd0e5d69627425 Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Wed, 9 Nov 2022 10:52:42 -0500
Subject: [PATCH 22/27] fixing some doctrings

---
 .../properties/nlte_rate_equation_solver.py   | 58 +++++++++++++++++--
 tardis/plasma/properties/rate_matrix_index.py |  4 +-
 2 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index 48f3dc952ba..e05f4f805d4 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -30,21 +30,38 @@ def calculate(
         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.
         """
 
         (
@@ -105,14 +122,14 @@ def calculate_rate_matrix(
         total_rad_recomb_coefficients : DataFrame
             Radiative recombination coefficients(should get multiplied by electron density)
         total_coll_ion_coefficients : DataFrame
-            Collisionional ionization coefficients(should get multiplied by electron density)
+            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 used for NLTE solver.
         """
         rate_matrix = pd.DataFrame(
             0.0, columns=rate_matrix_index, index=rate_matrix_index
@@ -304,7 +321,8 @@ def ion_matrix(ion_coefficients, atomic_number, 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."""
+        """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))
@@ -324,7 +342,39 @@ def prepare_ion_recomb_coefficients_nlte_ion(
         level_boltzmann_factor,
     ):
         """
-        Prepares the ionization and recombination coefficients by grouping them for ion numbers.
+        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]),
diff --git a/tardis/plasma/properties/rate_matrix_index.py b/tardis/plasma/properties/rate_matrix_index.py
index 8836765290f..0fc7842ff5d 100644
--- a/tardis/plasma/properties/rate_matrix_index.py
+++ b/tardis/plasma/properties/rate_matrix_index.py
@@ -14,7 +14,9 @@ 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.
+        """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
         ----------

From b4d494bf717123e4cd71daef768479ed93ccbaae Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Wed, 9 Nov 2022 11:02:03 -0500
Subject: [PATCH 23/27] swtiched from using numbers to index names

---
 tardis/plasma/properties/nlte_rate_equation_solver.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index e05f4f805d4..928dd7c65e9 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -144,11 +144,11 @@ def calculate_rate_matrix(
             total_coll_recomb_coefficients * electron_density**2
         )
         atomic_numbers = (
-            rate_matrix_index.get_level_values(0).unique().drop("n_e")
+            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(
-                0
+                "ion_number"
             )
             phi_block = phi_shell.loc[atomic_number]
             rate_matrix_block = NLTERateEquationSolver.lte_rate_matrix_block(
@@ -156,12 +156,12 @@ def calculate_rate_matrix(
             )
 
             nlte_ion_numbers = ion_numbers[
-                rate_matrix.loc[atomic_number].index.get_level_values(1)
+                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(1)
+                rate_matrix.loc[atomic_number].index.get_level_values("level_number")
                 == "lte_ion"
             ]
             # <<<

From 420f90209ccccf7b08e5d8dfd75da102f40deaac Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Wed, 9 Nov 2022 11:06:26 -0500
Subject: [PATCH 24/27] switched the return statement to NotImplementedError

---
 tardis/plasma/properties/nlte_rate_equation_solver.py | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index 928dd7c65e9..7b46e43a7ed 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -92,11 +92,9 @@ def calculate(
             total_coll_ion_coefficients[0],
             total_coll_recomb_coefficients[0],
         )
-        return (
-            -1,
-            -1,
-        )  # function still empty, that's why return statement is arbitrary at this point
-
+        
+        raise NotImplementedError("NLTE ionization hasn't been fully implemented yet!") 
+        
     @staticmethod
     def calculate_rate_matrix(
         phi_shell,

From 0307dba3f5f2e731044614164d26088d6d5d6fdd Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Wed, 9 Nov 2022 11:11:04 -0500
Subject: [PATCH 25/27] changed groupby from 0, 1 to atomic number, ion number

---
 tardis/plasma/properties/nlte_rate_equation_solver.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index 7b46e43a7ed..5a7655886d6 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -94,7 +94,7 @@ def calculate(
         )
         
         raise NotImplementedError("NLTE ionization hasn't been fully implemented yet!") 
-        
+
     @staticmethod
     def calculate_rate_matrix(
         phi_shell,
@@ -399,7 +399,7 @@ def prepare_ion_recomb_coefficients_nlte_ion(
                 level_population_fraction.loc[coll_ion_coeff.index]
                 * coll_ion_coeff
             )
-            .groupby(level=(0, 1))
+            .groupby(level=("atomic_number", "ion_number"))
             .sum()
         )
         total_coll_recomb_coefficients = (

From 20ca228812cae4097e13aedf99a9e21481fd2fe9 Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Wed, 9 Nov 2022 11:29:16 -0500
Subject: [PATCH 26/27] fixed an issue in the test

---
 tardis/plasma/tests/test_nlte_solver.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tardis/plasma/tests/test_nlte_solver.py b/tardis/plasma/tests/test_nlte_solver.py
index 1095617d9d1..ab23a9cbd19 100644
--- a/tardis/plasma/tests/test_nlte_solver.py
+++ b/tardis/plasma/tests/test_nlte_solver.py
@@ -24,7 +24,8 @@ def test_rate_matrix():
             (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]

From 121683a5dcb3ba60ea4dd552fae92effa90fffcd Mon Sep 17 00:00:00 2001
From: Sona Chitchyan <chitchyan.sona@gmail.com>
Date: Wed, 9 Nov 2022 11:30:24 -0500
Subject: [PATCH 27/27] ran black

---
 .../properties/nlte_rate_equation_solver.py   | 20 +++++++++++++------
 tardis/plasma/properties/rate_matrix_index.py |  4 ++--
 tardis/plasma/tests/test_nlte_solver.py       |  2 +-
 3 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
index 5a7655886d6..3a12e01f860 100644
--- a/tardis/plasma/properties/nlte_rate_equation_solver.py
+++ b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -92,8 +92,10 @@ def calculate(
             total_coll_ion_coefficients[0],
             total_coll_recomb_coefficients[0],
         )
-        
-        raise NotImplementedError("NLTE ionization hasn't been fully implemented yet!") 
+
+        raise NotImplementedError(
+            "NLTE ionization hasn't been fully implemented yet!"
+        )
 
     @staticmethod
     def calculate_rate_matrix(
@@ -142,7 +144,9 @@ def calculate_rate_matrix(
             total_coll_recomb_coefficients * electron_density**2
         )
         atomic_numbers = (
-            rate_matrix_index.get_level_values("atomic_number").unique().drop("n_e")
+            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(
@@ -154,12 +158,16 @@ def calculate_rate_matrix(
             )
 
             nlte_ion_numbers = ion_numbers[
-                rate_matrix.loc[atomic_number].index.get_level_values("level_number")
+                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")
+                rate_matrix.loc[atomic_number].index.get_level_values(
+                    "level_number"
+                )
                 == "lte_ion"
             ]
             # <<<
@@ -319,7 +327,7 @@ def ion_matrix(ion_coefficients, atomic_number, ion_number):
 
     @staticmethod
     def prepare_charge_conservation_row(atomic_numbers):
-        """Prepares the last row of the rate_matrix. This row corresponds to the charge 
+        """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:
diff --git a/tardis/plasma/properties/rate_matrix_index.py b/tardis/plasma/properties/rate_matrix_index.py
index 0fc7842ff5d..4be7552c114 100644
--- a/tardis/plasma/properties/rate_matrix_index.py
+++ b/tardis/plasma/properties/rate_matrix_index.py
@@ -14,8 +14,8 @@ 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 
+        """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
diff --git a/tardis/plasma/tests/test_nlte_solver.py b/tardis/plasma/tests/test_nlte_solver.py
index ab23a9cbd19..ab5552f694e 100644
--- a/tardis/plasma/tests/test_nlte_solver.py
+++ b/tardis/plasma/tests/test_nlte_solver.py
@@ -25,7 +25,7 @@ def test_rate_matrix():
             (2, 2, "lte_ion"),
             ("n_e", "n_e", "n_e"),
         ],
-        names=('atomic_number', 'ion_number', 'level_number')
+        names=("atomic_number", "ion_number", "level_number"),
     )
 
     simple_photo_ion_rates = [0.03464792, 0.68099508]