Implement some stats and better handling of inner solver kwargs

Author: vchuravy

examples/bratu.jl

@@ -43,67 +43,97 @@ const dx = 1 / (N + 1) # Grid-spacing
 x = LinRange(0.0 + dx, 1.0 - dx, N)
 u₀ = sin.(x .* π)
 
+lines(x, u₀, label = "Inital guess")
+
 # ## Reference solution evaluated over domain
 reference = true_sol_bratu.(x)
 
-# ## Solving using inplace and out-of-place variants
-uₖ_1 = newton_krylov!(
+f, ax = lines(x, u₀, label = "Inital guess")
+lines!(ax, x, reference, label = "Reference solution")
+axislegend(ax, position = :cb)
+f
+
+# ## Solving using inplace variant and CG
+uₖ, _ = newton_krylov!(
     (res, u) -> bratu!(res, u, dx, λ),
     copy(u₀), similar(u₀);
     Solver = CgSolver,
 )
 
-uₖ_2 = newton_krylov(
+ϵ = abs2.(uₖ .- reference)
+
+let
+    fig = Figure(size = (800, 800))
+    ax = Axis(fig[1, 1], title = "", ylabel = "", xlabel = "")
+
+    lines!(ax, x, reference, label = "True solution")
+    lines!(ax, x, u₀, label = "Initial guess")
+    lines!(ax, x, uₖ, label = "Newton-Krylov solution")
+
+    axislegend(ax, position = :cb)
+
+    ax = Axis(fig[1, 2], title = "Error", ylabel = "abs2 err", xlabel = "")
+    lines!(ax, abs2.(uₖ .- reference))
+
+    fig
+end
+
+# ## Solving using the out of place variant
+
+_, stats = newton_krylov(
     (u) -> bratu(u, dx, λ),
     copy(u₀);
     Solver = CgSolver
 )
-
-ϵ1 = abs2.(uₖ_1 .- reference)
-ϵ2 = abs2.(uₖ_1 .- reference)
+stats
 
 # ## Solving with a fixed forcing
-newton_krylov!(
+_, stats = newton_krylov!(
     (res, u) -> bratu!(res, u, dx, λ),
     copy(u₀), similar(u₀);
     Solver = CgSolver,
-    forcing = NewtonKrylov.Fixed()
+    forcing = NewtonKrylov.Fixed(0.1)
 )
+stats
 
 # ## Solving with no forcing
-newton_krylov!(
+_, stats = newton_krylov!(
     (res, u) -> bratu!(res, u, dx, λ),
     copy(u₀), similar(u₀);
     Solver = CgSolver,
     forcing = nothing
 )
+stats
 
-# ## Solve using GMRES -- very slow
+# ## Solve using GMRES -- doesn't converge
 # ```julia
-# @time newton_krylov!(
-# 	(res, u) -> bratu!(res, u, dx, λ),
-# 	copy(u₀), similar(u₀);
-# 	Solver = GmresSolver,
+# _, stats = newton_krylov!(
+#     (res, u) -> bratu!(res, u, dx, λ),
+#     copy(u₀), similar(u₀);
+#     Solver = GmresSolver,
 # )
+# stats
 # ```
 
 # ## Solve using GMRES + ILU Preconditoner
-@time newton_krylov!(
+_, stats = newton_krylov!(
     (res, u) -> bratu!(res, u, dx, λ),
     copy(u₀), similar(u₀);
     Solver = GmresSolver,
     N = (J) -> ilu(collect(J)), # Assembles the full Jacobian
-    ldiv = true,
+    krylov_kwargs = (; ldiv = true)
 )
+stats
 
 # ## Solve using FGMRES + ILU Preconditoner
-@time newton_krylov!(
+_, stats = newton_krylov!(
     (res, u) -> bratu!(res, u, dx, λ),
     copy(u₀), similar(u₀);
     Solver = FgmresSolver,
     N = (J) -> ilu(collect(J)), # Assembles the full Jacobian
-    ldiv = true,
+    krylov_kwargs = (; ldiv = true)
 )
+stats
 
 # ## Solve using FGMRES + GMRES Preconditoner
 struct GmresPreconditioner{JOp}
@@ -116,20 +146,21 @@ function LinearAlgebra.mul!(y, P::GmresPreconditioner, x)
     return copyto!(y, sol)
 end
 
-@time newton_krylov!(
+_, stats = newton_krylov!(
     (res, u) -> bratu!(res, u, dx, λ),
     copy(u₀), similar(u₀);
     Solver = FgmresSolver,
-    N = (J) -> GmresPreconditioner(J, 30),
+    N = (J) -> GmresPreconditioner(J, 5),
 )
+stats
 
 # ## Explodes..
 # ```julia
 # newton_krylov!(
 # 	(res, u) -> bratu!(res, u, dx, λ),
 # 	copy(u₀), similar(u₀);
-# 	verbose = 1,
 # 	Solver = CglsSolver, # CgneSolver
+#   krylov_kwargs = (; verbose=1)
 # )
 # ```
 #
@@ -138,7 +169,7 @@ end
 # 	(res, u) -> bratu!(res, u, dx, λ),
 # 	copy(u₀), similar(u₀);
 # 	verbose = 1,
-# 	Solver = BicgstabSolver,
+# 	Solver = BicgstabSolver, # L=2
 # 	η_max = nothing
 # )
 # ```

src/NewtonKrylov.jl

@@ -128,6 +128,17 @@ function newton_krylov!(F!, u₀, M::Int = length(u₀); kwargs...)
     return newton_krylov!(F!, u₀, res; kwargs...)
 end
 
+struct Stats
+    outer_iterations::Int
+    inner_iterations::Int
+end
+function update(stats::Stats, inner_iterations)
+    return Stats(
+        stats.outer_iterations + 1,
+        stats.inner_iterations + inner_iterations
+    )
+end
+
 """
 
 ## Arguments
@@ -151,8 +162,9 @@ function newton_krylov!(
         Solver = CgSolver,
         M = nothing,
         N = nothing,
-        ldiv = false
+        krylov_kwargs = (;)
     )
+    t₀ = time_ns()
     F!(res, u) # res = F(u)
     n_res = norm(res)
     tol = tol_rel * n_res + tol_abs
@@ -166,10 +178,10 @@ function newton_krylov!(
     J = JacobianOperator(F!, res, u)
     solver = Solver(J, res)
 
-    n_iter = 1
-    while n_res > tol && n_iter <= max_niter
+    stats = Stats(0, 0)
+    while n_res > tol && stats.outer_iterations <= max_niter
         # Handle kwargs for Preconditoners
-        kwargs = (; ldiv, verbose = verbose - 1)
+        kwargs = krylov_kwargs
         if N !== nothing
             kwargs = (; N = N(J), kwargs...)
         end
@@ -198,22 +210,19 @@ function newton_krylov!(
         n_res = norm(res)
 
         if isinf(n_res) || isnan(n_res)
-            error("Inner solver blew up at iter=$n_iter")
+            @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_iter n_res η
-        n_iter += 1
+        verbose > 0 && @info "Newton" iter = n_res η stats
+        stats = update(stats, solver.stats.niter)
     end
-    return u
+    t = (time_ns() - t₀) / 1.0e9
+    return u, (; solved = n_res <= tol, stats, t)
 end
 
-# TODO:
-# - Better statistic
-# - Allow choice of Krylov solver
-#   - maxit_krylov?
-
 end # module NewtonKrylov