JumpProblem with large number of jumps crashes

[gzagatti]

Following my post in discourse and advice from @isaacsas, I am raising this issue here.

I am trying to solve a JumpProblem defined on a network.

The problem consists of a SIR model in which I am trying to model every node of the network separately. For future, reference I am trying to simulate systems like the one described in this paper.

I have a working code that works with a small network (around 100 nodes), as I try to increase the size of nodes the code fails raising StackOverflow error.

See a minimal working version below:

using LightGraphs
using DifferentialEquations
using Plots
using Distributions
using SparseArrays
using LinearAlgebra

@info "Initializing system..."
degree = 5
kdist = Poisson(degree)
N = Int(100)
ks = 1
while sum(ks) % 2 != 0
    global ks = rand(kdist, N)
end

g = random_configuration_model(N, ks)

i₀ = rand(1:N, round(Int, N*0.1)) 
x₀ = [ones(N); zeros(N); zeros(N)]
x₀[i₀] .= 0
x₀[i₀ .+ N] .= 1
tspan = (0.0, 100.0)
p = [0.25, 0.05, degree, adjacency_matrix(g)]

function agg_sol(sol, N)
    agg = reshape(sol[:, :], (N, 3, :))
    agg = mean(agg; dims=1)
    agg = dropdims(agg; dims=1)
    agg = permutedims(agg, (2, 1))
    return agg
end

@info "Collecting jumps..."
function dNᵢ(i)

    function rate(u, p, t)
        u[i]*(p[1]/p[3])*dot(p[4][i, :], u[(N+1):2N])
    end

    function affect!(integrator)
        integrator.u[i] -= 1
        integrator.u[N+i] += 1
    end

    return rate, affect!

end

function dMᵢ(i)

    function rate(u, p, t)
        u[N+i]*p[2]
    end

    function affect!(integrator)
        integrator.u[N+i] -= 1

        integrator.u[2N+i] += 1
    end

    return rate, affect!

end

jumps = ConstantRateJump[]
for i in 1:N
    push!(jumps, ConstantRateJump(dNᵢ(i)...))
end

for i in 1:N
    push!(jumps, ConstantRateJump(dMᵢ(i)...))
end

@info "Building dependency graph..."
vtoj = Vector{Vector{Int64}}(undef, 3N)
for i in 1:N
    vtoj[i] = [i]
    vtoj[N+i] = [findnz(p[4][i, :])[1]; N+i]
    vtoj[2N+i] = []
end

jtov = Vector{Vector{Int64}}(undef, length(jumps))
for i in 1:N
    jtov[i] = [i, N+i]
    jtov[N+i] = [N+i, 2N+i]
end

jtoj = Vector{Vector{Int64}}(undef, 2N)
for i in 1:N
    jtoj[i] = [findnz(p[4][i, :])[1]; i; N+i]
    jtoj[N+i] = [findnz(p[4][i, :])[1]; i]
end

@info "Building discrete problem..."
discrete_prob = DiscreteProblem(x₀, tspan, p)

@info "Building jump problem..."
# see https://github.com/SciML/ModelingToolkit.jl/blob/f12f472f630fd85a6fab4ca547b1c679217c33df/src/systems/jumps/jumpsystem.jl
# see https://diffeq.sciml.ai/stable/types/jump_types/
jump_prob = JumpProblem(discrete_prob, RSSA(), JumpSet(jumps...);
    vartojumps_map=vtoj, jumptovars_map=jtov)
# DirectCR() also fails
# jump_prob = JumpProblem(discrete_prob, DirectCR(), JumpSet(jumps...); dep_graph=jtoj)

@info "Solving problem..."
jump_sol = solve(jump_prob)

@info "Plotting..."
plot(agg_sol(jump_sol, N));
ylims!(0, 1);
plot!(legend=:right);
savefig("test.png")

@info "Done."

If I try to increase the number of nodes above to k=1000 the code crashes. The end of the error (as it is otherwise too big to reproduce here) is copied below. I am not sure why StackOverflowError is being raised since the JumpSet is reasonably small (only 2000 functions).

My intuititon tells me that to simulate this model one only needs to update the rate of the whole model every time an event happens. Given that the dependency graph is provided, it should not be such an expensive operation or one that involves too many iterations.

The simulation could be slow if there are too many updates (as is the case here), but at least it should not crash after every update.

In discourse, @isaacsas suggested that "split_jumps is recursive and uses splatting" which could be the issue here.

...
_jl_invoke at /buildworker/worker/package_linux64/build/src/gf.c:2237 [inlined]
jl_apply_generic at /buildworker/worker/package_linux64/build/src/gf.c:2419
jl_apply at /buildworker/worker/package_linux64/build/src/julia.h:1703 [inlined]
true_main at /buildworker/worker/package_linux64/build/src/jlapi.c:560
repl_entrypoint at /buildworker/worker/package_linux64/build/src/jlapi.c:702
main at julia (unknown line)
__libc_start_main at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
unknown function (ip: 0x4007d8)
ERROR: LoadError: StackOverflowError:
Stacktrace:
 [1] split_jumps(::Tuple{}, ::NTuple{983, ConstantRateJump{var"#rate#1"{Int64}, var"#affect!#2"{
Int64}}}, ::Nothing, ::Nothing, ::ConstantRateJump{var"#rate#1"{Int64}, var"#affect!#2"{Int64}},
 ::ConstantRateJump{var"#rate#1"{Int64}, var"#affect!#2"{Int64}}, ::Vararg{Any, N} where N) (rep
eats 984 times)
   @ DiffEqJump ~/.julia/packages/DiffEqJump/5mkeM/src/jumps.jl:127
 [2] JumpSet(::ConstantRateJump{var"#rate#1"{Int64}, var"#affect!#2"{Int64}}, ::Vararg{DiffEqJum
p.AbstractJump, N} where N)
   @ DiffEqJump ~/.julia/packages/DiffEqJump/5mkeM/src/jumps.jl:106
 [3] top-level scope
   @ ~/phd/qe/test3.jl:106
in expression starting at /home/gzagatti/phd/qe/test3.jl:106

[isaacsas]

See my reply on Discourse. I think this is just more of a documentation issue. If you have lots of jumps you shouldn't use the splatting form of the JumpSet construction, but instead use the intrinsic constructor for the JumpSet struct, see the definition at

https://github.com/SciML/DiffEqJump.jl/blob/6e6ce598b164c10ce1c9c5e508b6b79e725d5482/src/jumps.jl#L90

Maybe it would make sense for us to add a kwarg only constructor version where users can pass the arrays manually? i.e. jset = JumpSet(; constantratejumps=crj_vec, variableratejumps=vrj_vec, etc)?

[isaacsas]

If you update to the latest release (6.14.2) I added the ability to just pass the jumps through a kwarg:

jset = JumpSet(; constant_jumps=jumps)
jump_prob = JumpProblem(discrete_prob, RSSA(), jset;
    vartojumps_map=vtoj, jumptovars_map=jtov)

[isaacsas]

I'm going to close this, but feel free to reopen if you are still having issues!

[gzagatti]

Thanks for the quick response. I tested the new API and it works with as much as 10k nodes in under 2s.