From 121683a5dcb3ba60ea4dd552fae92effa90fffcd Mon Sep 17 00:00:00 2001 From: Sona Chitchyan Date: Wed, 9 Nov 2022 11:30:24 -0500 Subject: [PATCH] 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]