← Back to Mesh-based solvers Module · scimba_jax

Nonlinear solvers

The discretizations presented so far assume a linear problem, where one matrix is assembled and one linear system is solved. This is no longer the case when the coefficients of the PDE depend on the unknown, for instance

\(-\nabla\cdot(A(u)\nabla u) = f\)

with a coefficient A that depends on u. The four strategies below handle this case.

1. Fixed point iterations

\(-\nabla\cdot(A(u_k)\nabla u_{k+1}) = f\)

This is the oldest strategy, even though it is not the default choice. The coefficient A is evaluated at the previous iterate and treated as a known function of x. The resulting linear problem is then solved with the FEM or DG solvers from the previous pages. These steps are repeated until two successive iterates agree. The convergence is linear at best, and is not guaranteed when the coefficient is not monotone in u. 

# Fixed point: A is evaluated at the previous iterate, and a
# linear problem is solved:
#   u_{k+1} = solve(-div(A(u_k) grad u) = f)
# There is no separate class in scimba_jax: Newton (below)
# uses the same
# linear solve at each step, and converges quadratically.
# Fixed point does not differentiate A(u), which helps if A is
# only known numerically.

2. Newton's method

\(J(u_k)\,\delta u = -R(u_k), \qquad u_{k+1} = u_k + \delta u\)

Newton's method computes a correction using the exact Jacobian, where A(u) is differentiated with respect to u. It converges quadratically once it is close to the solution. In practice, it is the fastest method whenever it converges, so the default nonlinear solver of scimba_jax is plain Newton. Acceleration or damping strategies (line search, Anderson acceleration, Levenberg-Marquardt) are only turned on when a specific problem needs them. 

from scimba_jax.linear_approximation.solvers.newton import (
    NewtonSolver,
)
# Exact linearisation: solve J(u_k) delta = -F(u_k), where the
# Jacobian J
# is computed by differentiating A(u).
solver = NewtonSolver(max_iter=20, tol=1e-10)
scheme, report = EllipticFEscheme.solve(
    scheme, solver=solver, return_report=True
)
# 6 iterations to 1e-10. A fixed point loop can need dozens of
# iterations,
# or diverge, depending on how strongly A varies with u.

3. Newton preconditioned by multigrid

On a fine mesh, the inner linear solve of Newton's method is the most expensive part, and Multigrid can be used to precondition it. However, the preconditioned operator changes at every Newton step, so the multigrid has to be updated at each step, without a full rebuild. 

from scimba_jax.linear_approximation.solvers.multigrid import (
    newton_krylov_relinearised,
)
# The linear solve of each Newton step is preconditioned by the multigrid,
# updated at every step since the operator changes.
# Only what depends on the Jacobian is recomputed (smoother diagonal, coarse
# operator); the meshes, transfers and colouring are reused.
dofsl, n_newton, n_krylov, residual = newton_krylov_relinearised(
    scheme, dofsl_init, mg, tol=1e-10, max_iter=20, every=1
)
# 64x64 mesh, 3 levels: 9 Newton steps, 30 CG iterations, 0.4 s, against
# 46 s (33 800 CG iterations) for a multigrid that is never updated.

4. FAS: nonlinear multigrid

The Full Approximation Scheme (FAS) does not linearise the problem. It applies a multigrid cycle directly to the nonlinear residual, and corrects the right-hand side of the coarse level so that the coarse problem is consistent with the fine one (Brandt's "tau" correction). This method requires neither Newton steps, nor a Jacobian, nor a Krylov loop. In exchange, the problem on the coarsest level is nonlinear, and the smoother must be applied to the current operator, which changes at each iteration. 

from scimba_jax.linear_approximation.solvers.multigrid import MG
from scimba_jax.linear_approximation.solvers.newton import NewtonSolver
from scimba_jax.linear_approximation.solvers.smoothers import DampedJacobiSmoother
# A linear multigrid transfers a correction (assuming A(u+e) = Au + Ae), while
# FAS transfers the solution and corrects the coarse RHS with Brandt's "tau" term.
mg = MG(
    hierarchy,
    DampedJacobiSmoother(omega=0.6, nonlinear=True),
    nu_pre=2,
    nu_post=2,
    fas=True,
    coarse_solver=NewtonSolver(max_iter=30, tol=1e-13),
)
zeros = scheme._initial_dofs()
dofsl, n_cycles, residual = mg.solve(jnp.zeros(zeros.size), tol=1e-9, relinearise=True)
# 14 FAS cycles, 1.3 s, mostly spent in the nonlinear smoother and coarsest solve.
# `relinearise=True` makes the smoother use A(u_k) at the current iterate;
# without it, the iterations diverge to nan on this problem.

\(-\nabla\cdot\big((1 + 10 u^2)\nabla u\big) = f\)

Measurements on the equation above, with a 64×64 mesh and 3 levels: Newton preconditioned by a multigrid updated at every step needs 9 outer steps and 30 conjugate gradient iterations in total, for 0.4 s. The same Newton loop with a multigrid built once and never updated needs 33 800 conjugate gradient iterations for the same result, for 46 s. FAS reaches the same accuracy in 14 cycles and 1.3 s, without any linear solve.