ModelingToolkit compilation time for many small simulations

Hello all,

I am in a situation where I want to run many small ODEs, to do an ensemble-type of simulation.
I really like using ModelingToolkit, and was trying to use it for this purpose.

However, my code is running really slowly, and I was wondering if I could be doing something smarter, or it there are inherent drawbacks to using ModelingToolkit (specifically, the compilation time).

Take a slightly modified example from its docs here:

using ModelingToolkit, OrdinaryDiffEq

@variables t x(t)   # independent and dependent variables
@parameters τ       # parameters
D = Differential(t) # define an operator for the differentiation w.r.t. time

# your first ODE, consisting of a single equation, the equality indicated by ~
@variables f(t)

trials = 50

@time for ω in 0.1:0.5:5
    for i in 1:trials
        @named fol_variable_f = ODESystem([f ~ sin((ω * rand()) * t), D(x) ~ (f - x) / τ])
        _sys = structural_simplify(fol_variable_f)

        prob = ODEProblem(_sys, [x => 0.0], (0.0, 10.0), [τ => 0.75])
        # parameter `τ` can be assigned a value, but constant `h` cannot
        sol = solve(prob, Tsit5())
    end
end

this gives on my machine:

264.382315 seconds (103.54 M allocations: 6.749 GiB, 2.51% gc time, 97.25% compilation time)

As you can see, this is almost all compilation time.
Is this compilation time inherent to using ModelingToolkit, or is there something I can do to drastically speed things up here?

1 Like

I am absolutely unfamiliar with ModelingToolkit, but it seems very likely to me that the issue is that you work in the global scope. Try wrapping everything in a function, e.g. main() and then run that and see how the timings look like.

I believe what you want to do here is to make a single ode with parameters and then solve it repeatedtly varying the parameters

That’s a good suggestion, but that’s what my initial approach was.
Seemed to give no improvement benefits in this case!

That would probably work… if I can figure out how to make the input functions of my systems depend on those parameters, without it requiring Julia to recompile everything.
I suspect that it is possible somehow, I just haven’t found a way of achieving that yet.

remake is the function you are looking for.

using ModelingToolkit, OrdinaryDiffEq

function create_problem()
    @variables t x(t) f(t) # independent and dependent variables
    @parameters τ=0.75 ω=1.0 # parameters
    D = Differential(t) # define an operator for the differentiation w.r.t. time

    eqs = [
        f ~ sin((ω * rand()) * t)
        D(x) ~ (f - x) / τ
    ]

    @named fol_variable_f = ODESystem(eqs)
    _sys = structural_simplify(fol_variable_f)

    prob = ODEProblem(_sys, [x => 0.0], (0.0, 10.0))
end

function run_simulations(prob; trials = 50)
    @parameters ω
    for ω_ in 0.1:0.5:5
        for i in 1:trials
            prob_ = remake(prob, p=[ω=>ω_])
            sol = solve(prob_, Tsit5())
        end
    end
end
julia> prob = create_problem()
ODEProblem with uType Vector{Float64} and tType Float64. In-place: true
timespan: (0.0, 10.0)
u0: 1-element Vector{Float64}:
 0.0

julia> @time run_simulations(prob)
  0.072954 seconds (359.97 k allocations: 30.387 MiB, 68.91% compilation time)

julia> @time run_simulations(prob)
  0.051170 seconds (249.95 k allocations: 23.109 MiB, 53.83% gc time)
3 Likes