Merge pull request #201 from avik-pal/ap/jnfk

Add support for JFNK

Author: ChrisRackauckas

Project.toml

@@ -41,11 +41,14 @@ julia = "1.6"
 [extras]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics"]
+test = ["BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra"]

docs/Project.toml

@@ -1,6 +1,7 @@
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 NonlinearSolveMINPACK = "c100e077-885d-495a-a2ea-599e143bf69d"
 SciMLNLSolve = "e9a6253c-8580-4d32-9898-8661bb511710"
@@ -12,6 +13,7 @@ Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 [compat]
 BenchmarkTools = "1"
 Documenter = "0.27"
+LinearSolve = "2"
 NonlinearSolve = "1"
 NonlinearSolveMINPACK = "0.1"
 SciMLNLSolve = "0.1"

docs/src/solvers/BracketingSolvers.md

@@ -7,7 +7,7 @@ Solves for ``f(t)=0`` in the problem defined by `prob` using the algorithm
 
 ## Recommended Methods
 
-`ITP()` is the recommended method for the scalar interval root-finding problems. It is particularly well-suited for cases where the function is smooth and well-behaved; and achieved superlinear convergence while retaining the optimal worst-case performance of the Bisection method. For more details, consult the detailed solver API docs. 
+`ITP()` is the recommended method for the scalar interval root-finding problems. It is particularly well-suited for cases where the function is smooth and well-behaved; and achieved superlinear convergence while retaining the optimal worst-case performance of the Bisection method. For more details, consult the detailed solver API docs.
 `Ridder` is a hybrid method that uses the value of function at the midpoint of the interval to perform an exponential interpolation to the root. This gives a fast convergence with a guaranteed convergence of at most twice the number of iterations as the bisection method.
 `Brent` is a combination of the bisection method, the secant method and inverse quadratic interpolation. At every iteration, Brent's method decides which method out of these three is likely to do best, and proceeds by doing a step according to that method. This gives a robust and fast method, which therefore enjoys considerable popularity.
 

docs/src/tutorials/nonlinear.md

@@ -31,3 +31,21 @@ uspan = (1.0, 2.0) # brackets
 probB = IntervalNonlinearProblem(f, uspan)
 sol = solve(probB, Falsi())
 ```
+
+## Using Jacobian Free Newton Krylov (JNFK) Methods
+
+If we want to solve the first example, without constructing the entire Jacobian
+
+```@example
+using NonlinearSolve, LinearSolve
+
+function f!(res, u, p)
+    @. res = u * u - p
+end
+u0 = [1.0, 1.0]
+p = 2.0
+probN = NonlinearProblem(f!, u0, p)
+
+linsolve = LinearSolve.KrylovJL_GMRES()
+sol = solve(probN, NewtonRaphson(; linsolve), reltol = 1e-9)
+```

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

@@ -24,6 +24,8 @@ function jacobian_finitediff!(J, f, x, jac_config, cache)
     2 * maximum(jac_config.colorvec))
 end
 
+# NoOp for Jacobian if it is not a Abstract Array -- For eg, JacVec Operator
+jacobian!(J, cache) = J
 function jacobian!(J::AbstractMatrix{<:Number}, cache)
     f = cache.f
     uf = cache.uf
@@ -52,14 +54,16 @@ function jacobian!(J::AbstractMatrix{<:Number}, cache)
     nothing
 end
 
-function build_jac_config(alg, f::F1, uf::F2, du1, u, tmp, du2) where {F1, F2}
+function build_jac_and_jac_config(alg, f::F1, uf::F2, du1, u, tmp, du2) where {F1, F2}
     haslinsolve = hasfield(typeof(alg), :linsolve)
 
-    if !SciMLBase.has_jac(f) && # No Jacobian if has analytical solution
-       ((concrete_jac(alg) === nothing && (!haslinsolve || (haslinsolve && # No Jacobian if linsolve doesn't want it
-           (alg.linsolve === nothing || LinearSolve.needs_concrete_A(alg.linsolve))))) ||
-        (concrete_jac(alg) !== nothing && concrete_jac(alg))) # Jacobian if explicitly asked for
-        jac_prototype = f.jac_prototype
+    has_analytic_jac = SciMLBase.has_jac(f)
+    linsolve_needs_jac = (concrete_jac(alg) === nothing &&
+                          (!haslinsolve || (haslinsolve && (alg.linsolve === nothing ||
+                             LinearSolve.needs_concrete_A(alg.linsolve)))))
+    alg_wants_jac = (concrete_jac(alg) !== nothing && concrete_jac(alg))
+
+    if !has_analytic_jac && (linsolve_needs_jac || alg_wants_jac)
         sparsity, colorvec = sparsity_colorvec(f, u)
 
         if alg_autodiff(alg)
@@ -70,25 +74,55 @@ function build_jac_config(alg, f::F1, uf::F2, du1, u, tmp, du2) where {F1, F2}
             else
                 typeof(ForwardDiff.Tag(uf, eltype(u)))
             end
-            jac_config = ForwardColorJacCache(uf, u, _chunksize; colorvec = colorvec,
-                sparsity = sparsity, tag = T)
+            jac_config = ForwardColorJacCache(uf, u, _chunksize; colorvec, sparsity,
+                tag = T)
         else
             if alg_difftype(alg) !== Val{:complex}
-                jac_config = FiniteDiff.JacobianCache(tmp, du1, du2, alg_difftype(alg),
-                    colorvec = colorvec,
-                    sparsity = sparsity)
+                jac_config = FiniteDiff.JacobianCache(tmp, du1, du2, alg_difftype(alg);
+                    colorvec, sparsity)
             else
                 jac_config = FiniteDiff.JacobianCache(Complex{eltype(tmp)}.(tmp),
-                    Complex{eltype(du1)}.(du1), nothing,
-                    alg_difftype(alg), eltype(u),
-                    colorvec = colorvec,
-                    sparsity = sparsity)
+                    Complex{eltype(du1)}.(du1), nothing, alg_difftype(alg), eltype(u);
+                    colorvec, sparsity)
             end
         end
     else
         jac_config = nothing
     end
-    jac_config
+
+    J = if !linsolve_needs_jac
+        # We don't need to construct the Jacobian
+        JacVec(uf, u; autodiff = alg_autodiff(alg) ? AutoForwardDiff() : AutoFiniteDiff())
+    else
+        if f.jac_prototype === nothing
+            ArrayInterface.undefmatrix(u)
+        else
+            f.jac_prototype
+        end
+    end
+
+    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{

src/levenberg.jl

@@ -226,27 +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)
-    J = ArrayInterface.undefmatrix(u)
-
-    linprob = LinearProblem(J, _vec(zero(u)); u0 = _vec(zero(u)))
-    weight = similar(u)
-    recursivefill!(weight, false)
-
-    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 = Pl, Pr = Pr)
-
-    du1 = zero(u)
-    du2 = zero(u)
-    tmp = zero(u)
-    jac_config = build_jac_config(alg, f, uf, du1, u, tmp, du2)
-
-    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,31 +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)
-    J = if f.jac_prototype === nothing
-        ArrayInterface.undefmatrix(u)
-    else
-        f.jac_prototype
-    end
-
-    linprob = LinearProblem(J, _vec(zero(u)); u0 = _vec(zero(u)))
-    weight = similar(u)
-    recursivefill!(weight, false)
-
-    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 = Pl, Pr = Pr)
-
-    du1 = zero(u)
-    du2 = zero(u)
-    tmp = zero(u)
-    jac_config = build_jac_config(alg, f, uf, du1, u, tmp, du2)
-
-    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,27 +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)
-    J = ArrayInterface.undefmatrix(u)
-
-    linprob = LinearProblem(J, _vec(zero(u)); u0 = _vec(zero(u)))
-    weight = similar(u)
-    recursivefill!(weight, false)
-
-    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 = Pl, Pr = Pr)
-
-    du1 = zero(u)
-    du2 = zero(u)
-    tmp = zero(u)
-    jac_config = build_jac_config(alg, f, uf, du1, u, tmp, du2)
-
-    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

src/utils.jl

@@ -51,8 +51,7 @@ function alg_difftype(alg::AbstractNewtonAlgorithm{
     FDT,
     ST,
     CJ,
-}) where {CS, AD, FDT,
-    ST, CJ}
+}) where {CS, AD, FDT, ST, CJ}
     FDT
 end
 
@@ -62,8 +61,7 @@ function concrete_jac(alg::AbstractNewtonAlgorithm{
     FDT,
     ST,
     CJ,
-}) where {CS, AD, FDT,
-    ST, CJ}
+}) where {CS, AD, FDT, ST, CJ}
     CJ
 end
 
@@ -73,9 +71,7 @@ function get_chunksize(alg::AbstractNewtonAlgorithm{
     FDT,
     ST,
     CJ,
-}) where {CS, AD,
-    FDT,
-    ST, CJ}
+}) where {CS, AD, FDT, ST, CJ}
     Val(CS)
 end
 
@@ -85,17 +81,15 @@ function standardtag(alg::AbstractNewtonAlgorithm{
     FDT,
     ST,
     CJ,
-}) where {CS, AD, FDT,
-    ST, CJ}
+}) where {CS, AD, FDT, ST, CJ}
     ST
 end
 
 DEFAULT_PRECS(W, du, u, p, t, newW, Plprev, Prprev, cachedata) = nothing, nothing
 
 function dolinsolve(precs::P, linsolve; A = nothing, linu = nothing, b = nothing,
-    du = nothing, u = nothing, p = nothing, t = nothing,
-    weight = nothing, cachedata = nothing,
-    reltol = nothing) where {P}
+    du = nothing, u = nothing, p = nothing, t = nothing, weight = nothing,
+    cachedata = nothing, reltol = nothing) where {P}
     A !== nothing && (linsolve.A = A)
     b !== nothing && (linsolve.b = b)
     linu !== nothing && (linsolve.u = linu)