Extemely slow creation of `ODEProblem` from `ReactionSystem`

[johannesnauta]

Describe the bug 🐞

Me again. I am simulating Lotka-Volterra systems, and have recently jumped back on Catalyst.jl for the flexibility and comparison with other solvers (SDEs, SSAs, etc.) other than solving ODEs. However, I find that creating (not solving) the ODEProblem from a ReactionSystem that is sufficiently large takes an incredible amount of time. The time it takes reaches the point where trying to create a problem of medium sizes (~50 species) basically does not finish. Additionally, when using an ODESystem using ModelingToolkit.jl directly, and creating the ODEProblem does not seem to have this issue.

Minimal Reproducible Example 👇
To illustrate, consider the following functions, one using Catalyst.jl, and one using ModelingToolkit.jl. They both create a simple Lotka-Volterra system, that is

$$ \frac{dx_i}{dt} = r_i x_i ( 1 - \sum_j A_{ij} x_j ) $$

Using Catalyst.jl

(note: if there are other/better ways to generate the ReactionSystem, let me know!)

function create_rs(S::Int)
    t = default_t()
    @parameters r[1:S] A[1:S,1:S]
    @species (x(t))[1:S]

    #/ Allocate and fill
    growthrxs = Array{Reaction}(undef, S)
    interactionrxs = Array{Reaction}(undef, S, S)
    for i in 1:S
        #~ Growth
        growthrxs[i] = Reaction(r[i], [x[i]], [x[i]], [1], [2])
        #~ Diagonal terms
        interactionrxs[i,i] = Reaction(r[i]*A[i,i], [x[i]], [x[i]], [2], [1])
        for j in (i+1):S
            #~ Off-diagonal terms
            interactionrxs[i,j] = Reaction(r[i]*A[i,j], [x[i], x[j]], [x[j]], [1, 1], [1])
            interactionrxs[j,i] = Reaction(r[i]*A[j,i], [x[j], x[i]], [x[i]], [1, 1], [1])
        end
    end

    #/ Gather reactions
    rxs = [growthrxs..., interactionrxs...]
    @named glvrs = ReactionSystem(rxs, t, combinatoric_ratelaws=false)
    return Catalyst.complete(glvrs)
end

Using ModelingToolkit.jl

function create_ODESystem(S::Int)
	  #/ Define variables and parameters
    @variables (x(t))[1:S]
    @parameters r[1:S] A[1:S,1:S]

    eqns = [D(x[i]) ~ r[i]*x[i]*(1.0 - sum([A[i,j]*x[j] for j in 1:S])) for i in 1:S]
    @named odesys = ODESystem(eqns, t)
    return ModelingToolkit.complete(odesys)
end

Then, running the following:

julia> using BenchmarkTools

julia> S = 16;

julia> rs = create_rs(S);

julia> osys = create_ODESystem(S);

julia> rsprob = create_ODEProblem(rs);

julia> osprob = create_ODEProblem(osys);

julia> @btime rsprob = create_ODEProblem(rs);
  2.235 s (4388516 allocations: 242.32 MiB)

julia> @btime rsprob = create_ODEProblem(rs);
  2.241 s (4388516 allocations: 242.32 MiB)

julia> @btime osprob = create_ODEProblem(osys);
  13.086 ms (163580 allocations: 9.05 MiB)

julia> @btime osprob = create_ODEProblem(osys);
  13.267 ms (163580 allocations: 9.05 MiB)

From this the problem is already apparent; about two orders of magnitude longer time is needed to create the ODEProblem from the ReactionSystem. This issue is even more pressing the larger the no. of species S.
Now, I know I have read somewhere that it is expected to take a bit longer. But even converting the ReactionSystem to an ODESystem, i.e. using convert(ODESystem, rs), and then making the ODEProblem does not change anything for me.

Can someone explain to me what is going on? Am I missing something? Why does creating the ODEProblem takes so much time (and memory)? Why is this different from the ModelingToolkit.jl implementation?
I would like to learn how to do this properly, as I really would like to use the ReactionSystem to create other types of problems, such as SDEProblem, and JumpProblem -- which I think is one of the main selling points for using Catalyst.jl.

Any help or comments are greatly appreciated!

Environment (please complete the following information):

Details

  [c7e460c6] ArgParse v1.2.0
  [ec485272] ArnoldiMethod v0.4.0
  [6e4b80f9] BenchmarkTools v1.5.0
  [336ed68f] CSV v0.10.14
  [13f3f980] CairoMakie v0.12.9
  [479239e8] Catalyst v14.4.1
  [5ae59095] Colors v0.12.11
  [861a8166] Combinatorics v1.0.2
  [a93c6f00] DataFrames v1.6.1
  [0c46a032] DifferentialEquations v7.14.0
  [31c24e10] Distributions v0.25.111
  [ffbed154] DocStringExtensions v0.9.3
  [e30172f5] Documenter v1.7.0
  [d872a56f] ElectronDisplay v1.0.1
  [d4d017d3] ExponentialUtilities v1.26.1
  [68837c9b] FHist v0.11.6
  [86223c79] Graphs v1.11.2
  [033835bb] JLD2 v0.5.1
  [ccbc3e58] JumpProcesses v9.13.7
  [b964fa9f] LaTeXStrings v1.3.1
  [23fbe1c1] Latexify v0.16.5
  [2fda8390] LsqFit v0.15.0
  [961ee093] ModelingToolkit v9.39.1
  [2774e3e8] NLsolve v4.5.1
  [77ba4419] NaNMath v1.0.2
  [67456a42] OhMyThreads v0.6.0
  [1dea7af3] OrdinaryDiffEq v6.89.0
  [91a5bcdd] Plots v1.40.8
  [1fd47b50] QuadGK v2.11.0
  [0bca4576] SciMLBase v2.53.0
  [a9a3c162] SparseArrayKit v0.3.1
  [276daf66] SpecialFunctions v2.4.0
  [0c5d862f] Symbolics v6.11.0
  [bd369af6] Tables v1.12.0
  [3a884ed6] UnPack v1.0.2
  [8ba89e20] Distributed
  [37e2e46d] LinearAlgebra
  [9a3f8284] Random
  [2f01184e] SparseArrays v1.10.0
  [10745b16] Statistics v1.10.0

julia> versioninfo()
Julia Version 1.10.4
Commit 48d4fd48430 (2024-06-04 10:41 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 12 × 12th Gen Intel(R) Core(TM) i5-1235U
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, alderlake)
Threads: 8 default, 0 interactive, 4 GC (on 12 virtual cores)
Environment:
  JULIA_NUM_THREADS = 4

[isaacsas]

Thanks for the example; we'll have to look into this. ModelingToolkit development moves fast these days, and it looks like they have changed around some of the input processing from how it was done in the past. We may need to make updates here to mimic how they now internally process inputs for systems, and how array type symbolics are handled.

I'll try to play around with this when I have some time, unfortunately it may not be until next week due to other deadlines.

[johannesnauta]

No worries. I just found it very odd as it was not something I expected and, I believe, it worked for me before without issue. I am not super familiar with the ModelingToolkit-codebase, but if you have some hunches/pointers on where to look to see (changes in) the internal processing feel free to let me know so I can have look.

[isaacsas]

ModelingToolkit 9 has made a lot of under the hood changes, including how array symbolics are handled I think. It may be that we are not using the right approach with how to handle them any more.

You could look at the ODESystem constructors in ModelingToolkit and compare them to the Catalyst ReactionSystem constructor to see differences. It might also be good to actually profile the code and see what step is taking so long. Finally, it might be worth trying the four-argument constructor to ReactionSystem where you pass the species and variables too and seeing if that helps speed things up (but you might need to scalarize your array variables manually to use it).

[ChrisRackauckas]

Could you share a StatProfilerHTMlL.jl zip file report? That will make it easy to see what the culprit is.

[johannesnauta]

statprof.zip

[isaacsas]

@johannesnauta just to make sure we understand. There is no issue with the time it takes to make the systems then, it is with creating an ODEProblem from them? i.e. create_rs is fine?

If so, can you give us your code for creating the problems that you actually benchmarked. I don't see it above.

[isaacsas]

Your flamegraph seems to show that all the slowdown is actually in ModelingToolkit building an ODEProblem from the converted ODESystem we produce. @ChrisRackauckas should we not be scalarizing vector variables and/or parameters any more?

[johannesnauta]

Yes exactly. The only issue is when creating the ODEProblem from the ReactionSystem. And you're right, I forgot to include these in my initial post, somehow. Here are the functions:

function create_ODEProblem(
    rs::ReactionSystem;
    S=length(species(rs)),
    rv=ones(S), Av=ones(S, S), xv=ones(S),
    tspan=(0.0,100.0)
    )
    pmap = [rs.r => rv, rs.A => Av]
    xmap = [rs.x => xv]
    prob = ODEProblem(rs, xmap, tspan, pmap)
    return prob
end

function create_ODEProblem(
    osys::ODESystem;
    S=length(equations(osys)),
    rv=ones(S), Av=ones(S, S), xv=ones(S),
    tspan=(0.0,100.0)
    )
    pmap = [osys.r => rv, osys.A => Av]
    xmap = [osys.x => xv]
    prob = ODEProblem(osys, xmap, tspan, pmap)
    return prob
end

[TorkelE]

@isaacsas's guess that it might have to do with the vectorised parameters/species sounds like a good guess, there have been quite a lot of movements there lately. There was quite a lot of discussions a while ago, have forgotten some of the turns, but I did test building rather large systems using vector species/parameters, so it did work fine a couple of months ago. Probably requires some digging to sort out properly. The first step is probably to check whether the issue persist when only parameters/only species are on vector form (I remember these were treated differently), and maybe also with neither on that form.

[johannesnauta]

Seems that it might be connected to symbolic arrays indeed, here's a crude test (it's cumbersome to time out all interaction parameters by hand, obviously, so just a small system of 2 species for now).

function create_scalarrs()
    t = default_t()
    @parameters r1 r2
    @parameters A11 A12 A21 A22
    @species x1(t) x2(t)

    rxs = [
        Reaction(r1, [x1], [x1], [1], [2]),
        Reaction(r2, [x2], [x2], [1], [2]),
        Reaction(r1*A11, [x1], [x1], [2], [1]),
        Reaction(r1*A12, [x1,x2], [x2], [1,1], [1]),
        Reaction(r2*A21, [x1,x2], [x1], [1,1], [1]),
        Reaction(r2*A22, [x2], [x1], [2], [1])
    ]

    x = [x1, x2]
    p = [r1, r2, A11, A12, A21, A22]
    @named glvrs = ReactionSystem(rxs, t, x, p, combinatoric_ratelaws=false)
    return complete(glvrs)
end

which gives

julia> scalrs = create_scalarrs()

julia> scalprob = create_scalarODEProblem(scalrs);

julia> @btime scalprob = create_scalarODEProblem(scalrs);
718.716 μs (7103 allocations: 481.45 KiB)

julia> @btime scalprob = create_scalarODEProblem(scalrs);
  720.363 μs (7103 allocations: 481.45 KiB)

while using symbolic arrays with the create_rs(S)-function as defined above, for S=2 gives:

julia> vecrs = create_rs(2);

julia> vecode = create_ODEProblem(vecrs);

julia> @btime vecode = create_ODEProblem(vecrs);
  1.309 ms (10633 allocations: 647.30 KiB)

julia> @btime vecode = create_ODEProblem(vecrs);
  1.309 ms (10633 allocations: 647.30 KiB)

[isaacsas]

From what I can tell ModelingToolkit no longer scalarizes, however, unknowns and parameters are handled differently:

julia> function create_ODESystem(S::Int)
                 #/ Define variables and parameters    
           t = ModelingToolkit.t_nounits
           D = ModelingToolkit.D_nounits
           @variables (x(t))[1:S]
           @parameters r[1:S] A[1:S,1:S]

           eqns = [D(x[i]) ~ r[i]*x[i]*(1.0 - sum([A[i,j]*x[j] for j in 1:S])) for i in 1:S]
           @named odesys = ODESystem(eqns, t)
           return ModelingToolkit.complete(odesys)
       end
create_ODESystem (generic function with 1 method)

julia> osys = create_ODESystem(3)
Model odesys with 3 equations
Unknowns (3):
  (x(t))[1]
  (x(t))[2]
  (x(t))[3]
Parameters (2):
  A
  r

Notice that while it isn't scalarizing it is detecting and explicitly storing each component of the vector unknown x. In contrast, the parameters are just stored in the matrix/vector form.

We only scalarize in a few places, but I have no idea how this might mess things up in Catalyst since now the number of parameters is not the same as the length of the parameter vector (which we might assume), and we'd now have to handle in places expanding a vector x into its components via indexing without scalarizing.

[isaacsas]

@johannesnauta I just wanted to point out that it looks like your systems aren't actually exactly the same? When I take the difference of the generated ODEs rhs and simplify them I get some leftover terms:

3-element Vector{Any}:
 0
  A[2, 1]*r[1]*(x(t))[1]*(x(t))[2] - A[2, 1]*r[2]*(x(t))[1]*(x(t))[2]
  A[3, 1]*r[1]*(x(t))[1]*(x(t))[3] - A[3, 1]*r[3]*(x(t))[1]*(x(t))[3] + A[3, 2]*r[2]*(x(t))[2]*(x(t))[3] - A[3, 2]*r[3]*(x(t))[2]*(x(t))[3]

(This is the ODE model minus the converted reaction system model when S = 3.)

I guess you aren't exactly setting the rate consistently in the two models?

[isaacsas]

So good news; I think it should be relatively easy to make Catalyst match MTK's new conventions, see #1052. Bad news; it appears our tests will need lots of modifications if we make such a change.

Unfortunately I don't think this is something I have the time to handle in the next few weeks due to other deadlines, but that PR seems to get things fixed in terms of the generated ODESystem if you want to test it out on your example. Now someone has to go through all the tests and update them accordingly to no longer assume all inputs are scalars.

[isaacsas]

This will probably also require a breaking release of Catalyst once updated since it does change how parameters are represented internally and stored in what is passed to users (only returning the vector/matrix symbolics instead of their components).