Fixing typos for nlte ion (#2154)

* fixing typos on rate_matrix * changed the np arange to use 0 instead of 0.0
tardis-sn · Nov 24, 2022 · e26797d · e26797d
1 parent 56c050d
commit e26797d
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 15 deletions.
diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py
@@ -50,18 +50,18 @@ def calculate(
         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",
+            (atomic_number, ion_number, treatment type)
+            If ion is treated in LTE or nebular 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
+        ion_number_densities_nlte : DataFrame
             Number density with NLTE ionization treatment.
-        electron_densities_nlte: Series
-            Electron density with NLTE ionizaion treatment.
+        electron_densities_nlte : Series
+            Electron density with NLTE ionization treatment.
         """
 
         (
@@ -120,9 +120,9 @@ def calculate_rate_matrix(
         total_photo_ion_coefficients : DataFrame
             Photo ionization coefficients
         total_rad_recomb_coefficients : DataFrame
-            Radiative recombination coefficients(should get multiplied by electron density)
+            Radiative recombination coefficients (should get multiplied by electron density)
         total_coll_ion_coefficients : DataFrame
-            Collisional 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)
 
@@ -147,7 +147,7 @@ def calculate_rate_matrix(
             rate_matrix_index.get_level_values("atomic_number")
             .unique()
             .drop("n_e")
-        )  # dropping the n_e index
+        )  # dropping the n_e index, because 1st index of rate_matrix_index is (atomic_numbers, "n_e").
         for atomic_number in atomic_numbers:
             ion_numbers = rate_matrix.loc[atomic_number].index.get_level_values(
                 "ion_number"
@@ -331,7 +331,7 @@ def prepare_charge_conservation_row(atomic_numbers):
         density equation."""
         charge_conservation_row = []
         for atomic_number in atomic_numbers:
-            charge_conservation_row.append(np.arange(0.0, atomic_number + 1))
+            charge_conservation_row.append(np.arange(0, atomic_number + 1))
         charge_conservation_row = np.hstack([*charge_conservation_row, -1])
         # TODO needs to be modified for use in nlte_excitation
         return charge_conservation_row
@@ -373,13 +373,13 @@ def prepare_ion_recomb_coefficients_nlte_ion(
             General Boltzmann factor.
         Returns
         -------
-        total_photo_ion_coefficients:
-            Photoinization coefficients grouped by atomic number and ion number.
-        total_rad_recomb_coefficients:
+        total_photo_ion_coefficients
+            Photoionization 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:
+        total_coll_ion_coefficients
             Collisional ionization coefficients grouped by atomic number and ion number.
-        total_coll_recomb_coefficients:
+        total_coll_recomb_coefficients
             Collisional recombination coefficients grouped by atomic number and ion number.
         """
         indexer = pd.Series(

diff --git a/tardis/plasma/properties/rate_matrix_index.py b/tardis/plasma/properties/rate_matrix_index.py
@@ -15,7 +15,7 @@ def __init__(self, plasma_parent, nlte_ionization_species=0):
 
     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
+        "lte_ion" if that ion_number is treated in LTE or nebular, "nlte_ion" for NLTE ionization and
         keeps the levels for the rest.
 
         Parameters