I am running into problems regarding memory allocation that depends on the actual runtime of the solver of a stochastic jump process (i.e. it depends on T that defines tspan(0.0,T)). This gives notable performance drops when multi-threading a JumpProblem() using EnsembleProblem(args, EnsembleThreads()).
To illustrate the problem, I have a MWE that I present below. It highlights the memory issue without defining an ensemble problem. Note that the real problem I am solving is more complex and there the memory issues actually lead to decreased performance when multi-threading and/or using distributed solutions (using EnsembleThreads() and/or EnsembleDistributed()).
(All benchmarks below are subsequent calls to the functions. My code has all this in a single module, but I made it into a MWE for clarity.)
Minimal working example
As an illustration/MWE of my problem, let us consider a very simple 2-species Lotka-Volterra system with a carrying capacity K. This is just the standard 2-species generalized Lotka-Volterra model with unit interaction a_{xy} = a_{yx} = 1 and K_x = K_y = K.
For this problem, we expect simple logistic growth with rate \eta until x+y = K, after which the system is in a stable state.
Let us define the function using Catalyst.jl, and solve the ODE directly:
function lotkavolterra(;T::Float64=100.0, nsaves::Int64=1)
function growth(Ī·,X,Y,K)
return Ī·*(1-(X+Y)/K)
end
lv_model = @reaction_network begin
growth(Ī·,X,Y,K), X --> 2X
growth(Ī·,Y,X,K), Y --> 2Y
end
p = (:Ī· => 0.1, :K => 1000)
u0 = [:X => 50, :Y => 100]
tspan = (0.0,T)
Īt = (tspan[end]-tspan[begin])/nsaves
tsave = tspan[begin]:Īt:tspan[end]
odeprob = ODEProblem(lv_model, u0, tspan, p)
odesol = solve(odeprob, Tsit5(); saveat=tsave)
return odesol
end
julia> sol = lotkavolterra(T=100.0, nsaves=100);
julia> plot(sol)
So far so good.
Additionally, we expect Catalyst.jl to build non-allocating functions (see also this answer), so the total memory allocation should only depend on nsaves, not on T.
(I am only leaving in the Memory estimate, as the time is at the moment not of interest.)
julia> using BenchmarkTools
julia> @benchmark lotkavolterra(T=100.0, nsaves=10)
BenchmarkTools.Trial: 10000 samples with 6 evaluations.
Memory estimate: 6.91 KiB, allocs estimate: 59.
julia> @benchmark lotkavolterra(T=100.0, nsaves=100)
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
Memory estimate: 17.50 KiB, allocs estimate: 150
julia> @benchmark lotkavolterra(T=1000.0, nsaves=100)
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
Memory estimate: 17.50 KiB, allocs estimate: 150.
So far so good.
Now we simulate the Lotka-Volterra system by defining a JumpProblem, and we are now only interested in the final state, thus:
function lotkavolterra(;T::Float64=100.0)
function growth(Ī·,X,Y,K)
return Ī·*(1-(X+Y)/K)
end
lv_model = @reaction_network begin
growth(Ī·,X,Y,K), X --> 2X
growth(Ī·,Y,X,K), Y --> 2Y
end
p = (:Ī· => 0.1, :K => 1000)
u0 = [:X => 50, :Y => 100]
tspan = (0.0,T)
Īt = (tspan[end]-tspan[begin])/nsaves
tsave = tspan[begin]:Īt:tspan[end]
discrete_prob = DiscreteProblem(lv_model, u0, tspan, p)
jump_prob = JumpProblem(lv_model, discrete_prob, Direct())
jsol = solve(jump_prob, SSAStepper())
return jsol
end
First to double check:
>julia sol = lotkavolterra(T=100.0); plot(sol)
Seems good.
Memory usage increases as total time increases
Now the problem. What I expect is to see Catalyst.jl build non-allocating functions, similar to solving it with an ODE, but now I actually observe a dependence on the total run time T. Again omitting runtime information and only saving the final state by adding save_positions=(false,false) to the definition of JumpProblem(...), i.e.
jump_prob = JumpProblem(lv_model, discrete_prob, Direct(); save_position=(false,false))
julia> @benchmark lotkavolterra(T=100.0)
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
Memory estimate: 41.64 KiB, allocs estimate: 1703.
julia> @benchmark lotkavolterra(T=10.0)
BenchmarkTools.Trial: 10000 samples with 5 evaluations.
Memory estimate: 8.78 KiB, allocs estimate: 301.
Why does the memory estimate increase?
Investigations
Interestingly, when I change the growth rate function to an in-place version:
function growth!(Ī·,X,Y,K)
return Ī·*(1-(X+Y)/K)
end
nothing really changes ā as in, I get the same increase in memory and allocations.
What is more interesting is that when the specify carrying capacity term is removed (as in the examples of JumpProcesses), the memory does not increase! Other solves, such as RSSA() also have the same problem.
Question
What is happening here? Why does the memory usage increase when the runtime increases when I define a JumpProblem() from a DiscreteProblem(), even though only the last state is saved by setting save_positions=(false,false). Does this have to do with the interpolation that is done internally? Is there any way to have terms, such as the carrying-capacity term in this example, without increasing the memory usage?
I ask because my real problem is very similar ā basically an n-species Lotka-Volterra model with different carrying capacity terms, and recall I am constantly running into memory-related issues, especially when turning the JumpProblem() into an EnsembleProblem() for multi-threading and/or distributed computing later down the line.
Any help or tips would be greatly appreciated, thanks!
Full code
Module
module TestJumpProcesses
export lotkavolterra
using JumpProcesses
using Catalyst
using DifferentialEquations
using BenchmarkTools
function lotkavolterra(;T::Float64=100.0, nsaves::Int64=1)
function growth(Ī·,X,Y,K)
return Ī·*(1-(X+Y)/K)
end
lv_model = @reaction_network begin
growth(Ī·,X,Y,K), X --> 2X
growth(Ī·,Y,X,K), Y --> 2Y
end
p = (:Ī· => 0.1, :K => 1000)
u0 = [:X => 50, :Y => 100]
tspan = (0.0,T)
Īt = (tspan[end]-tspan[begin])/nsaves
tsave = tspan[begin]:Īt:tspan[end]
@info "Benchmarking ODE problem..."
odeprob = ODEProblem(lv_model, u0, tspan, p)
odesol = solve(odeprob, Tsit5(); saveat=tsave)
bode = @benchmark solve($odeprob, Tsit5(); saveat=$tsave)
@info "Benchmarking jump problem..."
discrete_prob = DiscreteProblem(lv_model, u0, tspan, p)
jump_prob = JumpProblem(lv_model, discrete_prob, Direct(); save_positions=(false,false))
sol = solve(jump_prob, SSAStepper())
bjump = @benchmark solve($jump_prob, SSAStepper())
return bode, bjump
end