Merge pull request #105 from SciML/sparsedifftools_support

Setup with SparseDiffTools

Author: ChrisRackauckas

Project.toml

@@ -15,6 +15,8 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 SnoopPrecompile = "66db9d55-30c0-4569-8b51-7e840670fc0c"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 
@@ -39,7 +41,8 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 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"]
+test = ["BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics"]

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

@@ -1,68 +1,144 @@
-mutable struct JacobianWrapper{fType, pType}
+struct JacobianWrapper{fType, pType}
     f::fType
     p::pType
 end
 
 (uf::JacobianWrapper)(u) = uf.f(u, uf.p)
 (uf::JacobianWrapper)(res, u) = uf.f(res, u, uf.p)
 
-struct ImmutableJacobianWrapper{fType, pType}
-    f::fType
-    p::pType
-end
+struct NonlinearSolveTag 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);
+     maximum(jac_config.colorvec))
+end
+function jacobian_finitediff!(J, f, x, jac_config, cache)
+    (FiniteDiff.finite_difference_jacobian!(J, f, x, jac_config);
+     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}, cache)
+    f = cache.f
+    uf = cache.uf
+    x = cache.u
+    fx = cache.fu
+    jac_config = cache.jac_config
+    alg = cache.alg
+
+    if alg_autodiff(alg)
+        forwarddiff_color_jacobian!(J, uf, 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
+            uf(fx, x)
+            #cache.destats.nf += 1
+            tmp = jacobian_finitediff_forward!(J, uf, x, jac_config, fx,
+                                               cache)
+        else # not forward difference
+            tmp = jacobian_finitediff!(J, uf, 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(NonlinearSolveTag(), 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_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(cache, f::F) where {F}
+    x = cache.u
+    alg = cache.alg
+    uf = cache.uf
+    local tmp
 
-function jacobian!(J, f, x, fx, solver, jac_config)
-    alg = solver.alg
-    if alg_autodiff(alg)
-        ForwardDiff.jacobian!(J, f, fx, x, jac_config)
+    if DiffEqBase.has_jac(cache.f)
+        J = f.jac(cache.u, cache.p)
+    elseif alg_autodiff(alg)
+        J, tmp = jacobian_autodiff(uf, 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 = true
+        J, tmp = jacobian_finitediff(uf, x, alg_difftype(alg), dir, colorvec, sparsity,
+                                     jac_prototype)
     end
-    nothing
+    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 /
+                             SparseDiffTools.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

@@ -66,24 +66,15 @@ function jacobian_caches(alg::NewtonRaphson, f, u, p, ::Val{true})
                     Pl = Pl, Pr = Pr)
 
     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
 
 function jacobian_caches(alg::NewtonRaphson, f, u, p, ::Val{false})
-    nothing, nothing, nothing, nothing, nothing
+    JacobianWrapper(f, p), nothing, nothing, nothing, nothing
 end
 
 function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::NewtonRaphson,
@@ -116,10 +107,10 @@ end
 function perform_step!(cache::NewtonRaphsonCache{true})
     @unpack u, fu, f, p, alg = cache
     @unpack J, linsolve, du1 = cache
-    calc_J!(J, cache, cache)
+    jacobian!(J, cache)
 
     # u = u - J \ fu
-    linres = dolinsolve(alg.precs, linsolve, A = J, b = fu, linu = du1,
+    linres = dolinsolve(alg.precs, linsolve, A = J, b = _vec(fu), linu = _vec(du1),
                         p = p, reltol = cache.abstol)
     cache.linsolve = linres.cache
     @. u = u - du1
@@ -133,10 +124,9 @@ end
 
 function perform_step!(cache::NewtonRaphsonCache{false})
     @unpack u, fu, f, p = cache
-    J = calc_J(cache, ImmutableJacobianWrapper(f, p))
+    J = jacobian(cache, f)
     cache.u = u - J \ fu
-    fu = f(cache.u, p)
-    cache.fu = fu
+    cache.fu = f(cache.u, p)
     if iszero(cache.fu) || cache.internalnorm(cache.fu) < cache.abstol
         cache.force_stop = true
     end

src/utils.jl

@@ -48,6 +48,22 @@ function alg_difftype(alg::AbstractNewtonAlgorithm{CS, AD, FDT, ST, CJ}) where {
     FDT
 end
 
+function concrete_jac(alg::AbstractNewtonAlgorithm{CS, AD, FDT, ST, CJ}) where {CS, AD, FDT,
+                                                                                ST, CJ}
+    CJ
+end
+
+function get_chunksize(alg::AbstractNewtonAlgorithm{CS, AD, FDT, ST, CJ}) where {CS, AD,
+                                                                                 FDT,
+                                                                                 ST, CJ}
+    Val(CS)
+end
+
+function standardtag(alg::AbstractNewtonAlgorithm{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,
@@ -97,3 +113,14 @@ function wrapprecs(_Pl, _Pr, weight)
     end
     Pl, Pr
 end
+
+function _nfcount(N, ::Type{diff_type}) where {diff_type}
+    if diff_type === Val{:complex}
+        tmp = N
+    elseif diff_type === Val{:forward}
+        tmp = N + 1
+    else
+        tmp = 2N
+    end
+    tmp
+end

test/runtests.jl

@@ -9,4 +9,5 @@ const is_APPVEYOR = Sys.iswindows() && haskey(ENV, "APPVEYOR")
 
 if GROUP == "All" || GROUP == "Core"
     @time @safetestset "Basic Tests + Some AD" begin include("basictests.jl") end
+    @time @safetestset "Sparsity Tests" begin include("sparse.jl") end
 end end

test/sparse.jl

@@ -0,0 +1,54 @@
+using NonlinearSolve, LinearAlgebra, SparseArrays, Symbolics
+
+const N = 32
+const xyd_brusselator = range(0, stop = 1, length = N)
+brusselator_f(x, y) = (((x - 0.3)^2 + (y - 0.6)^2) <= 0.1^2) * 5.0
+limit(a, N) = a == N + 1 ? 1 : a == 0 ? N : a
+function brusselator_2d_loop(du, u, p)
+    A, B, alpha, dx = p
+    alpha = alpha / dx^2
+    @inbounds for I in CartesianIndices((N, N))
+        i, j = Tuple(I)
+        x, y = xyd_brusselator[I[1]], xyd_brusselator[I[2]]
+        ip1, im1, jp1, jm1 = limit(i + 1, N), limit(i - 1, N), limit(j + 1, N),
+                             limit(j - 1, N)
+        du[i, j, 1] = alpha * (u[im1, j, 1] + u[ip1, j, 1] + u[i, jp1, 1] + u[i, jm1, 1] -
+                       4u[i, j, 1]) +
+                      B + u[i, j, 1]^2 * u[i, j, 2] - (A + 1) * u[i, j, 1] +
+                      brusselator_f(x, y)
+        du[i, j, 2] = alpha * (u[im1, j, 2] + u[ip1, j, 2] + u[i, jp1, 2] + u[i, jm1, 2] -
+                       4u[i, j, 2]) +
+                      A * u[i, j, 1] - u[i, j, 1]^2 * u[i, j, 2]
+    end
+end
+p = (3.4, 1.0, 10.0, step(xyd_brusselator))
+
+function init_brusselator_2d(xyd)
+    N = length(xyd)
+    u = zeros(N, N, 2)
+    for I in CartesianIndices((N, N))
+        x = xyd[I[1]]
+        y = xyd[I[2]]
+        u[I, 1] = 22 * (y * (1 - y))^(3 / 2)
+        u[I, 2] = 27 * (x * (1 - x))^(3 / 2)
+    end
+    u
+end
+u0 = init_brusselator_2d(xyd_brusselator)
+prob_brusselator_2d = NonlinearProblem(brusselator_2d_loop, u0, p)
+sol = solve(prob_brusselator_2d, NewtonRaphson())
+
+du0 = copy(u0)
+jac_sparsity = Symbolics.jacobian_sparsity((du, u) -> brusselator_2d_loop(du, u, p), du0,
+                                           u0)
+
+ff = NonlinearFunction(brusselator_2d_loop; jac_prototype = float.(jac_sparsity))
+prob_brusselator_2d = NonlinearProblem(ff, u0, p)
+sol = solve(prob_brusselator_2d, NewtonRaphson())
+@test norm(sol.resid) < 1e-8
+
+sol = solve(prob_brusselator_2d, NewtonRaphson(autodiff = false))
+@test norm(sol.resid) < 1e-6
+
+cache = init(prob_brusselator_2d, NewtonRaphson())
+@test maximum(cache.jac_config.colorvec) == 12