Add support for JFNK

Project.toml

@@ -41,11 +41,12 @@ julia = "1.6"
 [extras]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 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"]

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/tutorials/nonlinear.md

@@ -31,3 +31,22 @@ 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(), reltol = 1e-9, maxiters=4)
+sol = solve(probN, NewtonRaphson(; linsolve), reltol = 1e-9, maxiters=4)
+```

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,34 @@ 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 = AutoFiniteDiff())
+    else
+        if f.jac_prototype === nothing
+            ArrayInterface.undefmatrix(u)
+        else
+            f.jac_prototype
+        end
+    end
+
+    return J, jac_config
 end
 
 function get_chunksize(jac_config::ForwardDiff.JacobianConfig{

src/levenberg.jl

@@ -228,21 +228,20 @@ end
 
 function jacobian_caches(alg::LevenbergMarquardt, f, u, p, ::Val{true})
     uf = JacobianWrapper(f, p)
-    J = ArrayInterface.undefmatrix(u)
+
+    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, false)
+    # 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 = Pl, Pr = Pr)
-
-    du1 = zero(u)
-    du2 = zero(u)
-    tmp = zero(u)
-    jac_config = build_jac_config(alg, f, uf, du1, u, tmp, du2)
+    linsolve = init(linprob, alg.linsolve; alias_A = true, alias_b = true, Pl, Pr)
 
     uf, linsolve, J, du1, jac_config
 end

src/raphson.jl

@@ -106,25 +106,20 @@ 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
+
+    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, false)
+    # 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 = Pl, Pr = Pr)
-
-    du1 = zero(u)
-    du2 = zero(u)
-    tmp = zero(u)
-    jac_config = build_jac_config(alg, f, uf, du1, u, tmp, du2)
+    linsolve = init(linprob, alg.linsolve; alias_A = true, alias_b = true, Pl, Pr)
 
     uf, linsolve, J, du1, jac_config
 end

src/trustRegion.jl

@@ -280,21 +280,20 @@ end
 
 function jacobian_caches(alg::TrustRegion, f, u, p, ::Val{true})
     uf = JacobianWrapper(f, p)
-    J = ArrayInterface.undefmatrix(u)
+
+    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, false)
+    # 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 = Pl, Pr = Pr)
-
-    du1 = zero(u)
-    du2 = zero(u)
-    tmp = zero(u)
-    jac_config = build_jac_config(alg, f, uf, du1, u, tmp, du2)
+    linsolve = init(linprob, alg.linsolve; alias_A = true, alias_b = true, Pl, Pr)
 
     uf, linsolve, J, du1, jac_config
 end

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)

test/basictests.jl

@@ -46,9 +46,9 @@ sol = benchmark_scalar(sf, csu0)
 # @test (@ballocated benchmark_mutable(ff, cu0)) < 200
 # @test (@ballocated benchmark_scalar(sf, csu0)) < 400
 
-function benchmark_inplace(f, u0)
+function benchmark_inplace(f, u0, linsolve)
     probN = NonlinearProblem{true}(f, u0)
-    solver = init(probN, NewtonRaphson(), abstol = 1e-9)
+    solver = init(probN, NewtonRaphson(; linsolve), abstol = 1e-9)
     sol = solve!(solver)
 end
 
@@ -57,9 +57,11 @@ function ffiip(du, u, p)
 end
 u0 = [1.0, 1.0]
 
-sol = benchmark_inplace(ffiip, u0)
-@test sol.retcode === ReturnCode.Success
-@test all(abs.(sol.u .* sol.u .- 2) .< 1e-9)
+for linsolve in (nothing, KrylovJL_GMRES())
+    sol = benchmark_inplace(ffiip, u0, linsolve)
+    @test sol.retcode === ReturnCode.Success
+    @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9)
+end
 
 u0 = [1.0, 1.0]
 probN = NonlinearProblem{true}(ffiip, u0)
@@ -209,8 +211,7 @@ end
 
 function benchmark_inplace(f, u0, radius_update_scheme)
     probN = NonlinearProblem{true}(f, u0)
-    solver = init(probN, TrustRegion(radius_update_scheme = radius_update_scheme),
-        abstol = 1e-9)
+    solver = init(probN, TrustRegion(; radius_update_scheme), abstol = 1e-9)
     sol = solve!(solver)
 end