Reject adaptive time steps when nonlinear solves fail

opencmp/solvers/adaptive_transient_solvers/adaptive_three_step.py

@@ -73,9 +73,11 @@ def _re_assemble(self) -> None:
         self._assemble()
         self._update_preconditioners(self.preconditioner_long)
 
-    def _single_solve(self) -> None:
+    def _single_solve(self) -> bool:
         # Single solve for the full time step.
-        self.model.solve_single_step(self.a_long, self.L_long, self.preconditioner_long, self.gfu_long, 0)
+        if self.model.solve_single_step(self.a_long, self.L_long, self.preconditioner_long,
+                                        self.gfu_long, 0) is False:
+            return self._discard_diverged_solve()
 
         # Update the linearization terms back to their t^n values.
         # gfu_short needs to be solved with the same values for W as gfu_long since they have the same initial time
@@ -87,7 +89,9 @@ def _single_solve(self) -> None:
             self.L_short[i].Assemble()
 
         self._update_preconditioners(self.preconditioner_short)
-        self.model.solve_single_step(self.a_short, self.L_short, self.preconditioner_short, self.gfu_short, 1)
+        if self.model.solve_single_step(self.a_short, self.L_short, self.preconditioner_short,
+                                        self.gfu_short, 1) is False:
+            return self._discard_diverged_solve()
 
         # Update the model component values at t^n+1/2 now that they have been solved for.
         # The linearization terms do not need to be updated since they are now for t^n+1/2 as expected.
@@ -99,7 +103,18 @@ def _single_solve(self) -> None:
             self.L[i].Assemble()
 
         self._update_preconditioners(self.preconditioner)
-        self.model.solve_single_step(self.a, self.L, self.preconditioner, self.gfu, 0)
+        if self.model.solve_single_step(self.a, self.L, self.preconditioner,
+                                        self.gfu, 0) is False:
+            return self._discard_diverged_solve()
+
+        return True
+
+    def _discard_diverged_solve(self) -> bool:
+        # Drop the diverged iterate so the retry doesn't linearize about garbage.
+        for candidate in (self.gfu_long, self.gfu_short, self.gfu):
+            candidate.vec.data = self.gfu_0_list[-1].vec
+        self.model.update_linearization(self.gfu_0_list[-1])
+        return False
 
     def _calculate_local_error(self) -> Tuple[List[float], List[float], List[str]]:
         # Include any variables specified by the model as included in local error.

opencmp/solvers/adaptive_transient_solvers/adaptive_two_step.py

@@ -69,9 +69,23 @@ def _create_preconditioners(self) -> None:
     def _re_assemble(self) -> None:
         self._assemble()
 
-    def _single_solve(self) -> None:
-        self.model.solve_single_step(self.a_pred, self.L_pred, self.preconditioner_pred, self.gfu_pred)
-        self.model.solve_single_step(self.a_corr, self.L_corr, self.preconditioner_corr, self.gfu)
+    def _single_solve(self) -> bool:
+        if self.model.solve_single_step(self.a_pred, self.L_pred, self.preconditioner_pred,
+                                        self.gfu_pred) is False:
+            return self._discard_diverged_solve()
+
+        if self.model.solve_single_step(self.a_corr, self.L_corr, self.preconditioner_corr,
+                                        self.gfu) is False:
+            return self._discard_diverged_solve()
+
+        return True
+
+    def _discard_diverged_solve(self) -> bool:
+        # Drop the diverged iterate so the retry doesn't linearize about garbage.
+        self.gfu_pred.vec.data = self.gfu_0_list[0].vec
+        self.gfu.vec.data = self.gfu_0_list[0].vec
+        self.model.update_linearization(self.gfu_0_list[0])
+        return False
 
     def _calculate_local_error(self) -> Tuple[List[float], List[float], List[str]]:
         # Include any variables specified by the model as included in local error.

opencmp/solvers/adaptive_transient_solvers/base_adaptive_transient_RK.py

@@ -39,12 +39,17 @@ class BaseAdaptiveTransientRKSolver(TransientRKSolver, ABC):
     def __init__(self, model_class: Type[Model], config: ConfigParser) -> None:
         super().__init__(model_class, config)
 
-    def _update_time_step(self) -> Tuple[bool, float, float, str]:
+    def _update_time_step(self, nonlinear_failed: bool = False) -> Tuple[bool, float, float, str]:
         dt_min_allowed = self.dt_range[0]
         dt_max_allowed = self.dt_range[1]
 
-        # Get the local error and the norm of the solution gridfunction.
-        local_error, gfu_norm, comp_names = self._calculate_local_error()
+        if nonlinear_failed:
+            # A diverged solve has no meaningful error estimate (and may be NaN, which compares
+            # as acceptable). Force a rejection through the normal path.
+            local_error, gfu_norm, comp_names = [1e30], [1.0], ['nonlinear solve']
+        else:
+            # Get the local error and the norm of the solution gridfunction.
+            local_error, gfu_norm, comp_names = self._calculate_local_error()
 
         # Create a list of all of the relative errors
         local_error_relative = local_error.copy()

opencmp/solvers/adaptive_transient_solvers/base_adaptive_transient_multistep.py

@@ -38,12 +38,17 @@ class BaseAdaptiveTransientMultiStepSolver(TransientMultiStepSolver, ABC):
     def __init__(self, model_class: Type[Model], config: ConfigParser) -> None:
         super().__init__(model_class, config)
 
-    def _update_time_step(self) -> Tuple[bool, float, float, str]:
+    def _update_time_step(self, nonlinear_failed: bool = False) -> Tuple[bool, float, float, str]:
         dt_min_allowed = self.dt_range[0]
         dt_max_allowed = self.dt_range[1]
 
-        # Get the local error and the norm of the solution gridfunction.
-        local_error, gfu_norm, comp_names = self._calculate_local_error()
+        if nonlinear_failed:
+            # A diverged solve has no meaningful error estimate (and may be NaN, which compares
+            # as acceptable). Force a rejection through the normal path.
+            local_error, gfu_norm, comp_names = [1e30], [1.0], ['nonlinear solve']
+        else:
+            # Get the local error and the norm of the solution gridfunction.
+            local_error, gfu_norm, comp_names = self._calculate_local_error()
 
         # Create a list of all the relative errors
         local_error_relative = local_error.copy()

opencmp/solvers/base_solver.py

@@ -544,10 +544,13 @@ def _solve(self) -> None:
 
                 self._re_assemble()
 
-                self._single_solve()
-
-                # Calculate local error, accept/reject current result, and update timestep
-                accept_this_iteration, local_error_abs, local_error_rel, component = self._update_time_step()
+                # A model may explicitly report that its nonlinear solve diverged.
+                if self._single_solve() is False and self.adaptive:
+                    accept_this_iteration, local_error_abs, local_error_rel, component = \
+                        self._update_time_step(nonlinear_failed=True)
+                else:
+                    # Calculate local error, accept/reject current result, and update timestep
+                    accept_this_iteration, local_error_abs, local_error_rel, component = self._update_time_step()
 
                 # Log information about the current timestep
                 self._log_timestep(accept_this_iteration, local_error_abs, local_error_rel, component)

pytests/test_adaptive_nonlinear_failure.py

@@ -0,0 +1,125 @@
+"""Tests for how the adaptive solvers react to a diverged nonlinear solve."""
+
+from opencmp.solvers.adaptive_transient_solvers.adaptive_three_step import AdaptiveThreeStep
+from opencmp.solvers.adaptive_transient_solvers.adaptive_two_step import AdaptiveTwoStep
+
+
+class _Parameter:
+    def __init__(self, value):
+        self.value = value
+
+    def Get(self):
+        return self.value
+
+    def Set(self, value):
+        self.value = value
+
+
+class _Vector:
+    def __init__(self, value):
+        self.data = value
+
+
+class _GridFunction:
+    def __init__(self, value):
+        self.vec = _Vector(value)
+
+
+class _FailingModel:
+    """Reports divergence on every solve."""
+
+    def __init__(self):
+        self.solve_calls = 0
+        self.linearization = None
+
+    def solve_single_step(self, *args, **kwargs):
+        self.solve_calls += 1
+        return False
+
+    def update_linearization(self, gfu):
+        self.linearization = gfu
+
+    def update_model_variables(self, gfu, time_step=None):
+        pass
+
+
+def _two_step():
+    solver = AdaptiveTwoStep.__new__(AdaptiveTwoStep)
+    solver.model = _FailingModel()
+    solver.a_pred = solver.L_pred = solver.preconditioner_pred = []
+    solver.a_corr = solver.L_corr = solver.preconditioner_corr = []
+    solver.gfu_0_list = [_GridFunction('accepted')]
+    solver.gfu_pred = _GridFunction('failed predictor')
+    solver.gfu = _GridFunction('failed corrector')
+    return solver
+
+
+def _three_step():
+    solver = AdaptiveThreeStep.__new__(AdaptiveThreeStep)
+    solver.model = _FailingModel()
+    solver.a_long = solver.L_long = solver.preconditioner_long = []
+    solver.scheme = 'adaptive three step'
+    solver.scheme_order = 2
+    solver.scheme_dt_coef = [1.0, 0.5]
+    solver.gfu_0_list = [_GridFunction('intermediate'), _GridFunction('accepted')]
+    solver.gfu_long = _GridFunction('failed long')
+    solver.gfu_short = _GridFunction('failed short')
+    solver.gfu = _GridFunction('failed final')
+    return solver
+
+
+def test_two_step_stops_and_discards_after_first_failure() -> None:
+    solver = _two_step()
+    accepted = solver.gfu_0_list[0]
+
+    assert solver._single_solve() is False
+    assert solver.model.solve_calls == 1
+    assert solver.gfu_pred.vec.data is accepted.vec
+    assert solver.gfu.vec.data is accepted.vec
+    assert solver.model.linearization is accepted
+
+
+def test_three_step_stops_and_discards_after_first_failure() -> None:
+    solver = _three_step()
+    accepted = solver.gfu_0_list[-1]
+
+    assert solver._single_solve() is False
+    assert solver.model.solve_calls == 1
+    assert all(candidate.vec.data is accepted.vec for candidate in
+               (solver.gfu_long, solver.gfu_short, solver.gfu))
+    assert solver.model.linearization is accepted
+
+
+def test_two_step_nonlinear_failure_rejects_and_shrinks_dt() -> None:
+    solver = _two_step()
+    solver.dt_range = [1e-8, 1.0]
+    solver.dt_abs_tol = solver.dt_rel_tol = 1e-4
+    solver.dt_param = [_Parameter(0.4), _Parameter(0.4)]
+    solver.t_param = [_Parameter(1.0), _Parameter(0.6)]
+    solver._dt_for_next_time_to_hit = lambda: 1.0
+
+    accepted, _, _, component = solver._update_time_step(nonlinear_failed=True)
+
+    assert accepted is False
+    assert component == 'nonlinear solve'
+    assert solver.dt_param[0].Get() < 0.4
+    assert solver.t_param[0].Get() == 0.6
+
+
+def test_three_step_nonlinear_failure_rejects_and_halves_dt() -> None:
+    solver = _three_step()
+    solver.dt_range = [1e-8, 1.0]
+    solver.dt_abs_tol = solver.dt_rel_tol = 1e-4
+    solver.dt_param = [_Parameter(0.4), _Parameter(0.2), _Parameter(0.4)]
+    solver.t_param = [_Parameter(1.0), _Parameter(1.0), _Parameter(0.6)]
+    solver.step = 2
+    solver._dt_for_next_time_to_hit = lambda: 1.0
+
+    accepted, _, _, component = solver._update_time_step(nonlinear_failed=True)
+
+    assert accepted is False
+    assert component == 'nonlinear solve'
+    assert solver.dt_param[0].Get() == 0.2
+    assert solver.dt_param[1].Get() == 0.1
+    assert solver.t_param[0].Get() == 0.6
+    assert solver.step == 1