implement halley

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

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}
+}
+

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

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

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