Solution indexing is very slow

When running the following code, the last two lines do the same thing: they get decay2.f at the last point in time. sol[decay2.f][end] is the modelingtoolkit preferred way, while sol.u[end][2] is using normal array indexing.

using ModelingToolkit
using ModelingToolkit: t_nounits as t, D_nounits as D

function decay(; name)
    @parameters a
    @variables x(t) f(t)
    ODESystem([
            D(x) ~ -a * x + f
        ], t;
        name = name)
end

@named decay1 = decay()
@named decay2 = decay()

connected = compose(
    ODESystem([decay2.f ~ decay1.x
               D(decay1.f) ~ 0], t; name = :connected), decay1, decay2)

equations(connected)
simplified_sys = structural_simplify(connected)
equations(simplified_sys)

x0 = [decay1.x => 1.0
      decay1.f => 0.0
      decay2.x => 1.0]
p = [decay1.a => 0.1
     decay2.a => 0.2]

using OrdinaryDiffEq
prob = ODEProblem(simplified_sys, x0, (0.0, 100.0), p)
sol = solve(prob, Tsit5())
@time sol[decay2.f][end]
@time sol.u[end][2]

When timing these two lines of code, I get the following results:
0.000734 seconds (1.59 k allocations: 114.008 KiB)
0.000014 seconds (2 allocations: 1024 bytes)

Somehow, sol[decay2.f][end] takes 50 times longer doing the same thing. Are there any possible alternatives or solutions to this? I am developing performance critical simulation code, where this time difference is hurting performance. And using normal array indexing is not safe, as an update of ModelingToolkit could change the order of the solution variables.

You want to cache the symbol lookup function. This is what getu is for. Optimizing through an ODE solve and re-creating MTK Problems · ModelingToolkit.jl is somewhat helpful. Try doing it, and post what you get to, and complain, and @cryptic.ax we should write a proper tutorial on this.

The following code works for me:

using ModelingToolkit
using ModelingToolkit: t_nounits as t, D_nounits as D
using SymbolicIndexingInterface

function decay(; name)
    @parameters a
    @variables x(t) f(t)
    ODESystem([
            D(x) ~ -a * x + f
        ], t;
        name = name)
end

@named decay1 = decay()
@named decay2 = decay()

connected = compose(
    ODESystem([decay2.f ~ decay1.x
               D(decay1.f) ~ 0], t; name = :connected), decay1, decay2)

equations(connected)
simplified_sys = structural_simplify(connected)
equations(simplified_sys)

x0 = [decay1.x => 1.0
      decay1.f => 0.0
      decay2.x => 1.0]
p = [decay1.a => 0.1
     decay2.a => 0.2]

using OrdinaryDiffEq
prob = ODEProblem(simplified_sys, x0, (0.0, 100.0), p)
sol = solve(prob, Tsit5())
get_decay2_f = getu(sol, decay2.f)
@time sol[decay2.f][end]
@time sol.u[end][2]
@time get_decay2_f(sol)[end]

When you get the first solution you can create a function like get_decay2_f(sol) for all variables you are interested in, and later you just use these functions which is fast.

Thanks, that works!
I am getting the following results with Fechners code:

@time sol[decay2.f][end]
@time sol.u[end][2]
@time get_decay2_f(sol)[end]
  0.000141 seconds (276 allocations: 15.031 KiB)
  0.000012 seconds (2 allocations: 1024 bytes)
  0.000028 seconds (147 allocations: 3.453 KiB)

get_decay2_f should be non-allocating. Is that timing its compilation?

If I run my code multiple times I get:

julia> include("mwes/mwe_13b.jl")
  0.000173 seconds (131 allocations: 11.828 KiB)
  0.000017 seconds (2 allocations: 1024 bytes)
  0.000014 seconds (2 allocations: 256 bytes)
4.542765615967835e-5

Two allocations, but neither zero and nor 147.

That’s because get_decay2_f isn’t const. If you created that function as const or used it in a function that global lookup should go away?

Well, it doesn’t go away.

get_decay2_f = getu(sol, decay2.f)

function process_result(sol, get_decay2_f)
    @time sol[decay2.f][end]
    @time sol.u[end][2]
    @time get_decay2_f(sol)[end]
end

process_result(sol)

Output:

julia> process_result(sol)
  0.000019 seconds (132 allocations: 12.812 KiB)
  0.000001 seconds
  0.000001 seconds (3 allocations: 1.234 KiB)
4.542765615967835e-5

But it doesn’t matter. Good enough.

We can open an issue and get that fixed. It’s pretty close though so at least most people will be happy, but embedded will need it fixed.

2 Likes