Slow indexing of constant-valued variables in MTK

Specific Domains Modelling & Simulations
,
1

Hi

I’m currently working with large hierarchical models in ModelingToolkit.jl and found that indexing the solution with constant-valued variables becomes incredibly slow for larger systems. This results in plotting of the solution taking multiple minutes whereas the actual solving only takes seconds.

Replacing this variable with a parameter or constant is the obvious solution (and does reduce the indexing time to normal levels), but doesn’t seem like an option for my case as the constant-valued variable in question is actually only constant for some subsystems and variable in others.

Please consider the below MWE, where B serves as the constant-valued variable:

using ModelingToolkit, OrdinaryDiffEq
import ModelingToolkit: t_nounits as t, D

n = 1000
@variables A(t)[1:n] = zeros(n) B(t)[1:n]
@parameters B_val[1:n] = rand(n)

eqs = [
    # some complex equation for A
    [D(A[i]) ~ sin(t + B[i]) for i in eachindex(A)]...,
    # constant value for B - in reality some B also have complex eqs
    [B[i] ~ B_val[i] for i in eachindex(B)]...,
]

@named sys = ODESystem(eqs, t)
sys_simpl = structural_simplify(sys);
prob = ODEProblem(sys_simpl, [], (0.0, 100))
sol = solve(prob);

@time sol[A[rand(1:n)]]; # ~0.0002s
@time sol[B[rand(1:n)]]; # ~0.8s

I feel like I must be doing something incredibly dumb but can’t figure it out. Any help would be greatly appreciated.

I’m using Julia 1.10.4 with following package versions:

  [961ee093] ModelingToolkit v9.41.0
  [1dea7af3] OrdinaryDiffEq v6.89.0
3

Welcome @bspanoghe! I’ve not dug into this, but there are two things here:

  • the simplified system doesn’t directly solve for B
  • solving for each B index involves compilation time the first time you do it and then the second time is fast.

I’m not sure how I’d apply the generic perf tips here, @nsajko if you have ideas please elaborate.

1 Like
4

Thank you for your response!

I’ve considered forcing the system to solve directly for B by instantiating it with the desired values
@variables B(t)[1:n] = rand(n)
and then setting the gradient to 0.
[D(B[i]) ~ 0 for i in eachindex(B)]...

This does solve the indexing time problem (at the cost of slightly higher solving time), but I was hoping for a cleaner solution.

Based on this thread it seems there is some way to mark variables you want to track, but the closest I’ve found is the irreducible metadata tag mentioned here:
@variables A(t)[1:n] = zeros(n) B(t)[1:n] [irreducible = true]

Sadly it slows down the solver immensely, making it an even worse solution for my problem.

5

Maybe it would be faster to use SymbolicIndexingInterface:

using SymbolicIndexingInterface
evalsol = getu(sol, B)
Bvals = evalsol(sol)  # should be a vector of vector of values of B at each time

You can reuse evalsol each time you want to pull out the B values (as long as the underlying problem is the same I believe), so it will only be slow the first time it is called and compiled.

2 Likes
6

I did a quick test and it seems that this method is equally quick to simply using sol[B] - I assume MTK might use SymbolicIndexingInterface under the hood for indexing?

However, I also learned some unexpected behavior for indexing multiple variables:
Indexing a subset of the variables simultaneously
sol[[B[i] for i in rand(1:n, 10)]]
is about equally quick as indexing just a single variable
sol[B[rand(1:n)]]
while indexing all variables sol[B] is twice as quick as indexing a single one - sol[B] ends up being ~1000x faster than indexing all B separately!

Indexing for all variables will be more complex for my actual use case but should get plotting times to reasonable levels, so thanks a lot!

1 Like
7

Glad that helped. It is using SymbolicIndexingInterface under the hood, but I was worried it was recompiling each time you called sol via a symbolic index (even if previously used). It is good to hear that it is apparently caching the function.

8

Every time you index a single (new) variable I imagine it is compiling a function to extract that variable, whereas when grabbing the whole vector it is only compiling one function. So that would give a noticeable performance difference.

9

Indeed, it’s very obvious in hindsight why it took so long to index the solution a few hundred times but simply hadn’t thought of it. Just changed the plotting function to do all the indexing in one step and it dropped plotting time from ~5 minutes to a few seconds. Thanks both for the help :slight_smile:

1 Like