Add tolerance to not_done condition of time step calculator.

torax/time_step_calculator/pydantic_model.py

@@ -30,14 +30,25 @@ class TimeStepCalculatorType(enum.Enum):
 
 
 class TimeStepCalculator(torax_pydantic.BaseModelFrozen):
-  """Config for a time step calculator."""
+  """Config for a time step calculator.
+
+  Attributes:
+    calculator_type: The type of time step calculator to use.
+    tolerance: The tolerance within the final time for which the simulation
+      will be considered done.
+  """
 
   calculator_type: TimeStepCalculatorType = TimeStepCalculatorType.CHI
+  tolerance: float = 1e-7
 
   @property
   def time_step_calculator(self) -> time_step_calculator.TimeStepCalculator:
     match self.calculator_type:
       case TimeStepCalculatorType.CHI:
-        return chi_time_step_calculator.ChiTimeStepCalculator()
+        return chi_time_step_calculator.ChiTimeStepCalculator(
+            tolerance=self.tolerance
+        )
       case TimeStepCalculatorType.FIXED:
-        return fixed_time_step_calculator.FixedTimeStepCalculator()
+        return fixed_time_step_calculator.FixedTimeStepCalculator(
+            tolerance=self.tolerance
+        )

torax/time_step_calculator/time_step_calculator.py

@@ -18,15 +18,14 @@
 """
 
 import abc
-from typing import Protocol, Union
 
 import jax
 from torax import state as state_module
 from torax.config import runtime_params_slice
 from torax.geometry import geometry
 
 
-class TimeStepCalculator(Protocol):
+class TimeStepCalculator(abc.ABC):
   """Iterates over time during simulation.
 
   Usage follows this pattern:
@@ -42,12 +41,15 @@ class TimeStepCalculator(Protocol):
       sim_state = <update sim_state with step of size dt>
   """
 
+  def __init__(self, tolerance: float = 1e-7):
+    self.tolerance = tolerance
+
   def not_done(
       self,
-      t: Union[float, jax.Array],
+      t: float | jax.Array,
       t_final: float,
-  ) -> Union[bool, jax.Array]:
-    return t < t_final
+  ) -> bool | jax.Array:
+    return t < (t_final - self.tolerance)
 
   @abc.abstractmethod
   def next_dt(