Skip to content

Commit

Permalink
NLTE Ionization solver polish (#2461)
Browse files Browse the repository at this point in the history
* assertion for first guess

* Add checks to decay

* Fix greater equal

* Fix deprecated series index
  • Loading branch information
AlexHls authored Nov 2, 2023
1 parent dce9041 commit 98c19ee
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 2 deletions.
13 changes: 12 additions & 1 deletion tardis/io/decay.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,12 @@
from radioactivedecay import Nuclide, Inventory
from radioactivedecay.utils import Z_to_elem
from astropy import units as u
import logging

logger = logging.getLogger(__name__)

class IsotopeAbundances(pd.DataFrame):

class IsotopeAbundances(pd.DataFrame):
_metadata = ["time_0"]

def __init__(self, *args, **kwargs):
Expand Down Expand Up @@ -92,9 +94,18 @@ def decay(self, t):
t_second = (
u.Quantity(t, u.day).to(u.s).value - self.time_0.to(u.s).value
)
logger.info(f"Decaying abundances for {t_second} seconds")
if t_second < 0:
logger.warning(
f"Decay time {t_second} is negative. This could indicate a miss-specified input model."
f" A negative decay time can potentially lead to negative abundances."
)
decayed_inventories = [item.decay(t_second) for item in inventories]
df = IsotopeAbundances.from_inventories(decayed_inventories)
df.sort_index(inplace=True)
assert (
df.ge(0.0).all().all()
), "Negative abundances detected. Please make sure your input abundances are correct."
return df

def as_atoms(self):
Expand Down
6 changes: 5 additions & 1 deletion tardis/plasma/properties/nlte_rate_equation_solver.py
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,10 @@ def calculate(
number_density[shell],
initial_electron_densities[shell],
)
# All first guess values have to be positive
assert (
np.greater_equal(first_guess, 0.0).all()
).all(), "First guess for NLTE solver has negative values, something went wrong."
solution = root(
self.population_objective_function,
first_guess,
Expand Down Expand Up @@ -591,7 +595,7 @@ def prepare_first_guess(
"""
first_guess = pd.Series(0.0, index=rate_matrix_index)
for atomic_number in atomic_numbers:
first_guess.loc[(atomic_number, 1)][0] = number_density.loc[
first_guess.loc[(atomic_number, 1)].iloc[0] = number_density.loc[
atomic_number
]
# TODO: After the first iteration, the new guess can be the old solution.
Expand Down

0 comments on commit 98c19ee

Please sign in to comment.