From c17ac30fa343732a67c9aabff55b8a30bda40726 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Fri, 2 Dec 2022 02:01:50 +0100 Subject: [PATCH 1/6] Setup with SparseDiffTools --- Project.toml | 1 + docs/src/solvers/NonlinearSystemSolvers.md | 10 +- src/NonlinearSolve.jl | 1 + src/jacobian.jl | 145 ++++++++++++++++----- src/raphson.jl | 16 +-- 5 files changed, 124 insertions(+), 49 deletions(-) diff --git a/Project.toml b/Project.toml index 8228a0d69..16ada9033 100644 --- a/Project.toml +++ b/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" 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..6ea6a4224 100644 --- a/src/jacobian.jl +++ b/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 diff --git a/src/raphson.jl b/src/raphson.jl index f5f507dcf..a210b8280 100644 --- a/src/raphson.jl +++ b/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 From 1364deb239388b1543f8e77ea379c27da12b6f49 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Sun, 4 Dec 2022 01:48:07 +0100 Subject: [PATCH 2/6] tests passing --- src/jacobian.jl | 59 +++++++++++++++++++++++++------------------------ src/raphson.jl | 9 ++++---- src/utils.jl | 26 ++++++++++++++++++++++ 3 files changed, 60 insertions(+), 34 deletions(-) diff --git a/src/jacobian.jl b/src/jacobian.jl index 6ea6a4224..7b6275739 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,12 +6,7 @@ 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 - -(uf::ImmutableJacobianWrapper)(u) = uf.f(u, uf.p) +struct NonlinearSolveTag end function sparsity_colorvec(f, x) sparsity = f.sparsity @@ -21,33 +16,35 @@ function sparsity_colorvec(f, x) end function jacobian_finitediff_forward!(J, f, x, jac_config, forwardcache, cache) - (FiniteDiff.finite_difference_jacobian!(J, f, x, jac_config, forwardcache, - dir = diffdir(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, - dir = diffdir(cache)); + (FiniteDiff.finite_difference_jacobian!(J, f, x, jac_config); 2 * maximum(jac_config.colorvec)) end -function jacobian!(J::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number}, - fx::AbstractArray{<:Number}, cache, - jac_config) - alg = unwrap_alg(cache, true) +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, f, x, jac_config) + forwarddiff_color_jacobian!(J, uf, x, jac_config) #cache.destats.nf += 1 else isforward = alg_difftype(alg) === Val{:forward} if isforward forwardcache = get_tmp_cache(cache, alg, unwrap_cache(cache, true))[2] - f(forwardcache, x) + uf(forwardcache, x) #cache.destats.nf += 1 - tmp = jacobian_finitediff_forward!(J, f, x, jac_config, forwardcache, + tmp = jacobian_finitediff_forward!(J, uf, x, jac_config, forwardcache, cache) else # not forward difference - tmp = jacobian_finitediff!(J, f, x, jac_config, cache) + tmp = jacobian_finitediff!(J, uf, x, jac_config, cache) end cache.destats.nf += tmp end @@ -57,7 +54,7 @@ end 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 + 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 @@ -68,7 +65,7 @@ function build_jac_config(alg, f::F1, uf::F2, du1, u, tmp, du2) where {F1, F2} _chunksize = get_chunksize(alg) === Val(0) ? nothing : get_chunksize(alg) # SparseDiffEq uses different convection... T = if standardtag(alg) - typeof(ForwardDiff.Tag(OrdinaryDiffEqTag(), eltype(u))) + typeof(ForwardDiff.Tag(NonlinearSolveTag(), eltype(u))) else typeof(ForwardDiff.Tag(uf, eltype(u))) end @@ -112,19 +109,23 @@ function jacobian_finitediff(f, x::AbstractArray, ::Type{diff_type}, dir, colorv jac_prototype = jac_prototype) return J, _nfcount(maximum(colorvec), diff_type) end -function jacobian(f, x, cache) - alg = unwrap_alg(cache, true) +function jacobian(cache, f::F) where F + x = cache.u + alg = cache.alg + uf = cache.uf local tmp - if alg_autodiff(alg) - J, tmp = jacobian_autodiff(f, x, cache.f, alg) + + 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 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, + dir = true + J, tmp = jacobian_finitediff(uf, x, alg_difftype(alg), dir, colorvec, sparsity, jac_prototype) end - cache.destats.nf += tmp J end @@ -135,7 +136,7 @@ function jacobian_autodiff(f, x::AbstractArray, nonlinfun, alg) 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 / 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), diff --git a/src/raphson.jl b/src/raphson.jl index a210b8280..8190c3e16 100644 --- a/src/raphson.jl +++ b/src/raphson.jl @@ -73,7 +73,7 @@ function jacobian_caches(alg::NewtonRaphson, f, u, p, ::Val{true}) 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, @@ -106,7 +106,7 @@ 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, @@ -123,10 +123,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..e781d1a40 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -48,6 +48,21 @@ 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 +112,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 From 0e2e49f82e769534eeae1f611a456499245d31de Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Sun, 4 Dec 2022 01:52:22 +0100 Subject: [PATCH 3/6] format --- src/jacobian.jl | 5 +++-- src/raphson.jl | 5 +++-- src/utils.jl | 11 ++++++----- 3 files changed, 12 insertions(+), 9 deletions(-) diff --git a/src/jacobian.jl b/src/jacobian.jl index 7b6275739..a8ee4050a 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -109,7 +109,7 @@ function jacobian_finitediff(f, x::AbstractArray, ::Type{diff_type}, dir, colorv jac_prototype = jac_prototype) return J, _nfcount(maximum(colorvec), diff_type) end -function jacobian(cache, f::F) where F +function jacobian(cache, f::F) where {F} x = cache.u alg = cache.alg uf = cache.uf @@ -136,7 +136,8 @@ function jacobian_autodiff(f, x::AbstractArray, nonlinfun, alg) 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 / + 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), diff --git a/src/raphson.jl b/src/raphson.jl index 8190c3e16..327b9a94d 100644 --- a/src/raphson.jl +++ b/src/raphson.jl @@ -65,7 +65,8 @@ 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); du2 = zero(u) + du1 = zero(u) + du2 = zero(u) tmp = zero(u) jac_config = build_jac_config(alg, f, uf, du1, u, tmp, du2) @@ -73,7 +74,7 @@ function jacobian_caches(alg::NewtonRaphson, f, u, p, ::Val{true}) end function jacobian_caches(alg::NewtonRaphson, f, u, p, ::Val{false}) - JacobianWrapper(f,p), nothing, nothing, nothing, nothing + JacobianWrapper(f, p), nothing, nothing, nothing, nothing end function SciMLBase.__init(prob::NonlinearProblem{uType, iip}, alg::NewtonRaphson, diff --git a/src/utils.jl b/src/utils.jl index e781d1a40..f4946091b 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -49,17 +49,18 @@ function alg_difftype(alg::AbstractNewtonAlgorithm{CS, AD, FDT, ST, CJ}) where { end function concrete_jac(alg::AbstractNewtonAlgorithm{CS, AD, FDT, ST, CJ}) where {CS, AD, FDT, - ST, CJ} -CJ + ST, CJ} + CJ end -function get_chunksize(alg::AbstractNewtonAlgorithm{CS, AD, FDT, ST, CJ}) where {CS, AD, FDT, - ST, CJ} +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, CJ} ST end From 8b74b25bd51d7209f774d98de724defc6499fd4a Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Sun, 4 Dec 2022 02:28:31 +0100 Subject: [PATCH 4/6] push new test --- Project.toml | 4 +++- test/runtests.jl | 1 + test/sparse.jl | 43 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 47 insertions(+), 1 deletion(-) create mode 100644 test/sparse.jl diff --git a/Project.toml b/Project.toml index 16ada9033..072dc0a1d 100644 --- a/Project.toml +++ b/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" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed" @@ -40,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/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..24c16408a --- /dev/null +++ b/test/sparse.jl @@ -0,0 +1,43 @@ +using NonlinearSolve, LinearAlgebra, SparseArrays + +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. +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., 10., 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()) + +using Symbolics +du0 = copy(u0) +jac_sparsity = Symbolics.jacobian_sparsity((du,u)->brusselator_2d_loop(du,u,p),du0,u0) + +f = NonlinearFunction(brusselator_2d_loop;jac_prototype=float.(jac_sparsity)) +prob_brusselator_2d = NonlinearProblem(f,u0,p) +sol = solve(prob_brusselator_2d, NewtonRaphson()) From 486bc1e3242cf4781f11795f979a6b3838ced2a1 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Sun, 4 Dec 2022 02:43:57 +0100 Subject: [PATCH 5/6] fix sparse test --- src/jacobian.jl | 7 +++---- src/raphson.jl | 2 +- test/sparse.jl | 14 ++++++++++---- 3 files changed, 14 insertions(+), 9 deletions(-) diff --git a/src/jacobian.jl b/src/jacobian.jl index a8ee4050a..e8f6fb2f9 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -38,15 +38,14 @@ function jacobian!(J::AbstractMatrix{<:Number}, cache) else isforward = alg_difftype(alg) === Val{:forward} if isforward - forwardcache = get_tmp_cache(cache, alg, unwrap_cache(cache, true))[2] - uf(forwardcache, x) + uf(fx, x) #cache.destats.nf += 1 - tmp = jacobian_finitediff_forward!(J, uf, x, jac_config, forwardcache, + 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 + #cache.destats.nf += tmp end nothing end diff --git a/src/raphson.jl b/src/raphson.jl index 327b9a94d..9e7ae6c88 100644 --- a/src/raphson.jl +++ b/src/raphson.jl @@ -110,7 +110,7 @@ function perform_step!(cache::NewtonRaphsonCache{true}) 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 diff --git a/test/sparse.jl b/test/sparse.jl index 24c16408a..22642fa59 100644 --- a/test/sparse.jl +++ b/test/sparse.jl @@ -1,4 +1,4 @@ -using NonlinearSolve, LinearAlgebra, SparseArrays +using NonlinearSolve, LinearAlgebra, SparseArrays, Symbolics const N = 32 const xyd_brusselator = range(0,stop=1,length=N) @@ -34,10 +34,16 @@ u0 = init_brusselator_2d(xyd_brusselator) prob_brusselator_2d = NonlinearProblem(brusselator_2d_loop,u0,p) sol = solve(prob_brusselator_2d, NewtonRaphson()) -using Symbolics du0 = copy(u0) jac_sparsity = Symbolics.jacobian_sparsity((du,u)->brusselator_2d_loop(du,u,p),du0,u0) -f = NonlinearFunction(brusselator_2d_loop;jac_prototype=float.(jac_sparsity)) -prob_brusselator_2d = NonlinearProblem(f,u0,p) +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 From edf3fbce8bc7a8615ae447ba7faceb01e15b6d25 Mon Sep 17 00:00:00 2001 From: Chris Rackauckas Date: Sun, 4 Dec 2022 02:54:43 +0100 Subject: [PATCH 6/6] format --- test/sparse.jl | 61 +++++++++++++++++++++++++++----------------------- 1 file changed, 33 insertions(+), 28 deletions(-) diff --git a/test/sparse.jl b/test/sparse.jl index 22642fa59..6ec40509f 100644 --- a/test/sparse.jl +++ b/test/sparse.jl @@ -1,48 +1,53 @@ 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. -limit(a, N) = a == N+1 ? 1 : a == 0 ? N : a +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] + 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., 10., step(xyd_brusselator)) +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 + 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) +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) +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) +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)) +sol = solve(prob_brusselator_2d, NewtonRaphson(autodiff = false)) @test norm(sol.resid) < 1e-6 cache = init(prob_brusselator_2d, NewtonRaphson())