fix: support adaptive time step retries

autotest/test_interface.py

@@ -319,26 +319,76 @@ def callback(sim, step):
         raise Exception(e)
 
 
-def test_ats_model(function_tmpdir):
+def test_ats_timestep_shrink(function_tmpdir):
+    """
+    Check that ATS time step length is reduced, reflecting dtadj,
+    between successful timesteps. Unrelated to retry-on-failure.
+    """
+
     def callback(sim, step):
         if step == Callbacks.stress_period_start:
             if sim.kper == 0 and sim.kstp == 0:
                 delt0 = sim.delt
 
         if step == Callbacks.timestep_start:
             if sim.kstp == 1:
-                if delt0 == sim.delt:
-                    raise AssertionError("ATS routines not reducing timestep length")
-
-        name = "ats0"
-        sim_pth = data_pth / name
-        test_pth = function_tmpdir / name
-        shutil.copytree(sim_pth, test_pth, dirs_exist_ok=True)
-
-        try:
-            run_simulation(so, test_pth, callback)
-        except Exception as e:
-            raise Exception(e)
+                assert delt0 > sim.delt, "timestep length wasn't reduced"
+
+    name = "ats0"
+    sim_pth = data_pth / name
+    test_pth = function_tmpdir / name
+    shutil.copytree(sim_pth, test_pth, dirs_exist_ok=True)
+
+    run_simulation(so, test_pth, callback)
+
+
+def test_ats_timestep_retry(function_tmpdir, monkeypatch):
+    """
+    Exercises run_simulation()'s ATS retry-loop control flow without
+    needing a MODFLOW 6 library build that actually implements retry.
+    """
+    # dict, not a bool, so fake_finish_retry can flip it without `nonlocal`
+    retried = {"done": False}
+    steps = []
+
+    def fake_finish_retry(_):
+        if retried["done"]:
+            return True
+        retried["done"] = True
+        return False
+
+    monkeypatch.setattr(ModflowApi, "has_ats_retry", True)
+    monkeypatch.setattr(ModflowApi, "prepare_retryloop", lambda self: None)
+    monkeypatch.setattr(ModflowApi, "start_retry", lambda self: None)
+    monkeypatch.setattr(ModflowApi, "finish_retry", fake_finish_retry)
+
+    def callback(_, step):
+        steps.append(step)
+
+    name = "dis_model"
+    sim_pth = data_pth / name
+    test_pth = function_tmpdir / name
+    shutil.copytree(sim_pth, test_pth, dirs_exist_ok=True)
+
+    run_simulation(so, test_pth, callback)
+
+    # retried once
+    assert retried["done"]
+    assert steps.count(Callbacks.timestep_retry) == 1
+
+    # sanity check multiple timesteps were solved
+    n_start = steps.count(Callbacks.timestep_start)
+    n_end = steps.count(Callbacks.timestep_end)
+    assert n_start == n_end
+    assert n_start > 1
+
+    # iteration_end should fire immediately before timestep_retry
+    retry_idx = steps.index(Callbacks.timestep_retry)
+    assert steps[retry_idx - 1] == Callbacks.iteration_end
+
+    # before retry, there should have been one full solve attempt
+    assert steps[:retry_idx].count(Callbacks.timestep_start) == 1
+    assert steps[:retry_idx].count(Callbacks.timestep_end) == 0
 
 
 def test_rhs_hcof_advanced(function_tmpdir):

docs/examples/notebooks/Extensions.py

@@ -315,6 +315,12 @@
 #       the user to make adjustments to stress packages and check values of stress packages at the end of each
 #       outer solution iteration.
 #    - `Callbacks.finalize`: the finalize callback is useful for finalizing models coupled with the modflowapi.
+#    - `Callbacks.timestep_retry`: the timestep_retry callback fires when an adaptive time stepping (ATS)
+#       timestep fails to converge and MODFLOW 6 is about to retry it with a reduced timestep length. This
+#       requires a MODFLOW 6 library built with ATS retry support (`modflowapi.ModflowApi.has_ats_retry`);
+#       on older MODFLOW 6 builds, a failed timestep is not retried and this callback never fires. By the time
+#       this callback fires, MODFLOW 6 has already reduced `delt` for the upcoming retry, so it can be read
+#       from the simulation object passed to the callback.
 #
 # The user can use any or all of these callbacks within their callback function
 

modflowapi/extensions/runner.py

@@ -1,3 +1,4 @@
+import warnings
 from enum import Enum
 
 from .. import ModflowApi
@@ -13,6 +14,7 @@ class Callbacks(Enum):
     iteration_start = 5
     iteration_end = 6
     finalize = 7
+    timestep_retry = 8
 
 
 def run_simulation(dll, sim_path, callback, verbose=False, _develop=False):
@@ -47,6 +49,14 @@ def run_simulation(dll, sim_path, callback, verbose=False, _develop=False):
     mf6.initialize()
     sim = ApiSimulation.load(mf6)
 
+    if sim.ats_active and not mf6.has_ats_retry:
+        warnings.warn(
+            "ATS is active but this MODFLOW 6 library does not support it."
+            "Non-converging timesteps will fail instead of being retried."
+            "Install a newer MODFLOW 6 library to enable ATS support.",
+            stacklevel=2,
+        )
+
     if _develop:
         with open("var_list.txt", "w") as foo:
             for name in mf6.get_input_var_names():
@@ -76,28 +86,25 @@ def run_simulation(dll, sim_path, callback, verbose=False, _develop=False):
                     models[model.name.lower()] = model
 
             sim_grp = ApiSimulation(mf6, models, solution, sim._exchanges, sim.tdis, sim.ats)
-            mf6.prepare_solve(sol_id)
             if sim.kper != kperold[sol_id - 1]:
                 callback(sim_grp, Callbacks.stress_period_start)
                 kperold[sol_id - 1] += 1
             elif current_time == 0:
                 callback(sim_grp, Callbacks.stress_period_start)
 
-            kiter = 0
             callback(sim_grp, Callbacks.timestep_start)
 
-            if sim_grp.ats_period[0]:
-                mindt = sim_grp.ats_period[-1]
-                while sim_grp.delt > mindt and kiter < maxiter:
-                    sim_grp.iteration = kiter
-                    callback(sim_grp, Callbacks.iteration_start)
-                    has_converged = mf6.solve(sol_id)
-                    callback(sim_grp, Callbacks.iteration_end)
-                    kiter += 1
-                    if has_converged and sim_grp.allow_convergence:
-                        break
+            supports_retry = mf6.has_ats_retry
+            if supports_retry:
+                mf6.prepare_retryloop()
 
-            else:
+            finished = False
+            while not finished:
+                if supports_retry:
+                    mf6.start_retry()
+
+                mf6.prepare_solve(sol_id)
+                kiter = 0
                 while kiter < maxiter:
                     sim_grp.iteration = kiter
                     callback(sim_grp, Callbacks.iteration_start)
@@ -106,9 +113,13 @@ def run_simulation(dll, sim_path, callback, verbose=False, _develop=False):
                     kiter += 1
                     if has_converged and sim_grp.allow_convergence:
                         break
+                mf6.finalize_solve(sol_id)
+
+                finished = mf6.finish_retry() if supports_retry else True
+                if not finished:
+                    callback(sim_grp, Callbacks.timestep_retry)
 
             callback(sim_grp, Callbacks.timestep_end)
-            mf6.finalize_solve(sol_id)
             if sim_grp.nstp == sim_grp.kstp + 1:
                 callback(sim_grp, Callbacks.stress_period_end)
 

modflowapi/modflowapi.py

@@ -1,6 +1,8 @@
+from ctypes import byref, c_bool
 from os import PathLike
 
 from xmipy import XmiWrapper
+from xmipy.utils import cd
 
 from .util import amend_libmf6_path
 
@@ -31,3 +33,41 @@ def __init__(
             timing=timing,
             logger_level=logger_level,
         )
+        self._has_ats_retry: bool | None = None
+
+    @property
+    def has_ats_retry(self) -> bool:
+        """
+        Whether the MODFLOW 6 library exposes adaptive time stepping
+        retry routines (prepare_retryloop, start_retry, finish_retry).
+        """
+        if self._has_ats_retry is None:
+            self._has_ats_retry = hasattr(self.lib, "prepare_retryloop")
+        return self._has_ats_retry
+
+    def prepare_retryloop(self) -> None:
+        """
+        Reset the ATS retry counter before running a timestep. Must be
+        called once per timestep, before the first start_retry() call.
+        """
+        with cd(self.working_directory):
+            self._execute_function(self.lib.prepare_retryloop)
+
+    def start_retry(self) -> None:
+        """
+        Signal the start of a (re)try of the current timestep. If not
+        the first attempt, timeseries input will be reevaluated. Must be
+        called before prepare_solve() on all attempts, including the first.
+        """
+        with cd(self.working_directory):
+            self._execute_function(self.lib.start_retry)
+
+    def finish_retry(self) -> bool:
+        """
+        Check whether the current timestep is finished, or must be retried
+        with a smaller delt. Must be called after finalize_solve().
+        """
+        finished = c_bool(False)
+        with cd(self.working_directory):
+            self._execute_function(self.lib.finish_retry, byref(finished))
+        return finished.value