diff --git a/autotest/test_interface.py b/autotest/test_interface.py index 79a0fcc..f5e8e8f 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, 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: @@ -327,18 +332,63 @@ 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") - - 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): 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..64e46ce 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,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) diff --git a/modflowapi/modflowapi.py b/modflowapi/modflowapi.py index 11b8d9b..da9a8b7 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 + 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