Nonlinear Poisson equation

Two ways to solve the same nonlinear problem — fixed-point (Picard) iteration and Newton's method — so the difference between linear and quadratic convergence is something measured, not just asserted. Every number and every plot below was produced by the code shown.

Problem

\[-\left(\alpha(u) u'\right)' = g \text{ in } (0,1), \qquad u(0) = u(1) = u_{\text{exact}},\]

with a diffusion coefficient that depends on the unknown itself,

\[\alpha(u) = 3 + \frac{1}{1+u^2},\]

and the manufactured solution $u_{\text{exact}}(x) = e^{x}$, with $g$ calculated so that it is exactly satisfied.

using Bramble
using Random

sol(x) = exp(x[1])
α(u) = 3 + 1 / (1 + u^2)
dαdu(u) = -2u / (1 + u^2)^2
rhs(x) = -dαdu(sol(x)) * sol(x)^2 - α(sol(x)) * sol(x)

Ω = domain(interval(0.0, 1.0))
Random.seed!(20260903)
Ωₕ = mesh(Ω, 40, false)
Wₕ = gridspace(Ωₕ)

bcs = dirichlet_constraints(Ω, :boundary => sol)
gₕ = element(Wₕ)
avgₕ!(gₕ, rhs)
l = form(Wₕ, v -> innerₕ(gₕ, v))
F = assemble(l; dirichlet = bcs)

The seed is what makes the numbers below reproducible: false draws the interior points from the global RNG, so without it the mesh – and the iteration counts quoted below – would differ from build to build, and the suite could not assert what the page prints.

The right-hand side never changes across the iteration — only the diffusion matrix does, since only it depends on the current guess for $u$. α is evaluated at the average of the previous iterate, M₋ₕ, the standard discretization for a nonlinear flux.

Fixed-point (Picard) iteration

Linearize by freezing $\alpha$ at the previous iterate, solve, repeat. The pattern of the diffusion matrix — which entries are ever nonzero — never changes between iterations, only the values in it do, so it is allocated once with allocate_system_matrix and refilled with assemble! rather than rebuilt with assemble every step. αvals is a plain VectorElement the form closes over, not a fresh vector computed each time: mutating it in place (αvals .= α.(M₋ₕ(uₙ))) is what assemble! picks up on the next refill, the same "live coefficient" the forms tutorial relies on:

uₙ = element(Wₕ, 0.0)
αvals = element(Wₕ)
αvals .= α.(M₋ₕ(uₙ))
a = form(Wₕ, Wₕ, (U, V) -> inner₊(αvals * ∇₋ₕ(U), ∇₋ₕ(V)))
A = allocate_system_matrix(a)

picard_steps = Float64[]
for it in 1:200
    assemble!(A, a; dirichlet = :boundary)
    unew = A \ F
    step = maximum(abs, unew .- parent(uₙ))
    push!(picard_steps, step)
    uₙ .= unew
    αvals .= α.(M₋ₕ(uₙ))
    step < 1e-12 && break
end
length(picard_steps), picard_steps[[1, 2, 3, end]]

(9, [2.718281828459045, 0.05862611946851448, 0.0005965799435467822, 6.039613253960852e-14])

The step size drops by one to two orders of magnitude each time here, reaching machine precision in 9 iterations — still only linear convergence (a roughly constant per-step ratio, not the per-step squaring Newton gets below), just a fast-converging instance of it for this particular coefficient and mesh.

Newton's method

The residual $R(u) = A(u) u - F$ is the same matrix, applied to the vector it was built from rather than solved against. Boundary rows come along for free: dirichlet already replaces them with the identity before the residual ever sees them, so $R_i(u) = u_i - u_{\text{exact}}(x_i)$ there, and the Jacobian's boundary rows are the identity too, with no separate case to write.

That Jacobian is sparse — R inherits the same local stencil A itself has, a handful of nonzeros per row rather than a dense matrix — so it is computed with DifferentiationInterface's sparse AD rather than a plain ForwardDiff.jacobian: SparseConnectivityTracer finds which entries can possibly be nonzero, SparseMatrixColorings groups the independent columns so one ForwardDiff sweep gets several of them at once, and prepare_jacobian does both once, reused every Newton step since the sparsity pattern does not change across iterations, only the values do.

The Picard loop above could allocate its matrix once because it never leaves Float64. The residual below cannot use that same trick directly: T is Float64 on a plain call but a ForwardDiff.Dual while prepare_jacobian/jacobian are probing it, and a matrix allocated for one element type cannot hold values of the other — so diffusion_matrix builds a fresh, T-typed matrix (pattern included) on every call, the same way it always did:

using ForwardDiff, DifferentiationInterface
import SparseConnectivityTracer, SparseMatrixColorings

const sparse_ad = AutoSparse(AutoForwardDiff();
    sparsity_detector = SparseConnectivityTracer.TracerSparsityDetector(),
    coloring_algorithm = SparseMatrixColorings.GreedyColoringAlgorithm())

function diffusion_matrix(uₕ)
    αvals_local = α.(M₋ₕ(uₕ))
    a = form(Wₕ, Wₕ, (U, V) -> inner₊(αvals_local * ∇₋ₕ(U), ∇₋ₕ(V)))
    return assemble(a; dirichlet = :boundary)
end

function residual(u_vec::AbstractVector{T}) where {T}
    uₕ = element(Wₕ, T)
    uₕ .= u_vec
    A = diffusion_matrix(uₕ)
    return A * u_vec .- F
end

u = zeros(ndofs(Wₕ))
prep = prepare_jacobian(residual, sparse_ad, u)
J = DifferentiationInterface.jacobian(residual, prep, sparse_ad, u)  # once, for its sparse structure
newton_residuals = Float64[]
for it in 1:20
    r = residual(u)
    push!(newton_residuals, sqrt(sum(abs2, r)))
    newton_residuals[end] < 1e-10 && break
    DifferentiationInterface.jacobian!(residual, J, prep, sparse_ad, u)
    u .-= J \ r
end
length(newton_residuals), newton_residuals

(5, [3.054654960712615, 0.29577091853112736, 0.0010514747533858445, 7.841474412743256e-9, 1.2494226777187792e-12])

Close to allocation-free, not quite: the two rebuilds this step avoids — the Jacobian's sparsity pattern, and the diffusion matrix's own pattern inside assemble — were the two largest costs, but diffusion_matrix still rebuilds a fresh matrix, values and pattern both, on every call, because residual has to stay generic over T (Float64 on a plain call, ForwardDiff.Dual while jacobian! is probing it) and a matrix allocated for one element type cannot hold the other. Measured behind a function barrier: a plain residual(u) call costs 17,536 B (rebuilding A once, at T = Float64); a full Newton step costs 112,896 B (that, plus rebuilding it again at T = Dual for every colour jacobian!'s sparse sweep needs).

Quadratic convergence — the residual's correct digits roughly double each step, against Picard's roughly-constant gain of one — visible directly in how fast that list reaches machine precision. Both methods reach the same solution, and both are close to the true one, measured the same way the linear example measures it:

uₕ_newton = element(Wₕ)
uₕ_newton .= u
uexact = Rₕ(Wₕ, sol)
norm₁ₕ(uₕ_newton .- uexact), norm₁ₕ(uₙ .- uexact)

(0.00029131009520692186, 0.0002913100951362921)

Closing the gap: caching the diffusion matrix by element type

diffusion_matrix rebuilds its pattern on every call for a real reason — T differs between a plain call and a jacobian! sweep, and a Float64 matrix cannot hold a Dual — but the pattern itself is exactly as fixed across element types as it is across Newton iterations: only α's values differ, and only because they were evaluated at a different T. type_cached_assemble! gives that pattern a place to live per type it is ever reached at, instead of rebuilding it from nothing every time. build_diffusion is named and defined once, the same reason a above is built once outside the Picard loop rather than inside it; refill! reaches for M₋ₓ! rather than M₋ₕ, which would allocate a fresh result every call and reintroduce exactly the cost this is meant to stop paying:

function build_diffusion(uₕ)
    Mu = element(Wₕ, eltype(uₕ))
    αvals = element(Wₕ, eltype(uₕ))
    a = form(Wₕ, Wₕ, (U, V) -> inner₊(αvals * ∇₋ₕ(U), ∇₋ₕ(V)))
    refill!(uₕ) = begin
        M₋ₓ!(Mu, uₕ)
        αvals .= α.(Mu)
    end
    return a, refill!
end

cache = Dict()
diffusion_matrix_cached(uₕ) = type_cached_assemble!(
    build_diffusion, cache, uₕ; dirichlet = :boundary)

function residual_cached(u_vec::AbstractVector{T}) where {T}
    uₕ = element(Wₕ, T)
    uₕ .= u_vec
    A = diffusion_matrix_cached(uₕ)
    return A * u_vec .- F
end

u_cached = zeros(ndofs(Wₕ))
prep_cached = prepare_jacobian(residual_cached, sparse_ad, u_cached)
J_cached = DifferentiationInterface.jacobian(residual_cached, prep_cached, sparse_ad, u_cached)
newton_residuals_cached = Float64[]
for it in 1:20
    r = residual_cached(u_cached)
    push!(newton_residuals_cached, sqrt(sum(abs2, r)))
    newton_residuals_cached[end] < 1e-10 && break
    DifferentiationInterface.jacobian!(residual_cached, J_cached, prep_cached, sparse_ad, u_cached)
    u_cached .-= J_cached \ r
end
newton_residuals_cached, maximum(abs.(u_cached .- u))

([3.054654960712615, 0.29577091853112736, 0.0010514747533858445, 7.841474412743256e-9, 1.2494226777187792e-12], 0.0)

Same convergence, same answer, and only residual_cached and diffusion_matrix_cached (the first call at each of T = Float64 and T = Dual still pays to build and to allocate_system_matrix) differ from residual/diffusion_matrix above. Measured the same way, behind the same function barrier: a plain residual_cached(u) call, once both types have been seen, costs 2,880 B against residual's 17,536 B; a full Newton step costs 74,016 B against 112,896 B. What is left is not zero — cache's value type is necessarily Any, since the cached (a, refill!, A) triple's own concrete type differs across T, so fetching it back out still pays a small, fixed dictionary/dynamic-dispatch cost — but that cost does not grow with the mesh, unlike the pattern rebuild it replaces (see type_cached_assemble!'s own docstring and test/form/type_cached_assemble.jl for the same comparison run at a mesh 100 times larger).

Skipping the tracer: a Bramble-native pattern

SparseConnectivityTracer above finds the Jacobian's sparsity pattern by tracing residual — running it once with a special value that records which inputs reach which outputs. That works for any Julia function, which is exactly why it needs to run the function at all: a tracing pass, on top of the coloring pass that follows it.

residual here is not an arbitrary function, though — it is A(u) * u - F, where A comes from allocate_system_matrix, whose own sparsity is already known directly from a's AST — no tracing needed for that part at all. The only piece missing from A's own pattern is the extra chain-rule term from αvals_local's own dependence on u through M₋ₕ. jacobian_pattern supplies exactly that piece — named the same way the coefficient itself was built, U -> M₋ₕ(U) — and hands the result to ADTypes.KnownJacobianSparsityDetector in place of the tracer:

using ADTypes: KnownJacobianSparsityDetector

αvals_pattern = α.(M₋ₕ(element(Wₕ, 0.0)))
a_for_pattern = form(Wₕ, Wₕ, (U, V) -> inner₊(αvals_pattern * ∇₋ₕ(U), ∇₋ₕ(V)))
pattern = jacobian_pattern(a_for_pattern, U -> M₋ₕ(U))

sparse_ad_manual = AutoSparse(AutoForwardDiff();
    sparsity_detector = KnownJacobianSparsityDetector(pattern),
    coloring_algorithm = SparseMatrixColorings.GreedyColoringAlgorithm())

ast_sparsity_detector spells the same thing more directly, once ADTypes.jl is loaded — no separate pattern variable, no KnownJacobianSparsityDetector wrapper, the same detector either way. This is the one actually driving the Newton loop below, not just sparse_ad_manual shown for what it desugars to:

native_ad = AutoSparse(AutoForwardDiff();
    sparsity_detector = ast_sparsity_detector(a_for_pattern, U -> M₋ₕ(U)),
    coloring_algorithm = SparseMatrixColorings.GreedyColoringAlgorithm())

a_for_pattern only needs some concrete coefficient to build a BilinearForm from — the pattern is a property of the AST, not of αvals_pattern's values, so evaluating it at u = 0 is as good as evaluating it at the true solution. Feeding native_ad into the same prepare_jacobian/jacobian! loop as before reaches the same pattern (118 nonzeros, both ways, on this mesh) and the same quadratic convergence:

u_native = zeros(ndofs(Wₕ))
prep_native = prepare_jacobian(residual, native_ad, u_native)
J_native = DifferentiationInterface.jacobian(residual, prep_native, native_ad, u_native)
newton_residuals_native = Float64[]
for it in 1:20
    r = residual(u_native)
    push!(newton_residuals_native, sqrt(sum(abs2, r)))
    newton_residuals_native[end] < 1e-10 && break
    DifferentiationInterface.jacobian!(residual, J_native, prep_native, native_ad, u_native)
    u_native .-= J_native \ r
end
newton_residuals_native

5-element Vector{Float64}:
 3.054654960712615
 0.29577091853112736
 0.0010514747533858445
 7.841474412743256e-9
 1.2494226777187792e-12

What changes is what prepare_jacobian has to pay for: no tracing pass, only coloring. Measured on this mesh, prepare_jacobian costs 0.140 ms with the tracer against 0.062 ms given the pattern directly — jacobian_pattern itself costs 0.023 ms of that 0.062, read straight off a's AST. The gap widens with the mesh: tracing cost scales with however long one residual call takes to run and record, while jacobian_pattern only ever walks the grid once, touching neither ForwardDiff nor the coefficient's actual values.

Checking the answer

The same nested-random-mesh pattern as the linear example — one random coarse mesh per dimension, refined in place with iterative_refinement! — using Newton at every level, since it needs by far the fewest solves to reach machine precision. A dense Jacobian would have made 2D and 3D here impractical (O(n^2) memory for a matrix that is actually O(n)-nonzero); the sparse one keeps every level below a few seconds even at tens of thousands of degrees of freedom:

function nonlinear_series(D::Int; n0::Int = 5, levels::Int)
    sol_d(x) = exp(sum(x))
    rhs_d(x) = -D * dαdu(sol_d(x)) * sol_d(x)^2 - D * α(sol_d(x)) * sol_d(x)
    Ωd = domain(reduce(×, ntuple(_ -> interval(0.0, 1.0), D)))
    Ωc = mesh(Ωd, ntuple(_ -> n0, D), ntuple(_ -> false, D))
    hs, errs = Float64[], Float64[]
    for level in 1:levels
        Wc = gridspace(Ωc)
        bcs_c = dirichlet_constraints(Ωd, :boundary => sol_d)
        g_c = element(Wc)
        avgₕ!(g_c, rhs_d)
        l_c = form(Wc, v -> innerₕ(g_c, v))
        F_c = assemble(l_c; dirichlet = bcs_c)

        Ac(uₕ) = begin
            Mu = M₋ₕ(uₕ)
            αv = D == 1 ? α.(Mu) : ntuple(i -> α.(Mu[i]), D)
            grad(U) = D == 1 ? αv * ∇₋ₕ(U) : ntuple(i -> αv[i] * ∇₋ₕ(U)[i], D)
            assemble(form(Wc, Wc, (U, V) -> inner₊(grad(U), ∇₋ₕ(V)));
                dirichlet = :boundary)
        end
        rc(uv::AbstractVector{T}) where {T} = begin
            uₕ = element(Wc, T)
            uₕ .= uv
            Ac(uₕ) * uv .- F_c
        end

        uc = zeros(ndofs(Wc))
        prep_c = prepare_jacobian(rc, sparse_ad, uc)
        J_c = DifferentiationInterface.jacobian(rc, prep_c, sparse_ad, uc)
        for it in 1:20
            r = rc(uc)
            sqrt(sum(abs2, r)) < 1e-10 && break
            DifferentiationInterface.jacobian!(rc, J_c, prep_c, sparse_ad, uc)
            uc .-= J_c \ r
        end
        uexact_c = Rₕ(Wc, sol_d)

        push!(hs, hₘₐₓ(Ωc))
        push!(errs, norm₁ₕ(element(Wc) .= uc .- parent(uexact_c)))
        level < levels && iterative_refinement!(Ωc)
    end
    return hs, errs
end

Random.seed!(20260903)
hs1, errs1 = nonlinear_series(1; n0 = 6, levels = 7)
Random.seed!(20260903)
hs2, errs2 = nonlinear_series(2; levels = 5)
Random.seed!(20260903)
hs3, errs3 = nonlinear_series(3; levels = 4)

order1 = log(errs1[end - 1] / errs1[end]) / log(hs1[end - 1] / hs1[end])
order2 = log(errs2[end - 1] / errs2[end]) / log(hs2[end - 1] / hs2[end])
order3 = log(errs3[end - 1] / errs3[end]) / log(hs3[end - 1] / hs3[end])
(order1, order2, order3)

(1.999946784465622, 1.9959056138135924, 1.9217474793134068)

order1 > 1.9 && order2 > 1.9 && order3 > 1.8

true

convergence_plot([(hs1, errs1, "1D", "#5B5FC7"), (hs2, errs2, "2D", "#0E7C86"), (hs3, errs3, "3D", "#B26A00")];

Second order in every dimension, same as the linear problem — the nonlinearity changes how many solves it takes to reach a given $u$, not the discretization's own accuracy once it has.

nonlinear_series above uses sparse_ad, the tracer, at every level and dimension — the same substitution shown earlier (ast_sparsity_detector(a, U -> M₋ₕ(U)) in place of sparse_ad's sparsity_detector) works here unchanged, D-tuple coefficient and all: jacobian_pattern flattens whatever M₋ₕ(U) returns — one node in 1D, a D-tuple in 2D/3D — the same way before taking its reach, so nothing about Ac/grad above needs to change to swap it in. Not re-run a second time here only to save the doc build the cost of solving the same nine problems twice for an answer already shown identical above.