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

Specific Domains Modelling & Simulations
,
2

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:

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,:].*xboth 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.

1 Like