Huge RAM usage with JumpProcesses.jl

General Usage Performance
5

I have created something that uses a large amount of RAM ~7GB (see MWE below). I think it is using so much memory for the same reasons my problem is (just a lot of stochastic reactions occurring). For instance the shorter the timespan, the less memory used for both problems by a significant amount. I’m wondering if there’s some specific code optimisation to be used with this problem or if there’s no way round this in a fully stochastic model due to the large number of jumps occurring.

using DifferentialEquations, JumpProcesses, BenchmarkTools

function params()
    return (
        k_fast_production = 1000.0,   
        k_slow_production = 50.0,     
        k_fast_decay = 50.0,         
        k_slow_decay = 0.5,         
        n_species = 10             
    )
end

function control_ode!(du, u, p, t)
    du[1] = -0.01 * u[1] 
end

initial_u = [100.0]               
initial_jump_species = zeros(params().n_species)  
tspan = (0.0, 1000.0)           

# Set up jump processes with varied rates and high memory demands
jumps = []
for i in 1:params().n_species
    if i == 1
        # High-expression species
        fast_production_rate(u, p, t) = p.k_fast_production
        fast_decay_rate(u, p, t) = p.k_fast_decay
        affect_fast_production!(integrator) = (integrator.u[i] += 1; nothing)
        affect_fast_decay!(integrator) = (integrator.u[i] -= 1; nothing)
        
        push!(jumps, ConstantRateJump(fast_production_rate, affect_fast_production!))
        push!(jumps, VariableRateJump(fast_decay_rate, affect_fast_decay!))
    else
        # Low-expression species with varied production and decay rates
        slow_production_rate(u, p, t) = p.k_slow_production * i
        slow_decay_rate(u, p, t) = p.k_slow_decay * i
        affect_slow_production!(integrator) = (integrator.u[i] += 1; nothing)
        affect_slow_decay!(integrator) = (integrator.u[i] -= 1; nothing)
        
        push!(jumps, ConstantRateJump(slow_production_rate, affect_slow_production!))
        push!(jumps, VariableRateJump(slow_decay_rate, affect_slow_decay!))
    end
end

control_prob = ODEProblem(control_ode!, initial_u, tspan, params())
jumpset = JumpSet(jumps...)
jump_prob = JumpProblem(control_prob, Direct(), jumpset, save_positions=(false,false))

@btime solve(jump_prob, Tsit5(), maxiters=1e10)