Adaptive timestepping

src/festim/stepsize.py

@@ -1,3 +1,6 @@
+import numpy as np
+
+
 class Stepsize:
     """
     A class for evaluating the stepsize of transient simulations.
@@ -11,6 +14,10 @@ class Stepsize:
         target_nb_iterations (int, optional): number of Newton iterations
             over (resp. under) which the stepsize is increased
             (resp. decreased)
+        max_stepsize (float or callable, optional): Maximum stepsize.
+            Can be a function of festim.t. Defaults to None.
+        milestones (list, optional): list of times by which the simulation must
+            pass. Defaults to None.
 
 
     Attributes:
@@ -23,6 +30,9 @@ class Stepsize:
             over (resp. under) which the stepsize is increased
             (resp. decreased)
         adaptive (bool): True if the stepsize is adaptive, False otherwise.
+        max_stepsize (float): Maximum stepsize.
+        milestones (list): list of times by which the simulation must
+            pass.
     """
 
     def __init__(
@@ -31,14 +41,29 @@ def __init__(
         growth_factor=None,
         cutback_factor=None,
         target_nb_iterations=None,
+        max_stepsize=None,
+        milestones=None,
     ) -> None:
         self.initial_value = initial_value
         self.growth_factor = growth_factor
         self.cutback_factor = cutback_factor
         self.target_nb_iterations = target_nb_iterations
+        self.max_stepsize = max_stepsize
+        self.milestones = milestones
 
         # TODO should this class hold the dt object used in the formulation
 
+    @property
+    def milestones(self):
+        return self._milestones
+
+    @milestones.setter
+    def milestones(self, value):
+        if value:
+            self._milestones = sorted(value)
+        else:
+            self._milestones = value
+
     @property
     def adaptive(self):
         return self.growth_factor or self.cutback_factor or self.target_nb_iterations
@@ -67,16 +92,59 @@ def cutback_factor(self, value):
 
         self._cutback_factor = value
 
+    @property
+    def max_stepsize(self):
+        return self._max_stepsize
+
+    @max_stepsize.setter
+    def max_stepsize(self, value):
+        if value is not None and not callable(value):
+            if value < self.initial_value:
+                raise ValueError(
+                    "maximum stepsize cannot be less than initial stepsize"
+                )
+
+        self._max_stepsize = value
+
+    def get_max_stepsize(self, t):
+        """
+        Returns the maximum stepsize at time t.
+
+        Args:
+            t (float): the current time
+
+        Returns:
+            float or None: the maximum stepsize at time t
+        """
+        if callable(self._max_stepsize):
+            return self._max_stepsize(t)
+        return self._max_stepsize
+
     def modify_value(self, value, nb_iterations, t=None):
         if not self.is_adapt(t):
             return value
 
+        max_step = self.get_max_stepsize(t)
+
         if nb_iterations < self.target_nb_iterations:
-            return value * self.growth_factor
+            updated_value = min(
+                value * self.growth_factor,
+                max_step if max_step is not None else value * self.growth_factor,
+            )
         elif nb_iterations > self.target_nb_iterations:
-            return value * self.cutback_factor
+            updated_value = value * self.cutback_factor
         else:
-            return value
+            updated_value = value
+
+        next_milestone = self.next_milestone(t)
+        if next_milestone is not None:
+            time_to_milestone = next_milestone - t
+            if updated_value > time_to_milestone and not np.isclose(
+                t, next_milestone, atol=0
+            ):
+                updated_value = time_to_milestone
+
+        return updated_value
 
     def is_adapt(self, t):
         """
@@ -90,3 +158,20 @@ def is_adapt(self, t):
             bool: True if needs to adapt, False otherwise.
         """
         return True
+
+    def next_milestone(self, current_time: float):
+        """Returns the next milestone that the simulation must pass.
+        Returns None if there are no more milestones.
+
+        Args:
+            current_time (float): current time.
+
+        Returns:
+            float: next milestone.
+        """
+        if self.milestones is None:
+            return None
+        for milestone in self.milestones:
+            if current_time < milestone:
+                return milestone
+        return None

test/test_stepsize.py

@@ -47,6 +47,31 @@ def test_adaptive_stepsize_shrinks(cutback_factor, target):
     expected_value = current_value * my_stepsize.cutback_factor
     assert np.isclose(new_value, expected_value)
 
+@pytest.mark.parametrize("max_stepsize, growth_factor, target", [(4,3,1)])
+def test_max_stepsize(max_stepsize, growth_factor, target):
+    """Checks that the stepsize is capped at max
+    stepsize.
+
+    Args:
+        max_stepsize (float): maximum stepsize
+        growth_factor (float): the growth factor
+        target (int): the target number of iterations
+    """
+
+    my_stepsize = F.Stepsize(initial_value=2)
+    my_stepsize.max_stepsize = max_stepsize
+    my_stepsize.growth_factor = growth_factor
+    my_stepsize.target_nb_iterations = target
+
+    current_value = 2
+    new_value = my_stepsize.modify_value(
+        value=current_value,
+        nb_iterations=my_stepsize.target_nb_iterations - 1,
+    )
+
+    expected_value = max_stepsize
+    assert new_value == expected_value
+
 
 def test_stepsize_is_unchanged():
     """
@@ -100,3 +125,39 @@ def test_cutback_factor_setter():
     # Test that setting cutback factor to None works
     stepsize.cutback_factor = None
     assert stepsize.cutback_factor is None
+
+def test_max_stepsize_setter():
+    """Checks that the maximum stepsize setter works correctly"""
+    stepsize = F.Stepsize(initial_value=1)
+
+    # Test that setting a maximum stepsize less than initial stepsize raises a ValueError
+    with pytest.raises(ValueError, match="maximum stepsize cannot be less than initial stepsize"):
+        stepsize.max_stepsize = 0.5
+
+    # Test that setting maximum stepsize to None works
+    stepsize.max_stepsize = None
+    assert stepsize.max_stepsize is None
+
+@pytest.mark.parametrize("milestones, current_time, expected_value", [([1.0,25.4], 20.1, 25.4), ([9.8], 10.0, None)])
+def test_next_milestone(milestones, current_time, expected_value):
+     """Checks that the next milestone is 
+     identified and set correctly.
+
+    Args:
+        milestone (float): next milestone
+        current_time (float): current time in simulation
+    """
+     stepsize = F.Stepsize(initial_value=0.5)
+
+     stepsize.milestones = milestones
+
+     next_milestone = stepsize.next_milestone(current_time=current_time)
+     assert expected_value == next_milestone
+
+def test_no_milestones():
+     """Checks that no milestones works correctly"""
+     stepsize = F.Stepsize(initial_value=0.5)
+
+    # Test that setting milestones to None works 
+     stepsize.milestones = None
+     assert stepsize.milestones is None