# First solve of \`ODEProblem\` and running out of memory: compilation problems?

**URL:** <https://discourse.julialang.org/t/first-solve-of-odeproblem-and-running-out-of-memory-compilation-problems/126160>\
**Category:** Modelling & Simulations\
**Tags:** modelingtoolkit, differentialequation\
**Created:** [February 21, 2025, 4:38pm UTC](https://discourse.julialang.org/t/first-solve-of-odeproblem-and-running-out-of-memory-compilation-problems/126160 "2025-02-21T16:38:33Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![johannesnauta](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/johannesnauta/32/47434_2.png) [@johannesnauta](https://discourse.julialang.org/u/johannesnauta)\
**Post date:** [February 21, 2025, 4:38pm UTC](https://discourse.julialang.org/t/first-solve-of-odeproblem-and-running-out-of-memory-compilation-problems/126160/1 "2025-02-21T16:38:33Z")

</div>

Hey it’s me again,  
I keep encountering memory-related issues with the first call to `solve(prob)`, where `prob` is an `ODEProblem` (see details below). While this is probably expected, as it is most likely related to the “time to first plot”-problem, it becomes problematic when the first call actually crashes due to running out-of-memory (OOM), or when the first call takes an extraordinary amount of time.  
In my specific case, this is sadly exactly what happens when trying to simulate (relatively) large systems of coupled ODEs. In some recent posts, I have already been investigating some issues related to [saving](https://discourse.julialang.org/t/storing-only-rolling-mean-and-variance-when-solving-sdeproblems/125900/5), [callbacks](https://discourse.julialang.org/t/savingcallback-when-using-ensemble-simulations/88483/10), [`EnsembleProblems`](https://discourse.julialang.org/t/solving-ensembleproblem-efficiently-for-large-systems-memory-issues/116146/9), and more, yet this underlying issue has not yet been resolved for me. Sadly, the systems I am interested in are of sizes for which a first call to `solve(...)` does not resolve as I appear to be running OOM.

To illustrate my problem, as always, I consider the generalized disordered Lotka-Volterra model:

\frac{dx\_i}{dt} = x\_i \big( 1 - x\_i - \sum\_{j \neq i} A\_{ij} x\_j \big)

where entries of the interaction matrix A are random. Code is as follows, using `ModelingToolkit.jl` and `DifferentialEquations.jl`;

```julia
using DifferentialEquations
using ModelingToolkit
using ModelingToolkit: t_nounits as t, D_nounits as D

using Random
using SparseArrays

# ODESystem
function create_odesystem(S::Int)
    @variables (x(t))[1:S]
    @parameters A[1:S,1:S]

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

# ODEProblem
function create_odeproblem(
    sys::ODESystem;
    c=0.5, tspan=(0.0,1e3), rng=Random.Xoshiro(1234)
)
    S = length(sys.x)
    #/ Create parameters
    Av = sprandn(rng, S, S, c) ./ sqrt(S)
    for i in 1:S
        Av[i,i] = 1.0
    end
    xv = rand(rng, S)
    pmap = [sys.A => Av]
    xmap = [sys.x => xv]
    #/ Create ODEProblem
    prob = ODEProblem(sys, xmap, tspan, pmap)
    return prob
end

# Solve ODEProblem
function solve_odeproblem(prob::ODEProblem; save=false)
    return solve(
        prob,
        AutoTsit5(TRBDF2()),
        callback=PositiveDomain(save=false),
        save_everystep=save,
        save_start=save
    )
end

```

Then (note that `include(...)` statements are excluded for brevity):

```julia
julia> S = 16;
julia> sys = create_odesystem(S);
julia> prob = create_odeproblem(sys);

julia> @time sol = solve_odeproblem(prob);
  4.444015 seconds (20.86 M allocations: 1.066 GiB, 7.36% gc time)

julia> @time sol = solve_odeproblem(prob);
  0.007990 seconds (168.19 k allocations: 9.213 MiB)

```

Note the large no. of allocations and memory usage on the first call. In this case, there’s not a big problem, of course; just do it once, and subsequent calls are very fast. This raises the question whether using `PackageCompiler.jl` to make a sysimage of `DifferentialEquations` (or perhaps `OrdinaryDiffEq`) may be worth it — yet trying it out I did not see any improvements (so far).

The problem is bigger in larger systems

```julia
julia> S = 100;
julia> sys = create_odesystem(S);
julia> prob = create_odeproblem(sys);

julia> @time sol = solve_odeproblem(prob);
  8.313252 seconds (42.05 M allocations: 2.202 GiB, 6.66% gc time, 96.34% compilation time)

```

For my current projects I intend to investigate systems like the one above, but with additional processes/interactions, yet these often need to be investigated relatively large systems, which with my current implementation I simply cannot do.

Note that I am also interested in running many distinct instances, so any potential solutions should also apply to `EnsembleProblem`s.

## Questions

- What is the main reason that the first call to `solve` takes so much memory? In addition, subsequent calls with a different `S` also appear to take up quite some time/memory, so I cannot simply call it for, say, `S=1`, and afterwards investigate a larger system.
- Is there a way to overcome the memory overhead of the first call somehow? Can I exploit things like `PackageCompiler.jl`, or other implementations of the code above to drastically reduce the overhead of the first call?

All advice, help and tips are greatly appreciated. Thanks a lot in advance!

---

<div class="post-metadata">

**Author:** ![JonasWickman](https://avatars.discourse-cdn.com/v4/letter/j/9de0a6/32.png) [@JonasWickman](https://discourse.julialang.org/u/JonasWickman)\
**Post date:** [February 21, 2025, 10:04pm UTC](https://discourse.julialang.org/t/first-solve-of-odeproblem-and-running-out-of-memory-compilation-problems/126160/2 "2025-02-21T22:04:05Z")

</div>

I don’t really know ModelingToolkit.jl, but if you’re willing to forgo it, there are certainly ways of improving both the compilation and run-time performance of your code. Here is an example:

```julia
using DifferentialEquations
using Random
using SparseArrays
using LinearAlgebra
using PreallocationTools

function lotka_volterra_ode!(dx, x, p, t)

    inter_spec_comp = PreallocationTools.get_tmp(p.cache, x)
    mul!(inter_spec_comp, p.A, x)

    @. dx = x*(1 - inter_spec_comp)
    
    return nothing

end

function make_parameters(S; c = 0.5, rng = Random.Xoshiro(1234))
    A = sprandn(rng, S, S, c) ./ sqrt(S)
    for i ∈ axes(A, 1)
        A[i,i] = 1.0
    end
    x_temp = zeros(S)
    cache = PreallocationTools.DiffCache(x_temp)
    p = (; A, cache)
    return p
end

function solve_odeproblem(prob::ODEProblem; save=false)
    return solve(
        prob,
        AutoTsit5(TRBDF2()),
        callback=PositiveDomain(save=false),
        save_everystep=save,
        save_start=save
    )
end
## --
S = 16
t_end = 1e3

p = make_parameters(S)
x0 = rand(Random.Xoshiro(1234), S)

ode_prob = ODEProblem{true}(lotka_volterra_ode!, x0, (0.0, t_end), p)
## --
@time sol = solve_odeproblem(ode_prob)
@time sol = solve_odeproblem(ode_prob)

```

The initial allocation is still pretty big (I got ~15 M allocations, ~750 MiB), but it no longer scales with system size S. I think this is because in your code, the RHS function of the ODE allocates quite a lot (`A[i,:]` and `A[i,:].*x`both allocate temporary arrays, and as the system size increases these allocations get larger).

Additionally, while this in and of itself doesn’t affect allocations, the way you’re multiplying a matrix with a vector in pieces by summing over a product of matrix rows and the vector is also not ideal for performance.

I’m not sure what best practices are for writing non-allocating ModelingToolkit code is, hopefully someone else can chime in.

Above, I’ve made the RHS non-allocating by using the `mul!` function from the `LinearAlgebra` standard library in conjunction with PreallocationTools.jl, which is a SciML package for managing autodiff-friendly caches.

---

<div class="post-metadata">

**Author:** ![johannesnauta](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/johannesnauta/32/47434_2.png) [@johannesnauta](https://discourse.julialang.org/u/johannesnauta)\
**Post date:** [February 23, 2025, 8:10pm UTC](https://discourse.julialang.org/t/first-solve-of-odeproblem-and-running-out-of-memory-compilation-problems/126160/3 "2025-02-23T20:10:10Z")

</div>

Thanks, this looks great. The reason I like `ModelingToolkit` so much, is that it becomes very easy to change the problem, make `EnsembleProblem`s, etc., which I always find a bit confusing with `DifferentialEquations` alone. That is, I typically make specific setters, i.e.

```julia
#/ Specify setters
@unpack x, A = prob.f.sys
xsetter = ModelingToolkit.setu(prob, x)
Asetter = ModelingToolkit.setp(prob, A)

```

and then modify these in the `prob_func`

```julia
#/ Define TaskLocalValue that deepcopy's the problem to the task if it does not exist
#~ see: https://juliafolds2.github.io/OhMyThreads.jl/stable/literate/tls/tls/#TLV
tlv_prob = TaskLocalValue{ODEProblem}(() -> deepcopy(prob))

function set_interactions(prob, i, repeat)        
    localprob = tlv_prob[]
    xsetter(localprob, _xv[i])
    Asetter(localprob, _Av[i])
    return localprob
end

```

but, in this case, I can most likely just `remake(prob, [p.A = A])`, i.e.

```julia
function create_ensembleproblem(
    prob::ODEProblem;
    c=0.5,
    ntrajectories = 32,
    rng=Random.Xoshiro(1234)
)
    S = length(prob.u0)
    #/ Specify parameters
    _xv = [rand(rng, S) for _ in 1:ntrajectories]
    _Av = [sprandn(rng, S, S, c) ./ sqrt(S) for _ in 1:ntrajectories]
    for n in 1:ntrajectories, i in 1:S
        _Av[n][i,i] = 1.0
    end

    #/ Define TaskLocalValue that deepcopy's the problem to the task if it does not exist
    #~ see: https://juliafolds2.github.io/OhMyThreads.jl/stable/literate/tls/tls/#TLV
    tlv_prob = TaskLocalValue{ODEProblem}(() -> deepcopy(prob))

    function set_interactions(prob, i, repeat)        
        localprob = tlv_prob[]
        remake(localprob, u0=_xv[i], p=(; A=_Av[i]))
        return localprob
    end

    output_func(sol, i) = (last(sol), false)

    #/ Create EnsembleProblem
    eprob = EnsembleProblem(
        prob, prob_func=set_interactions, output_func=output_func, safetycopy=false
    )
    return eprob
end

```

I will need to double-check this tomorrow if the results are as I expect, but this does seem to do the trick! Thanks a lot for bringing `PreallocationTools.jl` to my attention. I will make sure to read the docs on that.

---

<div class="post-metadata">

**Author:** ![johannesnauta](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/johannesnauta/32/47434_2.png) [@johannesnauta](https://discourse.julialang.org/u/johannesnauta)\
**Post date:** [February 24, 2025, 10:58am UTC](https://discourse.julialang.org/t/first-solve-of-odeproblem-and-running-out-of-memory-compilation-problems/126160/4 "2025-02-24T10:58:33Z")

</div>

> [@johannesnauta](#):
>
> ```julia
> function set_interactions(prob, i, repeat)        
> localprob = tlv_prob[]
> remake(localprob, u0=_xv[i], p=(; A=_Av[i]))
> return localprob
> end
> 
> ```

This part has to be changed slightly to accomodate values that do not change (if they exist), i.e. this needs to be

```julia
function set_interactions(prob, i, repeat)        
    localprob = tlv_prob[]
    return remake(localprob, u0=_xv[i], p=merge(localprob.p, (; A=_Av[i])))
end

```

This seems to do the trick.

By the way, is there any guarantee that the `DiffCache` works while multi-threading? I cannot seem to find information on it, but running the code and testing it against my old implementation seems that this way is ok to implement, and I do not need to take it explicitly into account. Please correct me if I am wrong.