From ca593cc836d062e2844e21a523c0d6a6eed59e6b Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Tue, 4 Mar 2025 13:22:01 +0100 Subject: [PATCH 1/7] implement halley --- examples/halley.jl | 113 ++++++++++++++++++++++++++++++++++++++++++++ src/NewtonKrylov.jl | 21 ++++++++ 2 files changed, 134 insertions(+) create mode 100644 examples/halley.jl diff --git a/examples/halley.jl b/examples/halley.jl new file mode 100644 index 0000000..a649614 --- /dev/null +++ b/examples/halley.jl @@ -0,0 +1,113 @@ +## Simple 2D example from (Kelley2003)[@cite] + +using NewtonKrylov, LinearAlgebra +using CairoMakie +using Krylov, Enzyme + +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 + +function halley(F!, u, res; + tol_rel = 1.0e-6, + tol_abs = 1.0e-12, + max_niter = 50, + Solver = GmresSolver, + ) + + F!(res, u) # res = F(u) + n_res = norm(res) + + tol = tol_rel * n_res + tol_abs + + J = NewtonKrylov.JacobianOperator(F!, res, u) + H = NewtonKrylov.HessianOperator(J) + solver = Solver(J, res) + + for i in :max_niter + if n_res <= tol + break + end + solve!(solver, J, copy(res)) # 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) # J \ hvvp + b = solver.x + + # update + @. u -= (a * a) / (a - b / 2) + + end +end + +# u = [2.0, 0.5] +# res = zeros(2) +# J = NewtonKrylov.JacobianOperator(F!,u,res) +# F!(res, u) +# a, stats = gmres(J, copy(res)) + +# J_cache = Enzyme.make_zero(J) +# out = similar(J.res) +# hvvp = Enzyme.make_zero(out) +# du = Enzyme.make_zero(J.u) +# autodiff(Forward, LinearAlgebra.mul!, +# DuplicatedNoNeed(out, hvvp), +# DuplicatedNoNeed(J, J_cache), +# DuplicatedNoNeed(du, a)) + +# hvvp + +# b, stats = gmres(J, hvvp) +# @. u -= (a * a) / (a - b / 2) + +# a + + +dg_ad(x, dx) = autodiff(Forward, flux, DuplicatedNoNeed(x, dx))[1] +ddg_ad(x, dx, ddx) = autodiff(Forward, dg_ad, DuplicatedNoNeed(x, dx), + DuplicatedNoNeed(dx, ddx))[1] + +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.5, 3.0] + 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_2) + +trace_3 = let x₀ = [3.0, 4.0] + 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, forcing = NewtonKrylov.EisenstatWalker(η_max = 0.68949), verbose = 1) + @show stats.solved + xs +end +lines!(ax, trace_3) + +fig diff --git a/src/NewtonKrylov.jl b/src/NewtonKrylov.jl index b342d02..9f4c718 100644 --- a/src/NewtonKrylov.jl +++ b/src/NewtonKrylov.jl @@ -89,6 +89,27 @@ function Base.collect(JOp::JacobianOperator) return J end +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 + du = Enzyme.make_zero(H.J.u) # TODO cache in H + + autodiff(Forward, mul!, + DuplicatedNoNeed(_out, out), + DuplicatedNoNeed(H.J, H.J_cache), + DuplicatedNoNeed(du, v)) + + return nothing +end + ## # Newton-Krylov ## From c54635511d2b75c978e7008a96ce32eb85cda099 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Tue, 4 Mar 2025 16:32:23 +0100 Subject: [PATCH 2/7] mirror newton_krylov --- examples/halley.jl | 106 +++++++++++++++++++++++++------------------- src/NewtonKrylov.jl | 10 ++--- 2 files changed, 65 insertions(+), 51 deletions(-) diff --git a/examples/halley.jl b/examples/halley.jl index a649614..4f2b3f2 100644 --- a/examples/halley.jl +++ b/examples/halley.jl @@ -2,7 +2,6 @@ using NewtonKrylov, LinearAlgebra using CairoMakie -using Krylov, Enzyme function F!(res, x) res[1] = x[1]^2 + x[2]^2 - 2 @@ -15,68 +14,89 @@ function F(x) return res end -function halley(F!, u, res; +import NewtonKrylov: Forcing, EisenstatWalker, inital, forcing, solve!, + JacobianOperator, HessianOperator, Stats, update, GmresSolver + +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 - J = NewtonKrylov.JacobianOperator(F!, res, u) - H = NewtonKrylov.HessianOperator(J) + 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) - for i in :max_niter - if n_res <= tol - break + 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 - solve!(solver, J, copy(res)) # J \ fx + 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) # J \ hvvp + solve!(solver, J, hvvp; kwargs...) # J \ hvvp b = solver.x - # update + # Update u @. u -= (a * a) / (a - b / 2) - end -end - -# u = [2.0, 0.5] -# res = zeros(2) -# J = NewtonKrylov.JacobianOperator(F!,u,res) -# F!(res, u) -# a, stats = gmres(J, copy(res)) + # Update residual and norm + n_res_prior = n_res -# J_cache = Enzyme.make_zero(J) -# out = similar(J.res) -# hvvp = Enzyme.make_zero(out) -# du = Enzyme.make_zero(J.u) -# autodiff(Forward, LinearAlgebra.mul!, -# DuplicatedNoNeed(out, hvvp), -# DuplicatedNoNeed(J, J_cache), -# DuplicatedNoNeed(du, a)) + F!(res, u) # res = F(u) + n_res = norm(res) + callback(u, res, n_res) -# hvvp - -# b, stats = gmres(J, hvvp) -# @. u -= (a * a) / (a - b / 2) - -# a + 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 -dg_ad(x, dx) = autodiff(Forward, flux, DuplicatedNoNeed(x, dx))[1] -ddg_ad(x, dx, ddx) = autodiff(Forward, dg_ad, DuplicatedNoNeed(x, dx), - DuplicatedNoNeed(dx, ddx))[1] + 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 xs = LinRange(-3, 8, 1000) ys = LinRange(-15, 10, 1000) @@ -93,21 +113,15 @@ trace_1 = let x₀ = [2.0, 0.5] end lines!(ax, trace_1) -trace_2 = let x₀ = [2.5, 3.0] +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 = newton_krylov!(F!, x₀, callback = hist) + x, stats = halley_krylov!(F!, x₀, similar(x₀), callback = hist, verbose=1, forcing=nothing) + @show stats xs end lines!(ax, trace_2) -trace_3 = let x₀ = [3.0, 4.0] - 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, forcing = NewtonKrylov.EisenstatWalker(η_max = 0.68949), verbose = 1) - @show stats.solved - xs -end -lines!(ax, trace_3) +trace_2 fig diff --git a/src/NewtonKrylov.jl b/src/NewtonKrylov.jl index 9f4c718..594c39b 100644 --- a/src/NewtonKrylov.jl +++ b/src/NewtonKrylov.jl @@ -273,15 +273,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 @@ -299,7 +299,7 @@ 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 From 2e3c3f15f09fdf251eabc52c1ec8394d75faa64b Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Wed, 5 Mar 2025 09:59:52 +0100 Subject: [PATCH 3/7] move halley into directory --- examples/halley.jl | 88 +--------------------------------- src/NewtonKrylov.jl | 114 +++++++++++++++++++++++++++++++++++++++++--- test/runtests.jl | 59 +++++++++++++++++++++-- 3 files changed, 162 insertions(+), 99 deletions(-) diff --git a/examples/halley.jl b/examples/halley.jl index 4f2b3f2..8c96375 100644 --- a/examples/halley.jl +++ b/examples/halley.jl @@ -14,90 +14,6 @@ function F(x) return res end -import NewtonKrylov: Forcing, EisenstatWalker, inital, forcing, solve!, - JacobianOperator, HessianOperator, Stats, update, GmresSolver - -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 - xs = LinRange(-3, 8, 1000) ys = LinRange(-15, 10, 1000) @@ -116,12 +32,10 @@ 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) + x, stats = halley_krylov!(F!, x₀, similar(x₀), callback = hist, verbose = 1, forcing = nothing) @show stats xs end lines!(ax, trace_2) -trace_2 - fig diff --git a/src/NewtonKrylov.jl b/src/NewtonKrylov.jl index 594c39b..913825a 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,11 @@ 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} @@ -100,12 +106,15 @@ Base.eltype(H::HessianOperator) = eltype(H.J) function mul!(out, H::HessianOperator, v) _out = similar(H.J.res) # TODO cache in H - du = Enzyme.make_zero(H.J.u) # TODO cache in H - - autodiff(Forward, mul!, - DuplicatedNoNeed(_out, out), - DuplicatedNoNeed(H.J, H.J_cache), - DuplicatedNoNeed(du, v)) + 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 @@ -299,11 +308,102 @@ function newton_krylov!( η = forcing(η, tol, n_res, n_res_prior) end - verbose > 0 && @info "Newton" iter = n_res η=(forcing !== nothing ? η : nothing) 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 +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 + +function halley_krylov!(F!, u₀::AbstractArray, M::Int = length(u₀); kwargs...) + res = similar(u₀, M) + return halley_krylov!(F!, u₀, res; kwargs...) +end + +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..ad51995 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -25,16 +25,65 @@ 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 + +# 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) + @test hvvp == hvvp2 + + J = JacobianOperator(F!, zeros(2), x) + H = HessianOperator(J) + + hvvp3 = similar(x) + mul!(hvvp3, H, v) + + @test hvvp == hvvp3 end From 97b4dd8f731ff756f06006a5a5d6df6b24d3fad0 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Wed, 5 Mar 2025 10:26:22 +0100 Subject: [PATCH 4/7] small fixes --- docs/src/index.md | 9 +++++++++ src/NewtonKrylov.jl | 30 ++++++++++++++++++++++++++++++ test/runtests.jl | 2 ++ 3 files changed, 41 insertions(+) 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/src/NewtonKrylov.jl b/src/NewtonKrylov.jl index 913825a..00838dd 100644 --- a/src/NewtonKrylov.jl +++ b/src/NewtonKrylov.jl @@ -315,16 +315,46 @@ function newton_krylov!( return u, (; solved = n_res <= tol, stats, t) end +""" + halley_krylov(F, u₀::AbstractArray, M::Int = length(u₀)) + +## 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₀)) + +## 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) + +## 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, diff --git a/test/runtests.jl b/test/runtests.jl index ad51995..a6abf1e 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -58,6 +58,8 @@ end @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 From 774302b42ad7f967054c71b70003dccd03786f0c Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Wed, 5 Mar 2025 10:32:59 +0100 Subject: [PATCH 5/7] use approx --- test/runtests.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index a6abf1e..1c68f9d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -79,7 +79,9 @@ end ddf!(hvvp, x, v) hvvp2 = ddf(x, v) - @test hvvp == hvvp2 + # 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) From 968694aff33f7051c24a372017a79714a5332549 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Wed, 5 Mar 2025 10:39:45 +0100 Subject: [PATCH 6/7] add citation --- src/NewtonKrylov.jl | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/NewtonKrylov.jl b/src/NewtonKrylov.jl index 00838dd..b602ae3 100644 --- a/src/NewtonKrylov.jl +++ b/src/NewtonKrylov.jl @@ -318,6 +318,8 @@ 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 @@ -333,6 +335,8 @@ 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 @@ -348,6 +352,8 @@ 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 From 5e4126598a67bdd7ed34c4e3f26f6020c6c4bb19 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Wed, 5 Mar 2025 10:40:04 +0100 Subject: [PATCH 7/7] fixup! add citation --- docs/src/refs.bib | 14 ++++++++++++++ 1 file changed, 14 insertions(+) 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} +} +