NLTE Ionization solver polish

tardis/io/decay.py

@@ -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):
@@ -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):

tardis/plasma/properties/nlte_rate_equation_solver.py

@@ -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,
@@ -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.