JuliaDiff / SparseDiffTools.jl

Fast jacobian computation through sparsity exploitation and matrix coloring
MIT License
237 stars 41 forks source link

Specialize default chunk size on static arrays #143

Closed ChrisRackauckas closed 3 years ago

ChrisRackauckas commented 3 years ago

Since the size is static and known to be sufficiently small, we might as well directly specialize on it. This does so, and it reduces the allocations and noticeably improves the static array ODE solver speed.

using OrdinaryDiffEq, BenchmarkTools
using StaticArrays, LinearAlgebra

# rober
function rober(u,p,t)
    y₁,y₂,y₃ = u
    k₁,k₂,k₃ = p
    du1 = -k₁*y₁+k₃*y₂*y₃
    du2 =  k₁*y₁-k₂*y₂^2-k₃*y₂*y₃
    du3 =  k₂*y₂^2
    SA[du1,du2,du3]
end
ff = ODEFunction(rober,tgrad = (u,p,t)->SA[0.0,0.0,0.0])
prob = ODEProblem{false}(rober,SA[1.0,0.0,0.0],(0.0,1e5),SA[0.04,3e7,1e4])

# Before
@btime solve(prob,Rodas5(),reltol=1.0e-8,abstol=1.0e-8, saveat = 1000) # 122.500 μs (581 allocations: 53.39 KiB)
# After
@btime solve(prob,Rodas5(),reltol=1.0e-8,abstol=1.0e-8, saveat = 1000) # 101.600 μs (32 allocations: 10.50 KiB)

using ModelingToolkit
prob2 = ODEProblem{false}(modelingtoolkitize(prob),SA[1.0,0.0,0.0],(0.0,1e5),SA[0.04,3e7,1e4],jac=true)
@btime solve(prob2,Rodas5(), reltol=1.0e-8, abstol=1.0e-8, saveat = 1000) # 75.000 μs (401 allocations: 51.38 KiB)

It's still better to modelingtoolkitize, but this is still a nice victory.