Add backtracking options to solver config.

PiperOrigin-RevId: 911274846

Author: theo-brown

docs/configuration.rst

@@ -2270,6 +2270,21 @@ specific solver are defined in the relevant section below.
   Number of corrector steps for the predictor-corrector linear solver. 0 means a
   pure linear solve with no corrector steps. Must be a positive integer.
 
+``atol`` (float | None [default = None])
+  Absolute tolerance for fixed-point iterations in the predictor-corrector solver.
+  If specified, iterations can exit early when the normalized residual falls below this threshold.
+
+``rtol`` (float | None [default = None])
+  Relative tolerance for fixed-point iterations in the predictor-corrector solver.
+  If specified, iterations can exit early when the normalized residual falls below this fraction of the initial residual.
+
+``use_backtracking`` (bool [default = True])
+  Enables backtracking linesearch to improve stability. Can be used with any
+  solver. For the Newton-Raphson solver, this option is always enforced as True.
+
+``delta_reduction_factor`` (float [default = 0.5])
+  Factor by which the step size is reduced during backtracking.
+
 ``use_pereverzev`` (bool [default = False])
   Use Pereverzev-Corrigan terms in the heat and particle flux when using the
   linear solver. Critical for stable calculation of stiff transport, at the cost

docs/solver_details.rst

@@ -234,6 +234,24 @@ these coefficients become known at every iteration step, describing a `linear`
 system of equations. :math:`\mathbf{x}_{t+\Delta t}^k` can then be solved using
 standard linear algebra methods implemented in JAX.
 
+Optionally, the fixed-point iteration can be configured to terminate early
+once a specified tolerance is achieved, rather than running for a fixed number
+of iterations. This is controlled by user-configurable absolute and relative
+tolerances on the residual norm, denoted by :math:`\varepsilon_{abs}` and
+:math:`\varepsilon_{rel}` respectively. The solve iterates until the normalized
+residual falls below the absolute tolerance
+:math:`\| \mathbf{R} \|_{norm} < \varepsilon_{abs}` or becomes smaller than the
+relative tolerance multiplied by the initial residual, i.e.,
+:math:`\| \mathbf{R} \|_{norm} < \varepsilon_{rel} \| \mathbf{R}_{0} \|_{norm}`.
+
+Additionally, a backtracking linesearch can be used to improve stability in
+the solvers. When enabled in fixed-point iteration, if an iteration results
+in an increase in the residual or an invalid state (e.g., NaN values), the
+solver will backtrack along the update direction by reducing the step size.
+For the Newton-Raphson solver, this backtracking linesearch is always required
+and enforced to ensure robustness.
+
+
 To further enhance the stability of the linear solver, particularly in the
 presence of stiff transport coefficients (e.g., when using the QLKNN turbulent
 transport model, see :ref:`physics_models`), the |pereverzev-corrigan-method|

torax/_src/solver/jax_fixed_point.py

@@ -18,6 +18,7 @@
 import jax
 import jax.numpy as jnp
 from torax._src import jax_utils
+from torax._src.solver import linesearch
 
 PyTree: TypeAlias = Any
 
@@ -29,6 +30,11 @@ def fixed_point(
     xtol: float | None = 1e-08,
     maxiter: int = 500,
     method: Literal['del2', 'iteration'] = 'del2',
+    atol: float | None = None,
+    rtol: float | None = None,
+    use_backtracking: bool = False,
+    delta_reduction_factor: float = 0.5,
+    max_backtrack_steps: int = 10,
 ) -> PyTree:
   """A JAX version of `scipy.optimize.fixed_point`.
 
@@ -47,6 +53,12 @@ def fixed_point(
       with Aitken’s Del^2 convergence acceleration, taken from Burden, Faires,
       “Numerical Analysis”, 5th edition, pg. 80. 'iteration' just iterates the
       function until the tolerance is reached.
+    atol: Absolute tolerance on the residual norm.
+    rtol: Relative tolerance on the residual norm.
+    use_backtracking: If true, use backtracking linesearch in 'iteration'
+      method.
+    delta_reduction_factor: Factor by which step_size is reduced each step.
+    max_backtrack_steps: Maximum number of backtracking steps.
 
   Returns:
     The fixed point `jax.Array`.
@@ -56,8 +68,22 @@ def fixed_point(
   if maxiter <= 0:
     raise ValueError(f'Invalid maxiter: {maxiter} must be positive.')
 
-  def body(x):
-    x, count, _ = x
+  def residual_fn(x):
+    return jax.tree.map(lambda a, b: a - b, func(x, *args), x)
+
+  def norm_fn(res):
+    return jnp.sqrt(sum(jnp.sum(leaf**2) for leaf in jax.tree.leaves(res)))
+
+  def residual_norm(x):
+    return norm_fn(residual_fn(x))
+
+  if rtol is not None:
+    initial_residual_norm = residual_norm(x0)
+  else:
+    initial_residual_norm = jnp.array(0.0)
+
+  def body(x_state):
+    x, count, _ = x_state
     out1 = func(x, *args)
     if method == 'del2':
       out2 = func(out1, *args)
@@ -69,9 +95,45 @@ def _del2(p0, p1, p2):
 
       out = jax.tree.map(_del2, x, out1, out2)
     else:
-      out = out1
-
-    if xtol:
+      if use_backtracking:
+        direction = jax.tree.map(lambda a, b: a - b, out1, x)
+
+        init_res = direction
+        init_norm = norm_fn(init_res)
+
+        decrease = 1e-4
+        current_norm_sq = init_norm**2
+
+        def accept_fn(step_size, trial_norm):
+          target = (1.0 - 2.0 * decrease * step_size) * current_norm_sq
+          return (trial_norm**2) <= target
+
+        ls_state = linesearch.backtracking_linesearch(
+            residual_fn=residual_fn,
+            x_init=x,
+            direction=direction,
+            accept_fn=accept_fn,
+            norm_fn=norm_fn,
+            initial_residual=init_res,
+            initial_residual_norm=init_norm,
+            delta_reduction_factor=delta_reduction_factor,
+            max_steps=max_backtrack_steps,
+        )
+        out = ls_state.x
+      else:
+        out = out1
+
+    # Terminate based on residual norm.
+    if atol is not None or rtol is not None:
+      res_norm = residual_norm(out)
+      converged = jnp.array(False, dtype=jnp.bool_)
+      if atol is not None:
+        converged = converged | (res_norm <= atol)
+      if rtol is not None:
+        converged = converged | (res_norm <= rtol * initial_residual_norm)
+      stop = converged
+    # Terminate based on relative error.
+    elif xtol:
 
       def _relative_error(actual, expected):
         relative_error = (actual - expected) / expected
@@ -93,7 +155,7 @@ def cond(x):
   stop = jnp.array(False, dtype=jnp.bool_)
   x_init = (x0, count, stop)
 
-  if xtol is None:
+  if xtol is None and atol is None and rtol is None:
     return jax.lax.fori_loop(0, maxiter, lambda i, val: body(val), x_init)[0]
   else:
     return jax.lax.while_loop(cond, body, x_init)[0]

torax/_src/solver/jax_root_finding.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """JAX root finding functions."""
+
 import dataclasses
 import functools
 from typing import Callable, Final
@@ -21,6 +22,7 @@
 import jax.numpy as jnp
 import numpy as np
 from torax._src import jax_utils
+from torax._src.solver import linesearch
 
 # Delta is a vector. If no entry of delta is above this magnitude, we terminate
 # the delta loop. This is to avoid getting stuck in an infinite loop in edge
@@ -123,9 +125,7 @@ def back(g, y):
 
   if use_jax_custom_root:
     if custom_jac is not None:
-      raise ValueError(
-          'custom_jac is not compatible with use_jax_custom_root.'
-      )
+      raise ValueError('custom_jac is not compatible with use_jax_custom_root.')
     x_out, metadata = jax.lax.custom_root(
         f=fun,
         initial_guess=x0,
@@ -199,109 +199,44 @@ def _body(
   dtype = input_state['x'].dtype
   a_mat = jacobian_fun(input_state['x'])
   rhs = -input_state['residual']
-  # delta = x_new - x_old
-  # tau = delta/delta0, where delta0 is the delta that sets the linearized
-  # residual to zero. tau < 1 when needed such that x_new meets
-  # conditions of reduced residual and valid state quantities.
-  # If tau < taumin while residual > tol, then the routine exits with an
-  # error flag, leading to either a warning or recalculation at lower dt
-  initial_delta_state = {
-      'x': input_state['x'],
-      'delta': jnp.linalg.solve(a_mat, rhs),
-      'residual_old': input_state['residual'],
-      'residual_new': input_state['residual'],
-      'tau': jnp.array(1.0, dtype=dtype),
-  }
-  output_delta_state = _compute_output_delta_state(
-      initial_delta_state, residual_fun, delta_reduction_factor
+
+  direction = jnp.linalg.solve(a_mat, rhs)
+
+  def norm_fn(res):
+    return jnp.mean(jnp.abs(res))
+
+  init_norm = norm_fn(input_state['residual'])
+
+  def accept_fn(step_size, trial_norm):
+    del step_size  # Unused
+    return (trial_norm <= init_norm) & (~jnp.isnan(trial_norm))
+
+  ls_state = linesearch.backtracking_linesearch(
+      residual_fn=residual_fun,
+      x_init=input_state['x'],
+      direction=direction,
+      accept_fn=accept_fn,
+      norm_fn=norm_fn,
+      initial_residual=input_state['residual'],
+      initial_residual_norm=init_norm,
+      delta_reduction_factor=delta_reduction_factor,
+      max_steps=100,
+      min_step_norm=MIN_DELTA,
   )
 
   output_state = {
-      'x': input_state['x'] + output_delta_state['delta'],
-      'residual': output_delta_state['residual_new'],
+      'x': ls_state.x,
+      'residual': ls_state.residual,
       'iterations': jnp.array(input_state['iterations'][...], dtype=dtype) + 1,
-      'last_tau': output_delta_state['tau'],
+      'last_tau': ls_state.step_size,
   }
+
   if log_iterations:
     jax.debug.print(
         'Iteration: {iteration:d}. Residual: {residual:.16f}. tau = {tau:.6f}',
         iteration=output_state['iterations'].astype(jax_utils.get_int_dtype()),
         residual=_residual_scalar(output_state['residual']),
-        tau=output_delta_state['tau'],
+        tau=ls_state.step_size,
     )
 
   return output_state
-
-
-def _compute_output_delta_state(
-    initial_state: dict[str, jax.Array],
-    residual_fun: Callable[[jax.Array], jax.Array],
-    delta_reduction_factor: float,
-):
-  """Updates output delta state."""
-  delta_body_fun = functools.partial(
-      _delta_body,
-      delta_reduction_factor=delta_reduction_factor,
-  )
-  delta_cond_fun = functools.partial(
-      _delta_cond,
-      residual_fun=residual_fun,
-  )
-  output_delta_state = jax.lax.while_loop(
-      delta_cond_fun, delta_body_fun, initial_state
-  )
-
-  x_new = output_delta_state['x'] + output_delta_state['delta']
-  residual_vec_x_new = residual_fun(x_new)
-  output_delta_state |= dict(
-      residual_new=residual_vec_x_new,
-  )
-  return output_delta_state
-
-
-def _delta_cond(
-    delta_state: dict[str, jax.Array],
-    residual_fun: Callable[[jax.Array], jax.Array],
-) -> bool:
-  """Check if delta obtained from Newton step is valid.
-
-  Args:
-    delta_state: see `delta_body`.
-    residual_fun: Residual function.
-
-  Returns:
-    True if the new value of `x` causes any NaNs or has increased the residual
-    relative to the old value of `x`.
-  """
-  x_old = delta_state['x']
-  x_new = x_old + delta_state['delta']
-  residual_vec_x_old = delta_state['residual_old']
-  residual_scalar_x_old = _residual_scalar(residual_vec_x_old)
-  # Avoid sanity checking inside residual, since we directly
-  # afterwards check sanity on the output (NaN checking)
-  # TODO(b/312453092) consider instead sanity-checking x_new
-  with jax_utils.enable_errors(False):
-    residual_vec_x_new = residual_fun(x_new)
-    residual_scalar_x_new = _residual_scalar(residual_vec_x_new)
-    delta_state['residual_new'] = residual_vec_x_new
-  return jnp.bool_(
-      jnp.logical_and(
-          jnp.max(jnp.abs(delta_state['delta'])) > MIN_DELTA,
-          jnp.logical_or(
-              residual_scalar_x_old < residual_scalar_x_new,
-              jnp.isnan(residual_scalar_x_new),
-          ),
-      ),
-  )
-
-
-def _delta_body(
-    input_delta_state: dict[str, jax.Array],
-    delta_reduction_factor: float,
-) -> dict[str, jax.Array]:
-  """Reduces step size for this Newton iteration."""
-  return input_delta_state | dict(
-      delta=input_delta_state['delta'] * delta_reduction_factor,
-      tau=jnp.array(input_delta_state['tau'][...], dtype=jax_utils.get_dtype())
-      * delta_reduction_factor,
-  )