From 2973f7f73af73b3b14bd797d23d26bb62f5f8cb9 Mon Sep 17 00:00:00 2001 From: wpbonelli Date: Sat, 11 Jul 2026 08:24:58 -0700 Subject: [PATCH 1/4] fix: support adaptive time stepping --- autotest/test_interface.py | 45 +++++++++++++++++++++++++++ docs/examples/notebooks/Extensions.py | 6 ++++ modflowapi/extensions/runner.py | 41 +++++++++++++++--------- modflowapi/modflowapi.py | 40 ++++++++++++++++++++++++ 4 files changed, 118 insertions(+), 14 deletions(-) diff --git a/autotest/test_interface.py b/autotest/test_interface.py index 79a0fcc..f7047ac 100644 --- a/autotest/test_interface.py +++ b/autotest/test_interface.py @@ -341,6 +341,51 @@ def callback(sim, step): raise Exception(e) +def test_timestep_retry_callback(function_tmpdir, monkeypatch): + """ + Exercises run_simulation()'s ATS retry-loop control flow without + needing a MODFLOW 6 library build that actually implements retry. + """ + monkeypatch.setattr(ModflowApi, "has_ats_retry", True) + monkeypatch.setattr(ModflowApi, "prepare_retryloop", lambda self: None) + monkeypatch.setattr(ModflowApi, "start_retry", lambda self: None) + + retried = {"done": False} + + def fake_finish_retry(self): + if not retried["done"]: + retried["done"] = True + return False + return True + + monkeypatch.setattr(ModflowApi, "finish_retry", fake_finish_retry) + + steps = [] + + def callback(sim, 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) + + assert retried["done"], "finish_retry was never called; retry loop not exercised" + assert steps.count(Callbacks.timestep_retry) == 1 + + n_start = steps.count(Callbacks.timestep_start) + n_end = steps.count(Callbacks.timestep_end) + assert n_start == n_end + assert n_start > 1 # more than one outer timestep ran + + retry_idx = steps.index(Callbacks.timestep_retry) + assert steps[retry_idx - 1] == Callbacks.iteration_end + 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): def callback(sim, step): model = sim.test_model diff --git a/docs/examples/notebooks/Extensions.py b/docs/examples/notebooks/Extensions.py index 6f97fc2..16d300e 100644 --- a/docs/examples/notebooks/Extensions.py +++ b/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 diff --git a/modflowapi/extensions/runner.py b/modflowapi/extensions/runner.py index 33191a9..285aa6f 100644 --- a/modflowapi/extensions/runner.py +++ b/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,15 @@ 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: + # MF6 has already reduced delt for the next attempt by + # this point, so sim_grp.delt reflects the retry value + 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) diff --git a/modflowapi/modflowapi.py b/modflowapi/modflowapi.py index 11b8d9b..33868cb 100644 --- a/modflowapi/modflowapi.py +++ b/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 + (ATS) 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 From f00d2d80cea989bb2eb0522eda053b8576547237 Mon Sep 17 00:00:00 2001 From: wpbonelli Date: Sat, 11 Jul 2026 10:40:26 -0700 Subject: [PATCH 2/4] better test naming/comments --- autotest/test_interface.py | 44 ++++++++++++++++++++++---------------- 1 file changed, 26 insertions(+), 18 deletions(-) diff --git a/autotest/test_interface.py b/autotest/test_interface.py index f7047ac..9ab8aad 100644 --- a/autotest/test_interface.py +++ b/autotest/test_interface.py @@ -319,7 +319,12 @@ 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 between successful time + steps. This is unrelated to ATS retry-on-failure. + """ + def callback(sim, step): if step == Callbacks.stress_period_start: if sim.kper == 0 and sim.kstp == 0: @@ -327,8 +332,7 @@ def callback(sim, step): if step == Callbacks.timestep_start: if sim.kstp == 1: - if delt0 == sim.delt: - raise AssertionError("ATS routines not reducing timestep length") + assert delt0 > sim.delt, "timestep length wasn't reduced" name = "ats0" sim_pth = data_pth / name @@ -341,28 +345,27 @@ def callback(sim, step): raise Exception(e) -def test_timestep_retry_callback(function_tmpdir, monkeypatch): +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. """ - monkeypatch.setattr(ModflowApi, "has_ats_retry", True) - monkeypatch.setattr(ModflowApi, "prepare_retryloop", lambda self: None) - monkeypatch.setattr(ModflowApi, "start_retry", lambda self: None) - + # dict, not a bool, so fake_finish_retry can flip it without `nonlocal` retried = {"done": False} + steps = [] - def fake_finish_retry(self): - if not retried["done"]: - retried["done"] = True - return False - return True + 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) - steps = [] - - def callback(sim, step): + def callback(_, step): steps.append(step) name = "dis_model" @@ -372,16 +375,21 @@ def callback(sim, step): run_simulation(so, test_pth, callback) - assert retried["done"], "finish_retry was never called; retry loop not exercised" + # 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 # more than one outer timestep ran + 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 From b8c500d8e19bf65cb55119a66cc34051b43e0734 Mon Sep 17 00:00:00 2001 From: wpbonelli Date: Sat, 11 Jul 2026 10:42:54 -0700 Subject: [PATCH 3/4] fix test indentation --- autotest/test_interface.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/autotest/test_interface.py b/autotest/test_interface.py index 9ab8aad..19bf50a 100644 --- a/autotest/test_interface.py +++ b/autotest/test_interface.py @@ -334,15 +334,12 @@ def callback(sim, step): if sim.kstp == 1: 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) - - try: - run_simulation(so, test_pth, callback) - except Exception as e: - raise Exception(e) + 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): From f0f59fec357b509880978090d825317ef93c2883 Mon Sep 17 00:00:00 2001 From: wpbonelli Date: Sat, 11 Jul 2026 10:52:01 -0700 Subject: [PATCH 4/4] cleanup --- autotest/test_interface.py | 4 ++-- modflowapi/extensions/runner.py | 2 -- modflowapi/modflowapi.py | 2 +- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/autotest/test_interface.py b/autotest/test_interface.py index 19bf50a..f5e8e8f 100644 --- a/autotest/test_interface.py +++ b/autotest/test_interface.py @@ -321,8 +321,8 @@ def callback(sim, step): def test_ats_timestep_shrink(function_tmpdir): """ - Check that ATS time step length is reduced between successful time - steps. This is unrelated to ATS retry-on-failure. + Check that ATS time step length is reduced, reflecting dtadj, + between successful timesteps. Unrelated to retry-on-failure. """ def callback(sim, step): diff --git a/modflowapi/extensions/runner.py b/modflowapi/extensions/runner.py index 285aa6f..64e46ce 100644 --- a/modflowapi/extensions/runner.py +++ b/modflowapi/extensions/runner.py @@ -117,8 +117,6 @@ def run_simulation(dll, sim_path, callback, verbose=False, _develop=False): finished = mf6.finish_retry() if supports_retry else True if not finished: - # MF6 has already reduced delt for the next attempt by - # this point, so sim_grp.delt reflects the retry value callback(sim_grp, Callbacks.timestep_retry) callback(sim_grp, Callbacks.timestep_end) diff --git a/modflowapi/modflowapi.py b/modflowapi/modflowapi.py index 33868cb..da9a8b7 100644 --- a/modflowapi/modflowapi.py +++ b/modflowapi/modflowapi.py @@ -39,7 +39,7 @@ def __init__( def has_ats_retry(self) -> bool: """ Whether the MODFLOW 6 library exposes adaptive time stepping - (ATS) routines (prepare_retryloop, start_retry, finish_retry). + retry routines (prepare_retryloop, start_retry, finish_retry). """ if self._has_ats_retry is None: self._has_ats_retry = hasattr(self.lib, "prepare_retryloop")