From 4dcc2f66212eb72b6d4a314f43e56a0229576404 Mon Sep 17 00:00:00 2001
From: Frank Schaefer <knarf-schaefle@web.de>
Date: Wed, 13 Sep 2023 09:29:46 -0400
Subject: [PATCH] reuse option in Newton Raphson method

---
 src/raphson.jl     | 37 ++++++++++++++++++++++++++-----------
 test/basictests.jl |  2 +-
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/src/raphson.jl b/src/raphson.jl
index 24e5799fd..924bea8d3 100644
--- a/src/raphson.jl
+++ b/src/raphson.jl
@@ -2,7 +2,7 @@
 ```julia
 NewtonRaphson(; chunk_size = Val{0}(), autodiff = Val{true}(),
               standardtag = Val{true}(), concrete_jac = nothing,
-              diff_type = Val{:forward}, linsolve = nothing, precs = DEFAULT_PRECS)
+              diff_type = Val{:forward}, linsolve = nothing, precs = DEFAULT_PRECS, reuse = true)
 ```
 
 An advanced NewtonRaphson implementation with support for efficient handling of sparse
@@ -42,7 +42,9 @@ for large-scale and numerically-difficult nonlinear systems.
     preconditioners. For more information on specifying preconditioners for LinearSolve
     algorithms, consult the
     [LinearSolve.jl documentation](https://docs.sciml.ai/LinearSolve/stable/).
-
+  - `reuse`: Determines if the Jacobian is reused between (quasi-)Newton steps. Defaults to
+    `true`. If `true` we check convergence, and automatically take a new Jacobian if
+    convergence slows or starts to diverge. If `false`, the Jacobian is updated in each step.
 !!! note
 
     Currently, the linear solver and chunk size choice only applies to in-place defined
@@ -52,15 +54,16 @@ struct NewtonRaphson{CS, AD, FDT, L, P, ST, CJ} <:
        AbstractNewtonAlgorithm{CS, AD, FDT, ST, CJ}
     linsolve::L
     precs::P
+    reuse::Bool
 end
 
 function NewtonRaphson(; chunk_size = Val{0}(), autodiff = Val{true}(),
     standardtag = Val{true}(), concrete_jac = nothing,
-    diff_type = Val{:forward}, linsolve = nothing, precs = DEFAULT_PRECS)
+    diff_type = Val{:forward}, linsolve = nothing, precs = DEFAULT_PRECS, reuse = true)
     NewtonRaphson{_unwrap_val(chunk_size), _unwrap_val(autodiff), diff_type,
         typeof(linsolve), typeof(precs), _unwrap_val(standardtag),
         _unwrap_val(concrete_jac)}(linsolve,
-        precs)
+        precs, reuse)
 end
 
 mutable struct NewtonRaphsonCache{iip, fType, algType, uType, duType, resType, pType,
@@ -105,7 +108,7 @@ mutable struct NewtonRaphsonCache{iip, fType, algType, uType, duType, resType, p
 end
 
 function jacobian_caches(alg::NewtonRaphson, f, u, p, ::Val{false})
-    JacobianWrapper(f, p), nothing, nothing, nothing, nothing
+    JacobianWrapper(f, p), nothing, ArrayInterface.undefmatrix(u), nothing, nothing
 end
 
 function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::NewtonRaphson,
@@ -129,7 +132,6 @@ function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::NewtonRaphson
         fu = f(u, p)
     end
     uf, linsolve, J, du1, jac_config = jacobian_caches(alg, f, u, p, Val(iip))
-
     return NewtonRaphsonCache{iip}(f, alg, u, fu, p, uf, linsolve, J, du1, jac_config,
         false, maxiters, internalnorm,
         ReturnCode.Default, abstol, prob, NLStats(1, 0, 0, 0, 0))
@@ -138,7 +140,13 @@ end
 function perform_step!(cache::NewtonRaphsonCache{true})
     @unpack u, fu, f, p, alg = cache
     @unpack J, linsolve, du1 = cache
-    jacobian!(J, cache)
+    @unpack reuse = alg
+
+    if !reuse || cache.stats.njacs == 0
+        jacobian!(J, cache)
+        cache.stats.njacs += 1
+        @info cache.stats.njacs J
+    end
 
     # u = u - J \ fu
     linres = dolinsolve(alg.precs, linsolve, A = J, b = _vec(fu), linu = _vec(du1),
@@ -151,22 +159,29 @@ function perform_step!(cache::NewtonRaphsonCache{true})
         cache.force_stop = true
     end
     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, fu, f, p = cache
-    J = jacobian(cache, f)
+    @unpack u, fu, f, p, alg = cache
+    @unpack J = cache
+    @unpack reuse = alg
+    @show cache.internalnorm(cache.fu)
+    if !reuse || cache.stats.njacs == 0
+        J = jacobian(cache, f)
+        cache.J = J
+        cache.stats.njacs += 1
+        @info cache.stats.njacs J
+    end
+
     cache.u = u - J \ fu
     cache.fu = f(cache.u, p)
     if iszero(cache.fu) || cache.internalnorm(cache.fu) < cache.abstol
         cache.force_stop = true
     end
     cache.stats.nf += 1
-    cache.stats.njacs += 1
     cache.stats.nsolve += 1
     cache.stats.nfactors += 1
     return nothing
diff --git a/test/basictests.jl b/test/basictests.jl
index 05a0152fa..87b9fc7c0 100644
--- a/test/basictests.jl
+++ b/test/basictests.jl
@@ -62,7 +62,7 @@ u0 = [1.0, 1.0]
 
 precs = [
     NonlinearSolve.DEFAULT_PRECS,
-    (args...) -> (Diagonal(rand!(similar(u0))), nothing)
+    (args...) -> (Diagonal(rand!(similar(u0))), nothing),
 ]
 
 for prec in precs, linsolve in (nothing, KrylovJL_GMRES())