move halley into directory

Author: vchuravy

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

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

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