diff --git a/docs/src/index.md b/docs/src/index.md index e4f8634..090cb76 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -4,11 +4,20 @@ Newton Method using Krylov.jl (montoison-orban-2023)[@cite] ## API +#### Newton-Krylov + ```@docs newton_krylov! newton_krylov ``` +#### Halley-Krylov + +```@docs +halley_krylov! +halley_krylov +``` + ### Parameters ```@docs diff --git a/docs/src/refs.bib b/docs/src/refs.bib index 1605a8a..3cefc22 100644 --- a/docs/src/refs.bib +++ b/docs/src/refs.bib @@ -51,3 +51,17 @@ @ARTICLE{Kan2022-ko urldate = {2024-10-21}, language = {en} } + +@ARTICLE{Tan2025-al, + title = {Scalable higher-order nonlinear solvers via higher-order + automatic differentiation}, + author = {Tan, Songchen and Miao, Keming and Edelman, Alan and + Rackauckas, Christopher}, + journaltitle = {arXiv [math.NA]}, + date = {2025-01-28}, + eprint = {2501.16895}, + eprinttype = {arXiv}, + eprintclass = {math.NA}, + urldate = {2025-03-05} +} + diff --git a/examples/halley.jl b/examples/halley.jl new file mode 100644 index 0000000..8c96375 --- /dev/null +++ b/examples/halley.jl @@ -0,0 +1,41 @@ +## Simple 2D example from (Kelley2003)[@cite] + +using NewtonKrylov, LinearAlgebra +using CairoMakie + +function F!(res, x) + res[1] = x[1]^2 + x[2]^2 - 2 + return res[2] = exp(x[1] - 1) + x[2]^2 - 2 +end + +function F(x) + res = similar(x) + F!(res, x) + return res +end + +xs = LinRange(-3, 8, 1000) +ys = LinRange(-15, 10, 1000) + +levels = [0.1, 0.25, 0.5:2:10..., 10.0:10:200..., 200:100:4000...] + +fig, ax = contour(xs, ys, (x, y) -> norm(F([x, y])); levels) + +trace_1 = let x₀ = [2.0, 0.5] + xs = Vector{Tuple{Float64, Float64}}(undef, 0) + hist(x, res, n_res) = (push!(xs, (x[1], x[2])); nothing) + x, stats = newton_krylov!(F!, x₀, callback = hist) + xs +end +lines!(ax, trace_1) + +trace_2 = let x₀ = [2.0, 0.5] + xs = Vector{Tuple{Float64, Float64}}(undef, 0) + hist(x, res, n_res) = (push!(xs, (x[1], x[2])); nothing) + x, stats = halley_krylov!(F!, x₀, similar(x₀), callback = hist, verbose = 1, forcing = nothing) + @show stats + xs +end +lines!(ax, trace_2) + +fig diff --git a/src/NewtonKrylov.jl b/src/NewtonKrylov.jl index b342d02..b602ae3 100644 --- a/src/NewtonKrylov.jl +++ b/src/NewtonKrylov.jl @@ -1,6 +1,7 @@ module NewtonKrylov export newton_krylov, newton_krylov! +export halley_krylov, halley_krylov! using Krylov using LinearAlgebra, SparseArrays @@ -89,6 +90,35 @@ function Base.collect(JOp::JacobianOperator) return J end +""" + HessianOperator + +Calculcates H(F, u) * v * v +""" +struct HessianOperator{F, A} + J::JacobianOperator{F, A} + J_cache::JacobianOperator{F, A} +end +HessianOperator(J) = HessianOperator(J, Enzyme.make_zero(J)) + +Base.size(H::HessianOperator) = size(H.J) +Base.eltype(H::HessianOperator) = eltype(H.J) + +function mul!(out, H::HessianOperator, v) + _out = similar(H.J.res) # TODO cache in H + Enzyme.make_zero!(H.J_cache) + H.J_cache.u .= v + autodiff( + Forward, + mul!, + DuplicatedNoNeed(_out, out), + DuplicatedNoNeed(H.J, H.J_cache), + Const(v) + ) + + return nothing +end + ## # Newton-Krylov ## @@ -252,15 +282,15 @@ function newton_krylov!( end # Solve: Jx = -res - # res is modifyed by J, so we create a copy `-res` - # TODO: provide a temporary storage for `-res` - solve!(solver, J, -res; kwargs...) + # res is modifyed by J, so we create a copy `res` + # TODO: provide a temporary storage for `res` + solve!(solver, J, copy(res); kwargs...) d = solver.x # Newton direction s = 1 # Newton step TODO: LineSearch # Update u - u .+= s .* d + @. u -= s * d # Update residual and norm n_res_prior = n_res @@ -278,11 +308,138 @@ function newton_krylov!( η = forcing(η, tol, n_res, n_res_prior) end - verbose > 0 && @info "Newton" iter = n_res η stats + verbose > 0 && @info "Newton" iter = n_res η = (forcing !== nothing ? η : nothing) stats stats = update(stats, solver.stats.niter) end t = (time_ns() - t₀) / 1.0e9 return u, (; solved = n_res <= tol, stats, t) end +""" + halley_krylov(F, u₀::AbstractArray, M::Int = length(u₀)) + +Halley method after (Tan2025-al)[@cite]. + +## Arguments + - `F`: `F(u)` solves `res = F(u) = 0` + - `u`: Initial guess + - `M`: Length of the output of `F`. Defaults to `length(u₀)`. + +$(KWARGS_DOCS) +""" +function halley_krylov(F, u₀::AbstractArray, M::Int = length(u₀); kwargs...) + F!(res, u) = (res .= F(u); nothing) + return halley_krylov!(F!, u₀, M; kwargs...) +end + +""" + halley_krylov!(F!, u₀::AbstractArray, M::Int = length(u₀)) + +Halley method after (Tan2025-al)[@cite]. + +## Arguments + - `F!`: `F!(res, u)` solves `res = F(u) = 0` + - `u`: Initial guess + - `M`: Length of the output of `F!`. Defaults to `length(u₀)`. + +$(KWARGS_DOCS) +""" +function halley_krylov!(F!, u₀::AbstractArray, M::Int = length(u₀); kwargs...) + res = similar(u₀, M) + return halley_krylov!(F!, u₀, res; kwargs...) +end + +""" + halley_krylov!(F!, u::AbstractArray, res::AbstractArray) + +Halley method after (Tan2025-al)[@cite]. + +## Arguments + - `F!`: `F!(res, u)` solves `res = F(u) = 0` + - `u`: Initial guess + - `res`: Temporary for residual + +$(KWARGS_DOCS) +""" +function halley_krylov!( + F!, u::AbstractArray, res::AbstractArray; + tol_rel = 1.0e-6, + tol_abs = 1.0e-12, + max_niter = 50, + forcing::Union{Forcing, Nothing} = EisenstatWalker(), + verbose = 0, + Solver = GmresSolver, + M = nothing, + N = nothing, + krylov_kwargs = (;), + callback = (args...) -> nothing, + ) + t₀ = time_ns() + F!(res, u) # res = F(u) + n_res = norm(res) + callback(u, res, n_res) + + tol = tol_rel * n_res + tol_abs + + if forcing !== nothing + η = inital(forcing) + end + + verbose > 0 && @info "Jacobian-Free Halley-Krylov" Solver res₀ = n_res tol tol_rel tol_abs η + + J = JacobianOperator(F!, res, u) + H = HessianOperator(J) + solver = Solver(J, res) + + stats = Stats(0, 0) + while n_res > tol && stats.outer_iterations <= max_niter + # Handle kwargs for Preconditoners + kwargs = krylov_kwargs + if N !== nothing + kwargs = (; N = N(J), kwargs...) + end + if M !== nothing + kwargs = (; M = M(J), kwargs...) + end + if forcing !== nothing + # ‖F′(u)d + F(u)‖ <= η * ‖F(u)‖ Inexact Newton termination + kwargs = (; rtol = η, kwargs...) + end + + solve!(solver, J, copy(res); kwargs...) # J \ fx + a = copy(solver.x) + + # calculate hvvp (2nd order directional derivative using the JVP) + hvvp = similar(res) + mul!(hvvp, H, a) + + solve!(solver, J, hvvp; kwargs...) # J \ hvvp + b = solver.x + + # Update u + @. u -= (a * a) / (a - b / 2) + + # Update residual and norm + n_res_prior = n_res + + F!(res, u) # res = F(u) + n_res = norm(res) + callback(u, res, n_res) + + if isinf(n_res) || isnan(n_res) + @error "Inner solver blew up" stats + break + end + + if forcing !== nothing + η = forcing(η, tol, n_res, n_res_prior) + end + + verbose > 0 && @info "Newton" iter = n_res η = (forcing !== nothing ? η : nothing) stats + stats = update(stats, solver.stats.niter) # TODO we do two calls to solver iterations + end + t = (time_ns() - t₀) / 1.0e9 + return u, (; solved = n_res <= tol, stats, t) +end + end # module NewtonKrylov diff --git a/test/runtests.jl b/test/runtests.jl index e515fc0..1c68f9d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -25,16 +25,69 @@ end import NewtonKrylov: JacobianOperator using Enzyme, LinearAlgebra +function df(x, a) + return autodiff(Forward, F, DuplicatedNoNeed(x, a)) |> first +end + +function df!(out, x, a) + res = similar(out) + autodiff(Forward, F!, DuplicatedNoNeed(res, out), DuplicatedNoNeed(x, a)) + return nothing +end + @testset "Jacobian" begin - J_Enz = jacobian(Forward, F, [3.0, 5.0]) |> only - J = JacobianOperator(F!, zeros(2), [3.0, 5.0]) + x = [3.0, 5.0] + v = rand(2) + + J_Enz = jacobian(Forward, F, x) |> only + J = JacobianOperator(F!, zeros(2), x) J_NK = collect(J) @test J_NK == J_Enz + jvp = similar(v) + mul!(jvp, J, v) + + jvp2 = df(x, v) + @test jvp == jvp2 + + jvp3 = similar(v) + df!(jvp3, x, v) + @test jvp == jvp3 + + @test jvp ≈ J_Enz * v +end + +import NewtonKrylov: HessianOperator + +# Differentiate F with respect to x twice. +function ddf(x, a) + return autodiff(Forward, df, DuplicatedNoNeed(x, a), Const(a)) |> first +end + +function ddf!(out, x, a) + _out = similar(out) + autodiff(Forward, df!, DuplicatedNoNeed(_out, out), DuplicatedNoNeed(x, a), Const(a)) + return nothing +end + +@testset "2nd order directional derivative" begin + x = [3.0, 5.0] v = rand(2) - out = similar(v) - mul!(out, J, v) - @test out ≈ J_Enz * v + hvvp = similar(x) + ddf!(hvvp, x, v) + + hvvp2 = ddf(x, v) + # Enzyme seems to disagree with itself only on MacOs/Windows so probably we need to look at + # second derivative of exp (which Enzyme will be using the system libm for iirc) + @test hvvp ≈ hvvp2 + + J = JacobianOperator(F!, zeros(2), x) + H = HessianOperator(J) + + hvvp3 = similar(x) + mul!(hvvp3, H, v) + + @test hvvp == hvvp3 end