diff --git a/tardis/plasma/properties/nlte_rate_equation_solver.py b/tardis/plasma/properties/nlte_rate_equation_solver.py index 2ef015b32ed..d37c4464679 100644 --- a/tardis/plasma/properties/nlte_rate_equation_solver.py +++ b/tardis/plasma/properties/nlte_rate_equation_solver.py @@ -18,7 +18,7 @@ NLTE_POPULATION_NEGATIVE_POPULATION_TOLERANCE = ( -1e-10 ) # Maximum negative population allowed before solver fails -NLTE_POPULATION_SOLVER_CHARGE_CONSERVATION_TOLERANCE = 1e-6 # Arbitrary tolerance for charge conservation, should be changed to a more reasonable value +NLTE_POPULATION_SOLVER_CHARGE_CONSERVATION_TOLERANCE = 1e-6 # Arbitrary tolerance for charge conservation, should be changed to a more reasonable value class NLTEPopulationSolverRoot(ProcessingPlasmaProperty): @@ -150,18 +150,18 @@ def calculate( solution.x[-1], number_density[shell], ion_numbers, - ) + ) # Check that the electron density is still in line with charge conservation # after removing negative populations assert ( np.abs( np.sum(ion_numbers * ion_number_density[shell]) - electron_densities[shell] - ) / electron_densities[shell] + ) + / electron_densities[shell] < NLTE_POPULATION_SOLVER_CHARGE_CONSERVATION_TOLERANCE ), "Charge conservation not fulfilled after correcting for negative populations, solver failed." - # TODO: change the jacobian and rate matrix to use shell id and get coefficients from the attribute of the class. return ion_number_density, electron_densities diff --git a/tardis/plasma/tests/test_nlte_solver.py b/tardis/plasma/tests/test_nlte_solver.py index 8a20187c529..6379e7a3e0a 100644 --- a/tardis/plasma/tests/test_nlte_solver.py +++ b/tardis/plasma/tests/test_nlte_solver.py @@ -3,7 +3,10 @@ import pytest from numpy.testing import assert_allclose -from tardis.plasma.properties import NLTEPopulationSolverRoot, NLTEPopulationSolverLU +from tardis.plasma.properties import ( + NLTEPopulationSolverRoot, + NLTEPopulationSolverLU, +) from tardis.plasma.properties.nlte_rate_equation_solver import ( calculate_jacobian_matrix, calculate_rate_matrix,