From 1e57e01415a9e1a4b06af741283f1704068bc8a0 Mon Sep 17 00:00:00 2001
From: kaelyndunnell <kaelyndunnell03@gmail.com>
Date: Tue, 5 Nov 2024 16:36:58 -0500
Subject: [PATCH] fix max stepsize check in modify_value & add test (thanks
 @RemDelaporteMathurin)

---
 src/festim/stepsize.py | 12 ++++++------
 test/test_stepsize.py  | 19 +++++++++++++++++++
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/src/festim/stepsize.py b/src/festim/stepsize.py
index 28840efb0..ccafc827f 100644
--- a/src/festim/stepsize.py
+++ b/src/festim/stepsize.py
@@ -124,18 +124,18 @@ 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:
-            updated_value = min(
-                value * self.growth_factor,
-                max_step if max_step is not None else value * self.growth_factor,
-            )
+            updated_value = value * self.growth_factor
         elif nb_iterations > self.target_nb_iterations:
             updated_value = value * self.cutback_factor
         else:
             updated_value = value
 
+        max_step = self.get_max_stepsize(t)
+        if max_step:
+            if updated_value > max_step:
+                updated_value = max_step
+
         next_milestone = self.next_milestone(t)
         if next_milestone is not None:
             time_to_milestone = next_milestone - t
diff --git a/test/test_stepsize.py b/test/test_stepsize.py
index 5baef2dde..288bb04f7 100644
--- a/test/test_stepsize.py
+++ b/test/test_stepsize.py
@@ -170,3 +170,22 @@ def test_no_milestones():
     # Test that setting milestones to None works
     stepsize.milestones = None
     assert stepsize.milestones is None
+
+
+def test_modify_for_stepsize():
+    """Tests that modify_value returns max_stepsize
+    when max_stepsize is less than next stepsize.
+    """
+    my_stepsize = F.Stepsize(initial_value=1)
+    my_stepsize.max_stepsize = 1
+    my_stepsize.growth_factor = 1.1
+    my_stepsize.target_nb_iterations = 4
+
+    current_value = 2
+    new_value = my_stepsize.modify_value(
+        value=current_value,
+        nb_iterations=my_stepsize.target_nb_iterations,
+    )
+
+    expected_value = 1
+    assert new_value == expected_value