mirror newton_krylov

Author: vchuravy

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

src/NewtonKrylov.jl

@@ -221,15 +221,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
@@ -247,7 +247,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