Setup with SparseDiffTools

Author: ChrisRackauckas

Project.toml

@@ -15,6 +15,7 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 SnoopPrecompile = "66db9d55-30c0-4569-8b51-7e840670fc0c"
+SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 

docs/src/solvers/NonlinearSystemSolvers.md

@@ -29,10 +29,18 @@ then `NLSolveJL`'s `:anderson` can be a good choice.
 
 These are the core solvers.
 
-- `NewtonRaphson(;autodiff=true,chunk_size=12,diff_type=Val{:forward},linsolve=DEFAULT_LINSOLVE)`:
+- `NewtonRaphson()`:
   A Newton-Raphson method with swappable nonlinear solvers and autodiff methods
   for high performance on large and sparse systems.
 
+#### Details on Controlling NonlinearSolve.jl Solvers
+
+```julia
+NewtonRaphson(; chunk_size = Val{0}(), autodiff = Val{true}(),
+                standardtag = Val{true}(), concrete_jac = nothing,
+                diff_type = Val{:forward}, linsolve = nothing, precs = DEFAULT_PRECS)
+```
+
 ### SimpleNonlinearSolve.jl
 
 These methods are included with NonlinearSolve.jl by default, though SimpleNonlinearSolve.jl

src/NonlinearSolve.jl

@@ -10,6 +10,7 @@ using RecursiveArrayTools
 import ArrayInterfaceCore
 import LinearSolve
 using DiffEqBase
+using SparseDiffTools
 
 @reexport using SciMLBase
 @reexport using SimpleNonlinearSolve

src/jacobian.jl

@@ -13,56 +13,131 @@ end
 
 (uf::ImmutableJacobianWrapper)(u) = uf.f(u, uf.p)
 
-function calc_J(solver, cache)
-    @unpack u, f, p, alg = solver
-    @unpack uf = cache
-    uf.f = f
-    uf.p = p
-    J = jacobian(uf, u, solver)
-    return J
+function sparsity_colorvec(f, x)
+    sparsity = f.sparsity
+    colorvec = DiffEqBase.has_colorvec(f) ? f.colorvec :
+               (isnothing(sparsity) ? (1:length(x)) : matrix_colors(sparsity))
+    sparsity, colorvec
 end
 
-function calc_J(solver, uf::ImmutableJacobianWrapper)
-    @unpack u, f, p, alg = solver
-    J = jacobian(uf, u, solver)
-    return J
+function jacobian_finitediff_forward!(J, f, x, jac_config, forwardcache, cache)
+    (FiniteDiff.finite_difference_jacobian!(J, f, x, jac_config, forwardcache,
+                                            dir = diffdir(cache));
+     maximum(jac_config.colorvec))
+end
+function jacobian_finitediff!(J, f, x, jac_config, cache)
+    (FiniteDiff.finite_difference_jacobian!(J, f, x, jac_config,
+                                            dir = diffdir(cache));
+     2 * maximum(jac_config.colorvec))
 end
 
-function jacobian(f, x::Number, solver)
-    if alg_autodiff(solver.alg)
-        J = ForwardDiff.derivative(f, x)
+function jacobian!(J::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number},
+                   fx::AbstractArray{<:Number}, cache,
+                   jac_config)
+    alg = unwrap_alg(cache, true)
+    if alg_autodiff(alg)
+        forwarddiff_color_jacobian!(J, f, x, jac_config)
+        #cache.destats.nf += 1
     else
-        J = FiniteDiff.finite_difference_derivative(f, x, alg_difftype(solver.alg),
-                                                    eltype(x))
+        isforward = alg_difftype(alg) === Val{:forward}
+        if isforward
+            forwardcache = get_tmp_cache(cache, alg, unwrap_cache(cache, true))[2]
+            f(forwardcache, x)
+            #cache.destats.nf += 1
+            tmp = jacobian_finitediff_forward!(J, f, x, jac_config, forwardcache,
+                                               cache)
+        else # not forward difference
+            tmp = jacobian_finitediff!(J, f, x, jac_config, cache)
+        end
+        cache.destats.nf += tmp
     end
-    return J
+    nothing
 end
 
-function jacobian(f, x, solver)
-    if alg_autodiff(solver.alg)
-        J = ForwardDiff.jacobian(f, x)
+function build_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
+        sparsity, colorvec = sparsity_colorvec(f, u)
+
+        if alg_autodiff(alg)
+            _chunksize = get_chunksize(alg) === Val(0) ? nothing : get_chunksize(alg) # SparseDiffEq uses different convection...
+
+            T = if standardtag(alg)
+                typeof(ForwardDiff.Tag(OrdinaryDiffEqTag(), eltype(u)))
+            else
+                typeof(ForwardDiff.Tag(uf, eltype(u)))
+            end
+            jac_config = ForwardColorJacCache(uf, u, _chunksize; colorvec = colorvec,
+                                              sparsity = sparsity, tag = T)
+        else
+            if alg_difftype(alg) !== Val{:complex}
+                jac_config = FiniteDiff.JacobianCache(tmp, du1, du2, alg_difftype(alg),
+                                                      colorvec = colorvec,
+                                                      sparsity = sparsity)
+            else
+                jac_config = FiniteDiff.JacobianCache(Complex{eltype(tmp)}.(tmp),
+                                                      Complex{eltype(du1)}.(du1), nothing,
+                                                      alg_difftype(alg), eltype(u),
+                                                      colorvec = colorvec,
+                                                      sparsity = sparsity)
+            end
+        end
     else
-        J = FiniteDiff.finite_difference_jacobian(f, x, alg_difftype(solver.alg), eltype(x))
+        jac_config = nothing
     end
-    return J
+    jac_config
 end
 
-function calc_J!(J, solver, cache)
-    @unpack f, u, p, alg = solver
-    @unpack du1, uf, jac_config = cache
+function get_chunksize(jac_config::ForwardDiff.JacobianConfig{T, V, N, D}) where {T, V, N, D
+                                                                                  }
+    Val(N)
+end # don't degrade compile time information to runtime information
 
-    uf.f = f
-    uf.p = p
-
-    jacobian!(J, uf, u, du1, solver, jac_config)
+function jacobian_finitediff(f, x, ::Type{diff_type}, dir, colorvec, sparsity,
+                             jac_prototype) where {diff_type}
+    (FiniteDiff.finite_difference_derivative(f, x, diff_type, eltype(x), dir = dir), 2)
 end
-
-function jacobian!(J, f, x, fx, solver, jac_config)
-    alg = solver.alg
+function jacobian_finitediff(f, x::AbstractArray, ::Type{diff_type}, dir, colorvec,
+                             sparsity, jac_prototype) where {diff_type}
+    f_in = diff_type === Val{:forward} ? f(x) : similar(x)
+    ret_eltype = eltype(f_in)
+    J = FiniteDiff.finite_difference_jacobian(f, x, diff_type, ret_eltype, f_in,
+                                              dir = dir, colorvec = colorvec,
+                                              sparsity = sparsity,
+                                              jac_prototype = jac_prototype)
+    return J, _nfcount(maximum(colorvec), diff_type)
+end
+function jacobian(f, x, cache)
+    alg = unwrap_alg(cache, true)
+    local tmp
     if alg_autodiff(alg)
-        ForwardDiff.jacobian!(J, f, fx, x, jac_config)
+        J, tmp = jacobian_autodiff(f, x, cache.f, alg)
     else
-        FiniteDiff.finite_difference_jacobian!(J, f, x, jac_config)
+        jac_prototype = cache.f.jac_prototype
+        sparsity, colorvec = sparsity_colorvec(cache.f, x)
+        dir = diffdir(cache)
+        J, tmp = jacobian_finitediff(f, x, alg_difftype(alg), dir, colorvec, sparsity,
+                                     jac_prototype)
     end
-    nothing
+    cache.destats.nf += tmp
+    J
+end
+
+jacobian_autodiff(f, x, nonlinfun, alg) = (ForwardDiff.derivative(f, x), 1, alg)
+function jacobian_autodiff(f, x::AbstractArray, nonlinfun, alg)
+    jac_prototype = nonlinfun.jac_prototype
+    sparsity, colorvec = sparsity_colorvec(nonlinfun, x)
+    maxcolor = maximum(colorvec)
+    chunk_size = get_chunksize(alg) === Val(0) ? nothing : get_chunksize(alg)
+    num_of_chunks = chunk_size === nothing ?
+                    Int(ceil(maxcolor / getsize(ForwardDiff.pickchunksize(maxcolor)))) :
+                    Int(ceil(maxcolor / _unwrap_val(chunk_size)))
+    (forwarddiff_color_jacobian(f, x, colorvec = colorvec, sparsity = sparsity,
+                                jac_prototype = jac_prototype, chunksize = chunk_size),
+     num_of_chunks)
 end

src/raphson.jl

@@ -65,20 +65,10 @@ function jacobian_caches(alg::NewtonRaphson, f, u, p, ::Val{true})
     linsolve = init(linprob, alg.linsolve, alias_A = true, alias_b = true,
                     Pl = Pl, Pr = Pr)
 
-    du1 = zero(u)
+    du1 = zero(u); du2 = zero(u)
     tmp = zero(u)
-    if alg_autodiff(alg)
-        jac_config = ForwardDiff.JacobianConfig(uf, du1, u)
-    else
-        if alg.diff_type != Val{:complex}
-            du2 = zero(u)
-            jac_config = FiniteDiff.JacobianCache(tmp, du1, du2, alg.diff_type)
-        else
-            jac_config = FiniteDiff.JacobianCache(Complex{eltype(tmp)}.(tmp),
-                                                  Complex{eltype(du1)}.(du1), nothing,
-                                                  alg.diff_type, eltype(u))
-        end
-    end
+    jac_config = build_jac_config(alg, f, uf, du1, u, tmp, du2)
+
     uf, linsolve, J, du1, jac_config
 end