Reuse jacobian_caches & use alg_autodiff to choose correct VecJac operator

Author: avik-pal

src/NonlinearSolve.jl

@@ -34,10 +34,10 @@ function SciMLBase.__solve(prob::NonlinearProblem,
 end
 
 include("utils.jl")
-include("jacobian.jl")
 include("raphson.jl")
 include("trustRegion.jl")
 include("levenberg.jl")
+include("jacobian.jl")
 include("ad.jl")
 
 import PrecompileTools

src/jacobian.jl

@@ -92,7 +92,7 @@ function build_jac_and_jac_config(alg, f::F1, uf::F2, du1, u, tmp, du2) where {F
 
     J = if !linsolve_needs_jac
         # We don't need to construct the Jacobian
-        JacVec(uf, u; autodiff = AutoFiniteDiff())
+        JacVec(uf, u; autodiff = alg_autodiff(alg) ? AutoForwardDiff() : AutoFiniteDiff())
     else
         if f.jac_prototype === nothing
             ArrayInterface.undefmatrix(u)
@@ -104,6 +104,27 @@ function build_jac_and_jac_config(alg, f::F1, uf::F2, du1, u, tmp, du2) where {F
     return J, jac_config
 end
 
+# Build Jacobian Caches
+function jacobian_caches(alg::Union{NewtonRaphson, LevenbergMarquardt, TrustRegion}, f, u,
+    p, ::Val{true})
+    uf = JacobianWrapper(f, p)
+
+    du1 = zero(u)
+    du2 = zero(u)
+    tmp = zero(u)
+    J, jac_config = build_jac_and_jac_config(alg, f, uf, du1, u, tmp, du2)
+
+    linprob = LinearProblem(J, _vec(zero(u)); u0 = _vec(zero(u)))
+    weight = similar(u)
+    recursivefill!(weight, true)
+
+    Pl, Pr = wrapprecs(alg.precs(J, nothing, u, p, nothing, nothing, nothing, nothing,
+            nothing)..., weight)
+    linsolve = init(linprob, alg.linsolve; alias_A = true, alias_b = true, Pl, Pr)
+
+    uf, linsolve, J, du1, jac_config
+end
+
 function get_chunksize(jac_config::ForwardDiff.JacobianConfig{
     T,
     V,

src/levenberg.jl

@@ -226,26 +226,6 @@ mutable struct LevenbergMarquardtCache{iip, fType, algType, uType, duType, resTy
     end
 end
 
-function jacobian_caches(alg::LevenbergMarquardt, f, u, p, ::Val{true})
-    uf = JacobianWrapper(f, p)
-
-    du1 = zero(u)
-    du2 = zero(u)
-    tmp = zero(u)
-    J, jac_config = build_jac_and_jac_config(alg, f, uf, du1, u, tmp, du2)
-
-    linprob = LinearProblem(J, _vec(zero(u)); u0 = _vec(zero(u)))
-    weight = similar(u)
-    # Q: Setting this to false leads to residual = 0 in GMRES?
-    recursivefill!(weight, true)
-
-    Pl, Pr = wrapprecs(alg.precs(J, nothing, u, p, nothing, nothing, nothing, nothing,
-            nothing)..., weight)
-    linsolve = init(linprob, alg.linsolve; alias_A = true, alias_b = true, Pl, Pr)
-
-    uf, linsolve, J, du1, jac_config
-end
-
 function jacobian_caches(alg::LevenbergMarquardt, f, u, p, ::Val{false})
     JacobianWrapper(f, p), nothing, ArrayInterface.undefmatrix(u), nothing, nothing
 end

src/raphson.jl

@@ -104,26 +104,6 @@ mutable struct NewtonRaphsonCache{iip, fType, algType, uType, duType, resType, p
     end
 end
 
-function jacobian_caches(alg::NewtonRaphson, f, u, p, ::Val{true})
-    uf = JacobianWrapper(f, p)
-
-    du1 = zero(u)
-    du2 = zero(u)
-    tmp = zero(u)
-    J, jac_config = build_jac_and_jac_config(alg, f, uf, du1, u, tmp, du2)
-
-    linprob = LinearProblem(J, _vec(zero(u)); u0 = _vec(zero(u)))
-    weight = similar(u)
-    # Q: Setting this to false leads to residual = 0 in GMRES?
-    recursivefill!(weight, true)
-
-    Pl, Pr = wrapprecs(alg.precs(J, nothing, u, p, nothing, nothing, nothing, nothing,
-            nothing)..., weight)
-    linsolve = init(linprob, alg.linsolve; alias_A = true, alias_b = true, Pl, Pr)
-
-    uf, linsolve, J, du1, jac_config
-end
-
 function jacobian_caches(alg::NewtonRaphson, f, u, p, ::Val{false})
     JacobianWrapper(f, p), nothing, nothing, nothing, nothing
 end

src/trustRegion.jl

@@ -278,26 +278,6 @@ mutable struct TrustRegionCache{iip, fType, algType, uType, resType, pType,
     end
 end
 
-function jacobian_caches(alg::TrustRegion, f, u, p, ::Val{true})
-    uf = JacobianWrapper(f, p)
-
-    du1 = zero(u)
-    du2 = zero(u)
-    tmp = zero(u)
-    J, jac_config = build_jac_and_jac_config(alg, f, uf, du1, u, tmp, du2)
-
-    linprob = LinearProblem(J, _vec(zero(u)); u0 = _vec(zero(u)))
-    weight = similar(u)
-    # Q: Setting this to false leads to residual = 0 in GMRES?
-    recursivefill!(weight, true)
-
-    Pl, Pr = wrapprecs(alg.precs(J, nothing, u, p, nothing, nothing, nothing, nothing,
-            nothing)..., weight)
-    linsolve = init(linprob, alg.linsolve; alias_A = true, alias_b = true, Pl, Pr)
-
-    uf, linsolve, J, du1, jac_config
-end
-
 function jacobian_caches(alg::TrustRegion, f, u, p, ::Val{false})
     J = ArrayInterface.undefmatrix(u)
     JacobianWrapper(f, p), nothing, J, zero(u), nothing