diff --git a/Project.toml b/Project.toml index 8228a0d69..072dc0a1d 100644 --- a/Project.toml +++ b/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"] diff --git a/docs/src/solvers/NonlinearSystemSolvers.md b/docs/src/solvers/NonlinearSystemSolvers.md index 7602c1936..e16a956cf 100644 --- a/docs/src/solvers/NonlinearSystemSolvers.md +++ b/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 diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index d82680481..99fc4d61c 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -10,6 +10,7 @@ using RecursiveArrayTools import ArrayInterfaceCore import LinearSolve using DiffEqBase +using SparseDiffTools @reexport using SciMLBase @reexport using SimpleNonlinearSolve diff --git a/src/jacobian.jl b/src/jacobian.jl index 96c07ada9..e8f6fb2f9 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -1,4 +1,4 @@ -mutable struct JacobianWrapper{fType, pType} +struct JacobianWrapper{fType, pType} f::fType p::pType end @@ -6,63 +6,139 @@ 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 diff --git a/src/raphson.jl b/src/raphson.jl index f5f507dcf..9e7ae6c88 100644 --- a/src/raphson.jl +++ b/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 diff --git a/src/utils.jl b/src/utils.jl index 2bcc6334d..f4946091b 100644 --- a/src/utils.jl +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index 7269ed6bf..37b9957e7 100644 --- a/test/runtests.jl +++ b/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 diff --git a/test/sparse.jl b/test/sparse.jl new file mode 100644 index 000000000..6ec40509f --- /dev/null +++ b/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