[WIP] reuse option in Newton Raphson method

src/raphson.jl

@@ -1,6 +1,6 @@
 """
     NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = LineSearch(),
-        precs = DEFAULT_PRECS, adkwargs...)
+        precs = DEFAULT_PRECS, reuse = true, reusetol = 1e-6, adkwargs...)
 
 An advanced NewtonRaphson implementation with support for efficient handling of sparse
 matrices via colored automatic differentiation and preconditioned linear solvers. Designed
@@ -29,32 +29,51 @@
   - `linesearch`: the line search algorithm to use. Defaults to [`LineSearch()`](@ref),
     which means that no line search is performed. Algorithms from `LineSearches.jl` can be
     used here directly, and they will be converted to the correct `LineSearch`.
+  - `reuse`: Determines if the Jacobian is reused between (quasi-)Newton steps. Defaults to
+    `false`. If `true` we check how far we stepped with the same Jacobian, and automatically
+    take a new Jacobian if we stepped more than `reusetol` or if convergence slows or starts
+    to diverge. If `false`, the Jacobian is updated in each step.
 """
 @concrete struct NewtonRaphson{CJ, AD} <:
                  AbstractNewtonAlgorithm{CJ, AD}
     ad::AD
     linsolve
     precs
     linesearch
+    reusetol
+    reuse::Bool
 end
 
 function set_ad(alg::NewtonRaphson{CJ}, ad) where {CJ}
-    return NewtonRaphson{CJ}(ad, alg.linsolve, alg.precs, alg.linesearch)
+    return NewtonRaphson{CJ}(ad,
+        alg.linsolve,
+        alg.precs,
+        alg.linesearch,
+        alg.reusetol,
+        alg.reuse)
 end
 
 function NewtonRaphson(; concrete_jac = nothing, linsolve = nothing,
-        linesearch = LineSearch(), precs = DEFAULT_PRECS, adkwargs...)
+        linesearch = LineSearch(), precs = DEFAULT_PRECS, reuse = false, reusetol = 1e-1,
+        adkwargs...)
     ad = default_adargs_to_adtype(; adkwargs...)
     linesearch = linesearch isa LineSearch ? linesearch : LineSearch(; method = linesearch)
-    return NewtonRaphson{_unwrap_val(concrete_jac)}(ad, linsolve, precs, linesearch)
+    return NewtonRaphson{_unwrap_val(concrete_jac)}(ad,
+        linsolve,
+        precs,
+        linesearch,
+        reusetol,
+        reuse)
 end
 
 @concrete mutable struct NewtonRaphsonCache{iip} <: AbstractNonlinearSolveCache{iip}
     f
     alg
     u
     u_prev
+    Δu
     fu1
+    res_norm_prev
     fu2
     du
     p
@@ -81,26 +100,55 @@
     alg = get_concrete_algorithm(alg_, prob)
     @unpack f, u0, p = prob
     u = alias_u0 ? u0 : deepcopy(u0)
+    u_prev = copy(u)
+    Δu = zero(u)
+
     fu1 = evaluate_f(prob, u)
+    res_norm_prev = internalnorm(fu1)
+
     uf, linsolve, J, fu2, jac_cache, du = jacobian_caches(alg, f, u, p, Val(iip);
         linsolve_kwargs)
 
     abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fu1, u,
         termination_condition)
 
-    return NewtonRaphsonCache{iip}(f, alg, u, copy(u), fu1, fu2, du, p, uf, linsolve, J,
+    return NewtonRaphsonCache{iip}(f, alg, u, u_prev, Δu, fu1, res_norm_prev, fu2, du, p,
+        uf,
+        linsolve, J,
         jac_cache, false, maxiters, internalnorm, ReturnCode.Default, abstol, reltol, prob,
         NLStats(1, 0, 0, 0, 0),
         init_linesearch_cache(alg.linesearch, f, u, p, fu1, Val(iip)), tc_cache)
 end
 
 function perform_step!(cache::NewtonRaphsonCache{true})
-    @unpack u, u_prev, fu1, f, p, alg, J, linsolve, du = cache
-    jacobian!!(J, cache)
+    @unpack u, u_prev, Δu, fu1, res_norm_prev, f, p, alg, J, linsolve, du = cache
+    @unpack reuse = alg
+
+    if reuse
+        # check if residual increased and check how far we stepped
+        res_norm = cache.internalnorm(fu1)
+        update = (res_norm > res_norm_prev) || (cache.internalnorm(Δu) > alg.reusetol)
+        if update || cache.stats.njacs == 0
+            jacobian!!(J, cache)
+            cache.stats.njacs += 1
+            Δu .*= false
+            # u = u - J \ fu
+            linres = dolinsolve(alg.precs, linsolve; A = J, b = _vec(fu1), linu = _vec(du),
+                p, reltol = cache.abstol)
+        else
+            # u = u - J \ fu
+            linres = dolinsolve(alg.precs, linsolve; b = _vec(fu1), linu = _vec(du),
+                p, reltol = cache.abstol)
+        end
+        cache.res_norm_prev = res_norm
+    else
+        jacobian!!(J, cache)
+        cache.stats.njacs += 1
+        # u = u - J \ fu
+        linres = dolinsolve(alg.precs, linsolve; A = J, b = _vec(fu1), linu = _vec(du),
+            p, reltol = cache.abstol)
+    end
 
-    # u = u - J \ fu
-    linres = dolinsolve(alg.precs, linsolve; A = J, b = _vec(fu1), linu = _vec(du),
-        p, reltol = cache.abstol)
     cache.linsolve = linres.cache
 
     # Line Search
@@ -109,26 +157,57 @@
     f(cache.fu1, u, p)
 
     check_and_update!(cache, cache.fu1, cache.u, cache.u_prev)
-
+    @. Δu += u - u_prev
     @. u_prev = u
     cache.stats.nf += 1
-    cache.stats.njacs += 1
     cache.stats.nsolve += 1
     cache.stats.nfactors += 1
     return nothing
 end
 
 function perform_step!(cache::NewtonRaphsonCache{false})
-    @unpack u, u_prev, fu1, f, p, alg, linsolve = cache
-
-    cache.J = jacobian!!(cache.J, cache)
-    # u = u - J \ fu
-    if linsolve === nothing
-        cache.du = fu1 / cache.J
+    @unpack u, u_prev, Δu, fu1, res_norm_prev, f, p, alg, linsolve = cache
+    @unpack reuse = alg
+
+    if reuse
+        # check if residual increased and check how far we stepped
+        res_norm = cache.internalnorm(fu1)
+        update = (res_norm > res_norm_prev) || (cache.internalnorm(Δu) > alg.reusetol)
+        if update || cache.stats.njacs == 0
+            cache.J = jacobian!!(cache.J, cache)
+            cache.stats.njacs += 1
+            # u = u - J \ fu
+            if linsolve === nothing
+                cache.du = fu1 / cache.J
+            else
+                linres = dolinsolve(alg.precs, linsolve; A = cache.J, b = _vec(fu1),
+                    linu = _vec(cache.du), p, reltol = cache.abstol)
+                cache.linsolve = linres.cache
+            end
+        else
+            # u = u - J \ fu
+            if linsolve === nothing
+                cache.du = fu1 / cache.J
+            else
+                linres = dolinsolve(alg.precs, linsolve; b = _vec(fu1),
+                    linu = _vec(cache.du), p, reltol = cache.abstol)
+                cache.linsolve = linres.cache
+            end
+        end
+        cache.res_norm_prev = res_norm
     else
-        linres = dolinsolve(alg.precs, linsolve; A = cache.J, b = _vec(fu1),
-            linu = _vec(cache.du), p, reltol = cache.abstol)
-        cache.linsolve = linres.cache
+        cache.J = jacobian!!(cache.J, cache)
+        # cache.Δu *= false
+        cache.stats.njacs += 1
+
+        # u = u - J \ fu
+        if linsolve === nothing
+            cache.du = fu1 / cache.J
+        else
+            linres = dolinsolve(alg.precs, linsolve; A = cache.J, b = _vec(fu1),
+                linu = _vec(cache.du), p, reltol = cache.abstol)
+            cache.linsolve = linres.cache
+        end
     end
 
     # Line Search
@@ -137,10 +216,9 @@
     cache.fu1 = f(cache.u, p)
 
     check_and_update!(cache, cache.fu1, cache.u, cache.u_prev)
-
+    cache.Δu = @. cache.u - cache.u_prev
     cache.u_prev = cache.u
     cache.stats.nf += 1
-    cache.stats.njacs += 1
     cache.stats.nsolve += 1
     cache.stats.nfactors += 1
     return nothing

test/23_test_problems.jl

@@ -11,7 +11,7 @@ function test_on_library(problems, dicts, alg_ops, broken_tests, ϵ = 1e-4)
         @testset "$idx: $(dict["title"])" begin
             for alg in alg_ops
                 try
-                    sol = solve(nlprob, alg;
+                    sol = solve(nlprob, alg; maxiters = 1_000_000,
                         termination_condition = AbsNormTerminationMode())
                     problem(res, sol.u, nothing)
 
@@ -30,12 +30,15 @@ function test_on_library(problems, dicts, alg_ops, broken_tests, ϵ = 1e-4)
     end
 end
 
-@testset "NewtonRaphson 23 Test Problems" begin
-    alg_ops = (NewtonRaphson(),)
+# NewtonRaphson
+@testset "NewtonRaphson test problem library" begin
+    alg_ops = (NewtonRaphson(; reuse = false),
+        NewtonRaphson(; reuse = true))
 
     # dictionary with indices of test problems where method does not converge to small residual
     broken_tests = Dict(alg => Int[] for alg in alg_ops)
     broken_tests[alg_ops[1]] = [1, 6]
+    broken_tests[alg_ops[2]] = [1, 6]
 
     test_on_library(problems, dicts, alg_ops, broken_tests)
 end