Adds docstrings

tardis/plasma/equilibrium/level_populations.py

@@ -2,7 +2,19 @@
 import pandas as pd
 
 
-def expand_series(series):
+def expand_series(series: pd.Series):
+    """Helper function to expand a series of numpy arrays into a DataFrame.
+
+    Parameters
+    ----------
+    series : pd.Series
+        A series of numpy arrays.
+
+    Returns
+    -------
+    pd.DataFrame
+        The expanded DataFrame where each row is an entry from the numpy array.
+    """
     # Expand the numpy arrays into rows
     expanded_series = pd.DataFrame(
         np.concatenate(series.values).reshape(-1, len(series))
@@ -11,11 +23,37 @@ def expand_series(series):
 
 
 class LevelPopulationSolver:
-    def __init__(self, rates_matrices: pd.DataFrame, levels):
+    def __init__(self, rates_matrices: pd.DataFrame, levels: pd.DataFrame):
+        """Solve the normalized level population values from the rate matrices.
+
+        Parameters
+        ----------
+        rates_matrices : pd.DataFrame
+            DataFrame of rate matrices indexed by atomic number and ion number,
+            with each column being a cell.
+        levels : pd.DataFrame
+            DataFrame of energy levels.
+        """
         self.rates_matrices = rates_matrices
         self.levels = levels
 
-    def __calculate_boltzmann_factor(self, rates_matrix, species_id):
+    def __calculate_boltzmann_factor(
+        self, rates_matrix: np.ndarray, species_id: tuple
+    ):
+        """Helper function to calculate the normalized, per-level boltzmann factor.
+
+        Parameters
+        ----------
+        rates_matrix : np.ndarray
+            The rate matrix for a given species and cell.
+        species_id : tuple
+            A tuple of the atomic number and ion number.
+
+        Returns
+        -------
+        np.ndarray
+            The normalized, per-level boltzmann factor.
+        """
         x = np.zeros(rates_matrix.shape[0])
         x[0] = 1.0
         level_boltzmann_factor = np.linalg.solve(rates_matrix[:, :], x)
@@ -27,13 +65,23 @@ def __calculate_boltzmann_factor(self, rates_matrix, species_id):
         return normalized_level_boltzmann_factor
 
     def solve(self):
+        """Solves the normalized level population values from the rate matrices.
+
+        Returns
+        -------
+        pd.DataFrame
+            Normalized level population values indexed by atomic number, ion
+            number and level number. Columns are cells.
+        """
         normalized_level_boltzmann_factors = pd.DataFrame(
             index=self.levels.index,
             columns=self.rates_matrices.columns,
             dtype=np.float64,
         )
 
         for species_id in self.rates_matrices.index:
+            # TODO: resolve the chained assignment here. Maybe an intermediate df
+            # is needed
             normalized_level_boltzmann_factors.loc[species_id, :].update(
                 expand_series(
                     self.rates_matrices.loc[species_id].apply(

tardis/plasma/equilibrium/rate_matrix.py

@@ -10,11 +10,31 @@ def __init__(
         levels: pd.DataFrame,
         electron_number_densities,
     ):
+        """Constructs the rate matrix from an arbitrary number of rate solvers.
+
+        Parameters
+        ----------
+        rate_solvers : list
+            List of rate solver objects.
+        levels : pd.DataFrame
+            DataFrame of energy levels.
+        electron_number_densities : float or array of floats
+            Electron number densities in g/cm^3. If an array, it should have the
+            same length as the number of cells in the simulation.
+        """
         self.rate_solvers = rate_solvers
         self.levels = levels
         self.electron_number_densities = electron_number_densities
 
     def solve(self):
+        """Construct the compiled rate matrix dataframe.
+
+        Returns
+        -------
+        pd.DataFrame
+            A DataFrame of rate matrices indexed by atomic number and ion number,
+            with each column being a cell.
+        """
         rates_df_list = [
             solver.solve(input) for solver, input in self.rate_solvers
         ]