DiffEq and SVector

Hello,

I’ve been rereading some of the docs and tutorials in SciML and in particular this https://tutorials.sciml.ai/html/introduction/03-optimizing_diffeq_code.html, and tried using StaticArrays as in the lorenz example, but on a problem of my own. In my case, the execution time went from about 70 ms using the modifying form with normal vectors, to about 300 ms using SVector. My problem has only 4 ODEs, and is needing about 22000 points to solve (mostly driven by a PeriodicCallback that fires 6000 times).

The solutions appear to be the same, but I do see a difference in the number of points saved, 22241 (modifying form) vs 22533 (SVector) which seems suspicious.

I’m not specifying the solver, but both report CompositeAlgorithm of Vern7, Rodas4 in the sol.alg field.

Due to proprietary stuff in my file I can’t just post it, so before I try to make an MWE I’ll try asking a couple of questions:

Have changes in Julia made the SVector optimization less valuable or possibly counterproductive?

Does the heavy use of PeriodicCallback do something that would cause this to backfire?

I’m using Julia 1.5.3 on Windows.

EDIT: The time point sets are generally different (other than those forced by the DiscreteCallback). I confirmed by all(isapprox(sol(t),sol2(t);atol=1e-6) for t ∈ 0.0:0.001:0.6) that the answers are equivalent to within the abstol I specified, but why would the change of mutating or not cause the simulation to take a (slightly) different path?

Thanks

Does SVector work for visualization?

I’m sorry, I don’t understand your question.

I am plotting the solution in either case, but I don’t see why changing the form of the ODE (modifying vector vs returning static vector) would affect the types in the solution object.

nevermind, I was confused by the term SVector. Thinking of .svg

Theres a lot of information that I can’t glean from your post but that could have performance implications.

How large are the SVectors? Beyond a certain size they perform less efficiently then Array.

Also how type stable is the code? If you put SVectors into containers or use them as feilds in structs it can be easy to end up in a type unstable situation, for example:

# This is an array type which's eltype is not concrete, and so requires dispatch when you try to do things with elements from it
Array{SVector{3,Float64}}


# This should be used instead
Array{SArray{Tuple{3}, Float64, 1, 3}}

# A different way to do the above which may be easier.
const Float64x3 = typeof(SVector(0,0,0))
Array{Float64x3}

As usual in Julia, the best way to deal with types is to not mention them, when possible:

julia> using StaticArrays

julia> x = SVector(1, 2, 3);

julia> typeof(x)
SVector{3, Int64} (alias for SArray{Tuple{3}, Int64, 1, 3})

julia> x
3-element SVector{3, Int64} with indices SOneTo(3):
 1
 2
 3

julia> v = [x];   # make a Vector containing just `x`

julia> typeof(v)
Vector{SVector{3, Int64}} (alias for Array{SArray{Tuple{3}, Int64, 1, 3}, 1})

(The nice printing is from Julia 1.6.)

2 Likes

Floating point is non-associative and SIMD changes the results, so SVector giving more SIMD can definitely lead to a different floating point result.

I think to understand what’s going on here I’d need to get an MWE. This seems really in the weeds and not a high level issue.

1 Like

It’s a rather simple system - 4 scalar ODEs, so the SVector is 4 elements long.

I was worried about how the @SVector macro knew about types, so I switched from the macro form and declared them explicitly to be Float64s and it performed the same. I think the code should be type stable, u0, p, and t are all Float64s and the ODEs are pure math. The ODEs are fairly simple, the complex code is in my callback which has a bunch of rounding, clamping, stuff done in smaller integer sizes, but then the result of the callback is always a Float64 that gets stuffed into one of the parameters. That should be irrelevant to the SVectors, I’ll cut it out and do more comparisons tomorrow.

Ah, that makes sense. I’ll trim it down tomorrow.

The @SVector macro should be perfectly fine to use.

julia> using StaticArrays

julia> isconcretetype(SVector{3,Float64})
true

So the eltype of Array{SVector{3,Float64}} is concrete, although Array{SVector{3,Float64}} itself isn’t, because it’s missing the dimensionality of the array.

julia> isconcretetype(Array{SVector{3,Float64}})
false

julia> isconcretetype(Array{SVector{3,Float64},5})
true
2 Likes

MWE:

### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ 88d20350-3af7-11eb-0983-791f4579d925
begin
	using DifferentialEquations
	using StaticArrays
end

# ╔═╡ f2ce46b0-3af7-11eb-2a6a-4716734a17d5
function sim(odef)
	Δt = 0.1e-3
    u0 = [0.0, 0.0, 0.0, 0.0]
	p = [0.178,19.0,10.0,0.25,0.25,20.0,140e3,1.37e3,1.00e3,0.1e-6,0.5,10*1000e-12]
    tspan = (0.0, 0.1)
    prob = ODEProblem(odef, u0, tspan, p)
    sol = solve(prob, reltol=1e-6, abstol=1e-6)
end

# ╔═╡ e6f9b200-3afe-11eb-3df1-d5a083f77901
function ode(u,p,t)
    i, Vsa, Vcfb, Vrc = u
	Lv,Rv,Vv,Rs,Rds,Ga,fc,Ri,Rfb,Cfb,Vdiv,τrc = p
    di = 1/Lv*(Vv-u[1]*(Rv+Rs+Rds))
    dVsa = 2π*fc*(i*Rs*Ga-Vsa)
	Vv    = Vdiv*Vsa
	dVcfb = 1/(Cfb*Rfb)*(Vv*Rfb/Ri-Vcfb)
	dVrc = 1/τrc*(Vcfb+Vv-Vrc)
	@SVector [di, dVsa, dVcfb, dVrc]
end

# ╔═╡ 0a24afa0-3aff-11eb-2a88-89b666495af8
@time sol_sv=sim(ode);

# ╔═╡ ec697740-3af7-11eb-2c48-f9c34e498151
function ode!(du,u,p,t)
    i, Vsa, Vcfb, Vrc = u
	Lv,Rv,Vv,Rs,Rds,Ga,fc,Ri,Rfb,Cfb,Vdiv,τrc = p
    di = 1/Lv*(Vv-u[1]*(Rv+Rs+Rds))
    dVsa = 2π*fc*(i*Rs*Ga-Vsa)
	Vv    = Vdiv*Vsa
	dVcfb = 1/(Cfb*Rfb)*(Vv*Rfb/Ri-Vcfb)
	dVrc = 1/τrc*(Vcfb+Vv-Vrc)
	du .= [di, dVsa, dVcfb, dVrc]
end

# ╔═╡ 30761000-3afe-11eb-028c-9732640ebe6f
@time sol=sim(ode!);

# ╔═╡ 36b22c50-3aff-11eb-2120-efdff8eada1c
all(isapprox(sol(t),sol_sv(t);atol=1e-6) for t ∈ 0.0:0.001:sol.t[end])

# ╔═╡ Cell order:
# ╠═36b22c50-3aff-11eb-2120-efdff8eada1c
# ╠═30761000-3afe-11eb-028c-9732640ebe6f
# ╠═0a24afa0-3aff-11eb-2a88-89b666495af8
# ╠═f2ce46b0-3af7-11eb-2a6a-4716734a17d5
# ╠═e6f9b200-3afe-11eb-3df1-d5a083f77901
# ╠═ec697740-3af7-11eb-2c48-f9c34e498151
# ╠═88d20350-3af7-11eb-0983-791f4579d925

I’ve been testing within Pluto, but just in case there’s some odd interaction, I included this file at the repl and got similar timings:

julia> using BenchmarkTools

julia> @btime sim(ode!);
  814.061 μs (5164 allocations: 362.88 KiB)

julia> @btime sim(ode);
  2.698 ms (28782 allocations: 1.20 MiB)

ode! can be easily improved too by doing du .= (di, dVsa, dVcfb, dVrc) so as to not allocate the array.

1 Like

Your initial condition is not a static vector here.

New MWE with improved ode! and using @SVector for u0.

### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ 88d20350-3af7-11eb-0983-791f4579d925
begin
	using DifferentialEquations
	using StaticArrays
end

# ╔═╡ eee69550-3b30-11eb-3592-076654279c24
function sim_sv(odef)
	Δt = 0.1e-3
    u0 = @SVector[0.0, 0.0, 0.0, 0.0]
	p = [0.178,19.0,10.0,0.25,0.25,20.0,140e3,1.37e3,1.00e3,0.1e-6,0.5,10*1000e-12]
    tspan = (0.0, 0.1)
    prob = ODEProblem(odef, u0, tspan, p)
    sol = solve(prob, reltol=1e-6, abstol=1e-6)
end

# ╔═╡ f2ce46b0-3af7-11eb-2a6a-4716734a17d5
function sim(odef)
	Δt = 0.1e-3
    u0 = [0.0, 0.0, 0.0, 0.0]
	p = [0.178,19.0,10.0,0.25,0.25,20.0,140e3,1.37e3,1.00e3,0.1e-6,0.5,10*1000e-12]
    tspan = (0.0, 0.1)
    prob = ODEProblem(odef, u0, tspan, p)
    sol = solve(prob, reltol=1e-6, abstol=1e-6)
end

# ╔═╡ e6f9b200-3afe-11eb-3df1-d5a083f77901
function ode(u,p,t)
    i, Vsa, Vcfb, Vrc = u
	Lv,Rv,Vv,Rs,Rds,Ga,fc,Ri,Rfb,Cfb,Vdiv,τrc = p
    di = 1/Lv*(Vv-u[1]*(Rv+Rs+Rds))
    dVsa = 2π*fc*(i*Rs*Ga-Vsa)
	Vv    = Vdiv*Vsa
	dVcfb = 1/(Cfb*Rfb)*(Vv*Rfb/Ri-Vcfb)
	dVrc = 1/τrc*(Vcfb+Vv-Vrc)
	@SVector [di, dVsa, dVcfb, dVrc]
end

# ╔═╡ 0a24afa0-3aff-11eb-2a88-89b666495af8
@time sol_sv=sim_sv(ode);

# ╔═╡ ec697740-3af7-11eb-2c48-f9c34e498151
function ode!(du,u,p,t)
    i, Vsa, Vcfb, Vrc = u
	Lv,Rv,Vv,Rs,Rds,Ga,fc,Ri,Rfb,Cfb,Vdiv,τrc = p
    di = 1/Lv*(Vv-u[1]*(Rv+Rs+Rds))
    dVsa = 2π*fc*(i*Rs*Ga-Vsa)
	Vv    = Vdiv*Vsa
	dVcfb = 1/(Cfb*Rfb)*(Vv*Rfb/Ri-Vcfb)
	dVrc = 1/τrc*(Vcfb+Vv-Vrc)
	du .= (di, dVsa, dVcfb, dVrc)
end

# ╔═╡ 30761000-3afe-11eb-028c-9732640ebe6f
@time sol=sim(ode!);

# ╔═╡ 36b22c50-3aff-11eb-2120-efdff8eada1c
all(isapprox(sol(t),sol_sv(t);atol=1e-6) for t ∈ 0.0:0.001:sol.t[end])

# ╔═╡ Cell order:
# ╠═36b22c50-3aff-11eb-2120-efdff8eada1c
# ╠═30761000-3afe-11eb-028c-9732640ebe6f
# ╠═0a24afa0-3aff-11eb-2a88-89b666495af8
# ╠═eee69550-3b30-11eb-3592-076654279c24
# ╠═f2ce46b0-3af7-11eb-2a6a-4716734a17d5
# ╠═e6f9b200-3afe-11eb-3df1-d5a083f77901
# ╠═ec697740-3af7-11eb-2c48-f9c34e498151
# ╠═88d20350-3af7-11eb-0983-791f4579d925

and results at REPL (ode! allocations reduced by replacing array w/ tuple):

julia> include("svector_mwe.jl")
 13.307219 seconds (37.73 M allocations: 1.982 GiB, 5.72% gc time)
 23.287191 seconds (70.86 M allocations: 3.400 GiB, 5.05% gc time)
true

julia> using BenchmarkTools

julia> @btime sim(ode!);
  525.576 μs (3527 allocations: 183.83 KiB)

julia> @btime sim_sv(ode);
  2.454 ms (28168 allocations: 1.15 MiB)

Because you’re using the default algorithm, autodiff is turned off and thus you’re hitting https://github.com/JuliaDiff/FiniteDiff.jl/issues/120 . If you specify Rodas5() as the algorithm (since the equations are stiff and do the switch), then you get ~200us which is the fastest.

2 Likes

Thank you!

I’ve applied the learnings to my original problem. Just for the record, I started this thread with what I thought was a fine solution needing 70 ms, tried to make it faster (which gave me 300 ms), then with changing algorithm and removing an array allocation got to 30 ms, and then adding SVector like I was trying to do am now down to 14 ms!

6 Likes