ran black

tardis-sn · Nov 9, 2022 · 121683a · 121683a
1 parent 20ca228
commit 121683a
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 9 deletions.
diff --git 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
@@ -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
@@ -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]