Performance overhead observed variables ModelingToolkit

Hi all, I am using ModelingToolkit to simulate a system and noticed that there seems to be a serious performance cost associated with using the nice observed variables functionality of ModelingToolkit. Here is a minimal working example of what I mean.

Problem setup:

using DifferentialEquations, ModelingToolkit, BenchmarkTools

@parameters t σ ρ β
@variables x(t) y(t) z(t) v_x(t) v_y(t) v_z(t)
D = Differential(t) 

eqs = [
    v_x ~ σ *(y - x),
    v_y ~ x * (ρ - z) - y,
    v_z ~ x*y - β * z,
    D(x) ~ v_x,
    D(y) ~ v_y,
    D(z) ~ v_z
]

sys = ODESystem(eqs)
sys = structural_simplify(sys) # get rid of v_x, v_y, v_z

ps = Dict(σ => 10.0, ρ => 28.0, β => 2.66)
u0 = Dict(x => 1.0, y => 0.5, z => 0.8)

tspan = (0.0,1.0)
prob = ODEProblem(sys, u0, tspan, ps)
sol = solve(prob, saveat=[0.1]) # more saveat locations makes everything worse

Then when using the observed variables:

f(sol) = sol[x][end] * sol[v_x][end] + sol[y][end] * sol[v_y][end] + sol[z][end] * sol[v_z][end]

@benchmark f(sol)

julia> @benchmark f(sol)
BechmarkTools.Trial: 10000 samples with 1 evaluations.
 Range (min … max):  38.800 μs …  10.927 ms  ┊ GC (min … max): 0.00% … 98.40%
 Time  (median):     46.200 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   54.137 μs ± 149.940 μs  ┊ GC (mean ± σ):  3.79% ±  1.40%

  ▆▆██▇▇▆▆▅▄▄▃▃▄▃▂▂▁        ▁▁▁                                ▂
  ████████████████████▇▇█████████▇▇▇▇▇▇▇█▇▇▇▇▅▇▇▆▆▆▅▆▆▅▅▅▅▄▅▃▆ █
  38.8 μs       Histogram: log(frequency) by time       133 μs <

 Memory estimate: 17.22 KiB, allocs estimate: 235.

Versus when not using them:

vec = [sol[x][end], sol[y][end], sol[z][end], sol[v_x][end], sol[v_y][end], sol[v_z][end]]
g(vec) = vec[1] * vec[4] + vec[2] * vec[5] + vec[3] * vec[6]

@benchmark g(vec)

julia> @benchmark g(vec)
BechmarkTools.Trial: 10000 samples with 979 evaluations.
 Range (min … max):  27.681 ns …  12.044 μs  ┊ GC (min … max): 0.00% … 99.43%
 Time  (median):     39.122 ns               ┊ GC (median):    0.00%
 Time  (mean ± σ):   48.855 ns ± 123.743 ns  ┊ GC (mean ± σ):  2.45% ±  0.99%

  ▇▆▃▅██▅▄▃▂▃▅▄▃▂▁▁▂▂▂▂▂▁    ▁▁▁                               ▂
  █████████████████████████████████▇▇▇█▇▇▇▇▇▆▆▆▆▆▅▆▅▆▆▄▂▄▅▆▄▆▅ █
  27.7 ns       Histogram: log(frequency) by time       164 ns <

 Memory estimate: 16 bytes, allocs estimate: 1.

Is there a better way to access observed variables that get eliminated through structural_simplify? This overhead becomes a problem when using ensemble simulations… Or am I doing something silly?

Open an issue. We should be able to fix this.

1 Like