Reducing memory allocation in DifferentialEquations.jl

Question

I'm using DifferentialEquations.jl to solve an ODE system as shown below. The result is not really relevant since p only contains test parameters for the purpose of producing a MWE, but the key is that I am seeing a lot of memory allocation despite using an in-place ODE function.

using DifferentialEquations

function ode_fun!(du,u,p,t)
    a,b,c,d,e = p

    X = @. u[1] * a * ((b-c)/b)
    Y = @. u[2] * d * ((b-e)/b)

    du[1] = -sum(X) + sum(Y) - u[1]*u[2]
    du[2] = sum(X) - sum(Y) - u[1]*u[2]
end

#exemplary parameters 
a = collect(10:-0.1:0.1)
b = a.^2
c = b*0.7
d = collect(0.01:0.01:1)
e = b*0.3

u0 = [1.0, 0.5]
p = [a,b,c,d,e]
tspan = [0.0, 100.0]
t = collect(0:0.01:100) 

prob = ODEProblem(ode_fun!,u0,tspan,p,saveat=t) 
@time sol = solve(prob)

1.837609 seconds (5.17 M allocations: 240.331 MiB, 2.31% gc time) #Julia 1.5.2

Since I need to solve this ODE system repeatedly I would like to reduce allocations as much as possible and am wondering if there is anything that can be done about them. I have been wondering if the issue lies with X and Y and have tried to preallocate these outside the ODE function, but have unfortunately not succeeded in reducing allocations that way.

Answer 1

I'm pretty sure this should be faster and half the allocations

function ode_fun!(du,u,p,t)
    a,b,c,d,e = p
    XmY = @. u[1] * a * (1-c/b) - u[2] * d * (1-e/b)
    sXmY = sum(XmY)
    du[1] = -sXmY - u[1]*u[2]
    du[2] = sXmY - u[1]*u[2]
end

There's probably a way to get rid of all of them, but I'm not a DifferentialEquations expert.

Answer 2

Preamble: I realize this is an old question and the asker may have learned how to avoid the large number of allocations. This answer is aimed at a reader who is still making the mistake in the post.

I believe the allocations you are seeing (as well as the slowdown) are mainly from compilation. When I run your script twice on Julia 1.9.1, I get¹

julia> include("your_script.jl")
2.287537 seconds (3.33 M allocations: 216.865 MiB, 3.31% gc time, 99.91% compilation time)
julia> include("your_script.jl")
0.069226 seconds (46.83 k allocations: 4.283 MiB, 97.70% compilation time: 100% of which was recompilation)

On the first call, Julia must compile all the functions needed to run the code, including the ones from DifferentialEquations.jl (see this section in the Performance Tips). This dominates the the actual runtime of the code, which only takes a fraction of a second. On subsequent calls, Julia only needs to perform a significantly quicker recompilation (which in this case still dominates the run time).

The second significant mistake is that Performance critical code should be inside a function. Simply wrapping the code in a function (main in the code below) yields significant savings in time and allocations:

# new_script.jl
using DifferentialEquations

function ode_fun!(du,u,p,t)
    # ...
end

function main()
    # rest of the code
end

julia> include("new_script.jl")
main (generic function with 1 method)
julia> main()
2.334540 seconds (3.33 M allocations: 216.774 MiB, 2.50% gc time, 99.92% compilation time)
julia> main()
0.001639 seconds (10.94 k allocations: 2.260 MiB)    # 40x speedup

As before, we have to let Julia compile the code before we can get the actual run time.

Reusing arrays

Finally, if you need to run the solver many times, it helps to preallocate the arrays and reuse as shown below. This results in relatively small savings compared to the changes described above.

using DifferentialEquations

function ode_fun!(du,u,p,t)
    a,b,c,d,e,X,Y = p              # X, Y from p

    @. X = u[1] * a * ((b-c)/b)    # avoid allocating new X, Y
    @. Y = u[2] * d * ((b-e)/b)

    du[1] = -sum(X) + sum(Y) - u[1]*u[2]
    du[2] = sum(X) - sum(Y) - u[1]*u[2]
end

function main()

    #exemplary parameters
    a = collect(10:-0.1:0.1)
    b = a.^2
    c = b*0.7
    d = collect(0.01:0.01:1)
    e = b*0.3
    X = similar(a)             # preallocate X, Y
    Y = similar(d)

    u0 = [1.0, 0.5]
    p = [a,b,c,d,e,X,Y]        # pass X, Y to solver through p
    tspan = [0.0, 100.0]
    t = collect(0:0.01:100)

    prob = ODEProblem(ode_fun!,u0,tspan,p,saveat=t)

    function solve_many_times()
        for i in 1:10000
            sol = solve(prob)
        end
    end

    @time solve_many_times()

    return

end

---

1. Note that in the currect version of Julia, @time reports also the percentage of time spent on compilation.