← Back to Mesh-based solvers Module · scimba_jax

Nonlinear solvers

Every discretization so far assumed a linear problem: assemble one matrix, solve one linear system, done. As soon as the PDE's own coefficients depend on the unknown —

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

instead of a fixed A — that stops being true, and the four strategies below are what to reach for, roughly in the order they should be tried.

1. Fixed point — freeze the coefficient, solve, repeat

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

The oldest strategy, and worth understanding even though it is not the one to reach for by default: evaluate A at the previous iterate, treat it as a known function of x, and solve the resulting LINEAR problem — exactly the FEM or DG solve from the earlier pages. Repeat until two iterates agree. It converges linearly at best, and is not guaranteed to converge at all on a coefficient that is not monotone in u. 

# The oldest strategy: FREEZE A at the previous iterate and solve the
# resulting LINEAR problem, repeating until two iterates agree.
#
#   u_{k+1} = solve(-div(A(u_k) grad u) = f)     -- a LINEAR solve
#
# Picard / fixed-point: converges LINEARLY at best, and need not converge
# at all if A is not monotone in u. scimba_jax ships no separate class for
# it -- Newton below needs the same per-step solve and converges
# QUADRATICALLY, so it is the entry point rather than a hand-written loop.
# What fixed point buys over Newton: it never differentiates A(u), useful
# if A is only known numerically (e.g. from a table).

2. Newton — differentiate instead of freezing

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

Newton solves for a correction using the TRUE Jacobian — A(u) differentiated with respect to u, not held fixed — and converges quadratically once it is close to the solution. In practice it wins whenever it converges at all: scimba_jax's default nonlinear solver is plain Newton, with none of the acceleration or damping strategies (line search, Anderson acceleration, Levenberg-Marquardt) turned on unless a specific problem needs one. 

from scimba_jax.linear_approximation.solvers.newton import NewtonSolver

# Exact linearisation: solve J(u_k) delta = -F(u_k) for a correction, with
# J the true Jacobian (differentiated through A(u), not frozen).
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 on the same problem can need
# dozens, or fail to converge, depending on how strongly A varies with u.

3. Newton, preconditioned by a multigrid that follows it

On a fine mesh, Newton's inner linear solve is the expensive part — which is exactly what Multigrid addresses. The complication is that the operator being preconditioned changes at every Newton step, so the multigrid has to be refreshed without paying for a full rebuild each time. 

from scimba_jax.linear_approximation.solvers.multigrid import (
    newton_krylov_relinearised,
)

# Newton's linear solve is the expensive part on a fine mesh -- precondition
# it with the SAME multigrid, refreshed each step since the operator moves.
# `relinearise` redoes only what depends on the Jacobian (smoother diagonal,
# coarse operator), reusing the meshes, transfers and colouring.
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 refreshed.

4. FAS — a multigrid with no Jacobian at all

The Full Approximation Scheme skips linearisation entirely: it runs a multigrid cycle directly on the nonlinear residual, correcting the coarse level's right-hand side so it solves a problem consistent with the fine one (Brandt's "tau" correction). No Newton step, no Jacobian, no Krylov loop — at the cost of a genuinely nonlinear solve on the coarsest level and a smoother that must sweep the current, moving operator rather than a frozen one. 

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

# FAS: a multigrid that never linearises. The ordinary cycle descends a
# CORRECTION (assumes A(u+e) = Au + Ae, false here); FAS descends the
# SOLUTION itself, correcting 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 -- no Newton step, no Jacobian, no Krylov loop; the
# cost moves to a NONLINEAR smoother and a real solve on the coarsest level.
# `relinearise=True` is not a tuning knob: the smoother sweeps A(u_k) at a
# fixed point, and sweeping the WRONG one diverges to nan here, measured.

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

Measured on the equation above, a 64×64 mesh, 3 levels: Newton preconditioned by a multigrid that is refreshed at every step needs 9 outer steps and 30 conjugate-gradient iterations in total, 0.4 s; the same Newton loop with a multigrid built once and never refreshed needs 33 800 conjugate-gradient iterations for the same answer, 46 s. FAS reaches the same accuracy in 14 cycles and 1.3 s, with no linear solve at all.