Skip to content
This repository was archived by the owner on Aug 22, 2025. It is now read-only.

Commit 8fa2d05

Browse files
Specialize default chunk size on static arrays
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. ```julia 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.
1 parent db3e916 commit 8fa2d05

File tree

3 files changed

+14
-2
lines changed

3 files changed

+14
-2
lines changed

Project.toml

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
1414
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
1515
Requires = "ae029012-a4dd-5104-9daa-d747884805df"
1616
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
17+
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
1718
VertexSafeGraphs = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f"
1819

1920
[compat]
@@ -25,8 +26,9 @@ FiniteDiff = "2"
2526
ForwardDiff = "0.10"
2627
LightGraphs = "1.3"
2728
Requires = "0.5, 1.0"
29+
StaticArrays = "1"
2830
VertexSafeGraphs = "0.1"
29-
julia = "1.4"
31+
julia = "1.6"
3032

3133
[extras]
3234
BandedMatrices = "aae01518-5342-5314-be14-df237901396f"

src/SparseDiffTools.jl

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ using Adapt
1212
using LinearAlgebra
1313
using SparseArrays, ArrayInterface
1414

15+
import StaticArrays
16+
1517
using ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD
1618
using DataStructures: DisjointSets, find_root!, union!
1719

src/differentiation/compute_jacobian_ad.jl

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,15 @@ end
7070
dx = sparsity === nothing && jac_prototype === nothing ? nothing : copy(x)) #if dx is nothing, we will estimate dx at the cost of a function call
7171

7272
if sparsity === nothing && jac_prototype === nothing
73-
cfg = chunksize === nothing ? ForwardDiff.JacobianConfig(f, x) : ForwardDiff.JacobianConfig(f, x, ForwardDiff.Chunk(getsize(chunksize)))
73+
cfg = if chunksize === nothing
74+
if typeof(x) <: StaticArrays.StaticArray
75+
ForwardDiff.JacobianConfig(f, x, ForwardDiff.Chunk{StaticArrays.Size(vec(x))[1]}())
76+
else
77+
ForwardDiff.JacobianConfig(f, x)
78+
end
79+
else
80+
ForwardDiff.JacobianConfig(f, x, ForwardDiff.Chunk(getsize(chunksize)))
81+
end
7482
return ForwardDiff.jacobian(f, x, cfg)
7583
end
7684
if dx isa Nothing

0 commit comments

Comments
 (0)