Modeling Toolkit: mtkcompile taking forever for large systems

Specific Domains Modelling & Simulations
1

I recently re-coded an advection-reaction biogeochemical box model using ModelingToolkit.jl, that allows a smooth transition form papers’ equations to actual code and make it easily understandable. I build the model in a very naive way, by writing multiple .jl files, each one containing the equations and variables, and in a script I just have the different include(...).

In the native version, I have about ~500 equations, and only ~90 are actual independant ODEs (others are fluxes, or forcing set via CallBacks, temporary variables to make it look nice, …). In this case, calling mtkcompile takes around ~20 seconds (without julia compilation) which is enough in my case (and solving the ODEs takes less than a sec).

However, in a second version, where I increase the discretization of my model (i.e., the number of oceanic boxes), it takes forever (I manually killed it after 3h). In that set-up, the number of equation is about ~4400 (that’s not that huge) before structural simplification (and the number of actual independent ODEs should be around ~650).

I have quite complex-nonlinear equation, that I tried to hide behind registered functions (using @register_symbolic) but it did not help. I also had a look at JuliaSimCompiler.jl but it does not seem to be compatible with the latest versions of ModelingToolkit.jl.

It would be difficult to provide a mwe but if you have advices I would be glad to hear them.

Thanks,
L.

julia> versioninfo()
Julia Version 1.11.5
Commit 760b2e5b739 (2025-04-14 06:53 UTC)

[961ee093] ModelingToolkit v10.10.0
1 Like
2

Open an issue on the repo with something that we can test. Specifically I’d like to test the impact of feat: add `SemilinearODEFunction` and `SemilinearODEProblem` by AayushSabharwal · Pull Request #3739 · SciML/ModelingToolkit.jl · GitHub

3

I downgraded my Julia version to v1.10.0, and managed to instal JuliaSimCompiler.jl. Now, after a bit of work on the code, it compiles well (and fast, less than ~10 sec for the large system. Now much of the time is spend when executing the first call to solve).

However, I face a new issue: it seems that callbacks are not working with JuliaSimCompiler.jl:

# ... define eqs ...

frequency = 1.0
evs = Any[]
push!(evs,
	frequency => [x ~ x * 1.5], # x(t) is a variable
)

@named sys = ODESystem(eqs, t;
	discrete_events = evs)

sys_compiled = structural_simplify(sys)

ode_prob = ODEProblem(sys_compiled, merge(u0_map, p_map), (0.0, 10.0);
	sparse = true,
)

# solve with Rodas5 ...

is working but

frequency = 1.0
evs = Any[]
push!(evs,
	frequency => [x ~ x * 1.5], # x(t) is a variable
)

@named sys = ODESystem(eqs, t;
	discrete_events = evs)

IRsys = IRSystem(sys)

IRsys_compiled = structural_simplify(IRsys)

IRode_prob = ODEProblem(IRsys_compiled, merge(u0_map, p_map), (0.0, 10.0);
	sparse = true,
)

# solve with Rodas5 ...

is not.

Is there a workaround to this issue (avoid the symbolic callback for instance)?

Btw, is JuliaSimCompiler.jl still maintained? It was not easy to install and depends on a older version of ModelingToolkit.jl.

julia> versioninfo()
Julia Version 1.10.0
Commit 3120989f39b (2023-12-25 18:01 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)

pkg> st
Status `[...]/Project.toml`
  [99985d1d] AbstractGPs v0.5.24
  [8bb1440f] DelimitedFiles v1.9.1
  [0c46a032] DifferentialEquations v7.16.1
  [033835bb] JLD2 v0.5.15
  [bc7fa6ce] JuliaHub v0.1.14
  [8391cb6b] JuliaSimCompiler v0.1.26
⌅ [961ee093] ModelingToolkit v9.59.0
  [6f286f6a] MultivariateStats v0.10.3
  [85f8d34a] NCDatasets v0.14.8
  [67456a42] OhMyThreads v0.8.3
  [91a5bcdd] Plots v1.40.16
  [2913bbd2] StatsBase v0.34.5
  [f3b207a7] StatsPlots v0.15.7
  [fd094767] Suppressor v0.2.8
  [5d786b92] TerminalLoggers v0.1.7
  [37e2e46d] LinearAlgebra
  [56ddb016] Logging
  [9a3f8284] Random
Info Packages marked with ⌅ have new versions available but compatibility constraints restrict them from upgrading. To see why use `status --outdated`
4

According to the function

function IRSystem(eqs::Vector{Equation}, iv; defaults = EMPTY_DEFAULTS,
        ps = [], sts = [], name_to_id = nothing, continuous_events = nothing)

it seems that discrete_events are always set to nothing, while continuous_events are handled correctly.

A workaround is to use:

frequency = 1.0
evs = Any[]
push!(evs,
	[t % frequency ~ 0] => [x ~ x * 1.5], # x(t) is a variable
)

@named sys = ODESystem(eqs, t;
	continuous_events = evs)
IRsys = IRSystem(sys)

[...]

and specify the tstops to the solver.

This is probably not the most elegant way of doing it, but it is working.

5

It is not, but we’re putting similar tricks directly to MTK. If you open an issue with the tests we can figure out what’s required and needed.

1 Like
6

Should I create an issue using my symbolic kite model? It has discretization as well, so it should work well for testing.

Example for testing (make sure you are in a clean directory):

using SymbolicAWEModels, KiteUtils, VortexStepMethod, ModelingToolkit

set_data_path("data")
if !ispath("data/settings_ram.yaml")
    SymbolicAWEModels.copy_model_settings()
end

function generate_sys(segments)
    set = Settings("system_ram.yaml")
    set.segments = segments
    set.physical_model = "ram"
    sam = SymbolicAWEModel(set)
    println("Creating sys with $segments segments")
    @time inputs = SymbolicAWEModels.create_sys!(sam, sam.sys_struct; 
        init_va_b=zeros(1,3), prn=false)
    println("Simplifying sys with $segments segments")
    @time sys = mtkcompile(sam.full_sys; inputs)
end

generate_sys(3)
generate_sys(3) # 50 seconds
generate_sys(6) # 80 seconds
generate_sys(10) # 130 seconds
nothing
7

That would help. Even better would be to make it into a benchmark like Multibody Robot, compilation and simulation performance · The SciML Benchmarks or since it’s a discretization, Thermal Fluid ODE Compilation and Perf · The SciML Benchmarks. The benchmarks that we currently have are what’s driving the development goals, so if you add yours to it then we can track it.

1 Like
8

Cool! I will do that.

9

Great !

Thanks for taking care of this @Bart_van_de_Lint.

Atm, I think I will stick to the JuliaSimCompiler version, and leave it as soon as the benchmarks are reasonable (I also remarked during my testing period that older version of MTK - I think ~8.7 - could simplify quite fast but was not reducing the system as much as more recent versions).

1 Like