diff --git a/Project.toml b/Project.toml index 6deb2438b..5aaad06d6 100644 --- a/Project.toml +++ b/Project.toml @@ -21,8 +21,10 @@ RandomNumbers = "e6cf234a-135c-5ec9-84dd-332b85af5143" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" [compat] +SparseDiffTools = ">= 0.3.0" julia = "1" [extras] diff --git a/src/StochasticDiffEq.jl b/src/StochasticDiffEq.jl index 4cf4be749..6e509c434 100644 --- a/src/StochasticDiffEq.jl +++ b/src/StochasticDiffEq.jl @@ -47,6 +47,8 @@ module StochasticDiffEq import DiffEqBase: iip_get_uf, oop_get_uf, build_jac_config + using SparseDiffTools: forwarddiff_color_jacobian!, ForwardColorJacCache + const CompiledFloats = Union{Float32,Float64} diff --git a/src/derivative_wrappers.jl b/src/derivative_wrappers.jl index d553b7297..91a98fed4 100644 --- a/src/derivative_wrappers.jl +++ b/src/derivative_wrappers.jl @@ -5,7 +5,7 @@ function derivative(f, x::Union{Number,AbstractArray{<:Number}}, if get_current_alg_autodiff(integrator.alg, integrator.cache) d = ForwardDiff.derivative(f, x) else - d = DiffEqDiffTools.finite_difference_gradient(f, x, alg.diff_type, eltype(x), Val{false}) + d = DiffEqDiffTools.finite_difference_gradient(f, x, alg.diff_type, eltype(x), Val(false)) end d end @@ -20,45 +20,57 @@ function derivative!(df::AbstractArray{<:Number}, f, x::Union{Number,AbstractArr nothing end +jacobian_autodiff(f,x,_,_)=ForwardDiff.derivative(f,x) +jacobian_autodiff(f,x::AbstractArray,_,hascolor::Val{false})=ForwardDiff.jacobian(f,x) +function jacobian_autodiff(f,x::AbstractArray,integrator,hascolor::Val{true}) + colorvec=integrator.f.colorvec + jac=integrator.f.jac_prototype + J=jac isa SparseMatrixCSC ? similar(jac) : fill(0.,size(jac)) + forwarddiff_color_jacobian!(J,f,x,color=colorvec,sparsity=jac) + J +end +jacobian_finitediff(f,x,difftype,_,_)=DiffEqDiffTools.finite_difference_derivative(f, x, difftype, eltype(x)) +jacobian_finitediff(f,x::AbstractArray,difftype,_,hascolor::Val{false})=DiffEqDiffTools.finite_difference_jacobian(f, x, difftype, eltype(x), Val{false}) +jacobian_finitediff(f,x::AbstractArray,difftype,integrator,hascolor::Val{true})= + DiffEqDiffTools.finite_difference_jacobian(f, x, difftype, eltype(x), Val(false),color=integrator.f.colorvec,sparsity=integrator.f.jac_prototype) + function jacobian(f, x, integrator::DiffEqBase.DEIntegrator) local J alg = unwrap_alg(integrator, true) - isarray = typeof(x) <: AbstractArray if get_current_alg_autodiff(integrator.alg, integrator.cache) - if isarray - J = ForwardDiff.jacobian(f,x) - else - J = ForwardDiff.derivative(f,x) - end + J = jacobian_autodiff(f,x,integrator,Val(DiffEqBase.has_colorvec(integrator.f))) else - if isarray - J = DiffEqDiffTools.finite_difference_jacobian(f, x, alg.diff_type, eltype(x), Val{false}) - else - J = DiffEqDiffTools.finite_difference_derivative(f, x, alg.diff_type, eltype(x)) - end + J = jacobian_finitediff(f,x,alg.diff_type,integrator,Val(DiffEqBase.has_colorvec(integrator.f))) end J end +jacobian_autodiff!(J,f,fx,x,jac_config::ForwardColorJacCache)=forwarddiff_color_jacobian!(J,f,x,jac_config) +jacobian_autodiff!(J,f,fx,x,jac_config::ForwardDiff.JacobianConfig)=ForwardDiff.jacobian!(J, f, fx, x, jac_config) + function jacobian!(J::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number}, fx::AbstractArray{<:Number}, integrator::DEIntegrator, jac_config) if alg_autodiff(integrator.alg) - ForwardDiff.jacobian!(J, f, fx, x, jac_config) + jacobian_autodiff!(J, f, fx, x, jac_config) else DiffEqDiffTools.finite_difference_jacobian!(J, f, x, jac_config) end nothing end +jac_cache_autodiff(alg,f,uf,du1,uprev,u,hascolor::Val{false})=ForwardDiff.JacobianConfig(uf,du1,uprev,ForwardDiff.Chunk{determine_chunksize(u,alg)}()) +jac_cache_autodiff(alg,f,uf,du1,uprev,u,hascolor::Val{true})=ForwardColorJacCache(uf,uprev,color=f.colorvec,sparsity=f.jac_prototype) + function DiffEqBase.build_jac_config(alg::StochasticDiffEqAlgorithm,f,uf,du1,uprev,u,tmp,du2) if !has_jac(f) if alg_autodiff(alg) - jac_config = ForwardDiff.JacobianConfig(uf,du1,uprev,ForwardDiff.Chunk{determine_chunksize(u,alg)}()) + jac_config = jac_cache_autodiff(alg,f,uf,du1,uprev,u,Val(DiffEqBase.has_colorvec(f))) else + colorvec= f.colorvec isa Nothing ? Base.OneTo(length(u)) : f.colorvec if alg.diff_type != Val{:complex} - jac_config = DiffEqDiffTools.JacobianCache(tmp,du1,du2,alg.diff_type) + jac_config = DiffEqDiffTools.JacobianCache(tmp,du1,du2,alg.diff_type,color=colorvec,sparsity=f.jac_prototype) else - jac_config = DiffEqDiffTools.JacobianCache(Complex{eltype(tmp)}.(tmp),Complex{eltype(du1)}.(du1),nothing,alg.diff_type) + jac_config = DiffEqDiffTools.JacobianCache(Complex{eltype(tmp)}.(tmp),Complex{eltype(du1)}.(du1),nothing,alg.diff_type,color=colorvec,sparsity=f.jac_prototype) end end else diff --git a/test/sparsediff_tests.jl b/test/sparsediff_tests.jl new file mode 100644 index 000000000..24ed40b4a --- /dev/null +++ b/test/sparsediff_tests.jl @@ -0,0 +1,60 @@ +using Test +using StochasticDiffEq +using SparseArrays +using LinearAlgebra + +#https://github.com/JuliaDiffEq/SparseDiffTools.jl/blob/master/test/test_integration.jl +function f(dx,x,p,t) + for i in 2:length(x)-1 + dx[i] = x[i-1] - 2x[i] + x[i+1] + end + dx[1] = -2x[1] + x[2] + dx[end] = x[end-1] - 2x[end] + nothing +end + +function g(dx,x,p,t) + dx .= 0 + nothing +end + +function second_derivative_stencil(N) + A = zeros(N,N) + for i in 1:N, j in 1:N + (j-i==-1 || j-i==1) && (A[i,j]=1) + j-i==0 && (A[i,j]=-2) + end + A +end + +function generate_sparsity_pattern(N::Integer) + dl = repeat([1.0],N-1) + du = repeat([1.0],N-1) + d = repeat([-2.0],N) + return Tridiagonal(dl,d,du) +end + +jac_sp = sparse(generate_sparsity_pattern(10)) +#jac = second_derivative_stencil(10) +colors = repeat(1:3,10)[1:10] +u0=[1.,2.,3,4,5,5,4,3,2,1] +tspan=(0.,10.) +sdefun_sp= SDEFunction(f,g,colorvec=colors,jac_prototype=jac_sp) +prob_sp = SDEProblem(sdefun_sp,g,u0,tspan) +prob_std = SDEProblem(f,g,u0,tspan) + +sol_sp=solve(prob_sp,SKenCarp(autodiff=false)) +@test sol_sp.retcode==:Success#test sparse finitediff +sol=solve(prob_std,SKenCarp(autodiff=false)) +@test sol_sp.u[end]≈sol.u[end] atol=1e-4 +@test length(sol_sp.t)==length(sol.t) + +sol_sp=solve(prob_sp,SKenCarp()) +sol=solve(prob_std,SKenCarp()) +@test sol_sp.u[end]≈sol.u[end] atol=1e-4 +@test length(sol_sp.t)==length(sol.t) + +#sol_sp=solve(prob_sp,SKenCarp(autodiff=false),abstol=1e-10,reltol=1e-10) +#sol=solve(prob_std,SKenCarp(autodiff=false),abstol=1e-10,reltol=1e-10) +#@test sol_sp.u[end]≈sol.u[end] +#@test length(sol_sp.t)==length(sol.t) \ No newline at end of file