Jump diffusion problem in network does not scale

[gzagatti]

As reported in this Julia Discourse thread. The following model setup does not scale well beyond about 25 nodes.

The model is intended to simulate an SIR type diffusion in an arbitrary network given by an adjacency matrix A. After model setup, I convert it to a JumpProblem. For future, reference I am trying to simulate systems like the one described in this paper.

The convenience and flexibility of Catalyst.jl in setting up diffusion problems like the one presented here is a big advantage of this library over lower-level setup via DiffEqJump.jl. However, the program is slow both during initialization when calling @variables and ReactionSystem as well as during conversion when calling convert. The program then fails when we attempt to solve it --- related with issue SciML/JumpProcesses.jl#180.

It would also be desirable and more convenient if we could map the adjacency matrix like Pair.([β, γ, k, A], p) instead of having to vectorize the adjacency matrix. Falling to vectorize the parameters will result in an error.

using LightGraphs
using DifferentialEquations
using Plots
using GraphPlot
using LinearAlgebra
using Distributions
using Catalyst
using ModelingToolkit

@info "Initializing system..."
degree = 5
kdist = Poisson(degree)
N = Int(20)
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(@view(sol[:, :]), (N, 3, :))
    agg = mean(agg; dims=1)
    agg = dropdims(agg; dims=1)
    agg = permutedims(agg, (2, 1))
    return agg
end

@info "Initializing variables..."
@parameters t
@variables β γ k A[1:N, 1:N] x[1:3N](t)

@info "Collecting reactions..."
rx = []
for i in 1:N
    push!(rx, Reaction(
            x[i]*(β/k)*dot(A[:, i], x[(N+1):2N]), 
            [x[i]], 
            [x[N+i]];
            only_use_rate=true)
    )
end

for i in 1:N
    push!(rx, Reaction(
            γ,
            [x[N+i]],
            [x[2N+i]];
            only_use_rate=false)
    )
end

@info "Building reaction system..."
rs = ReactionSystem(rx, t, x, reduce(vcat, [[β, γ, k], vec(A)]))

@info "Converting to jump system..."
jump_sys = convert(JumpSystem, rs; combinatoric_ratelaws=false)

@info "Building discrete problem..."
discrete_prob = DiscreteProblem(
    jump_sys, 
    Pair.(x, x₀), 
    tspan,
    Pair.(
        reduce(vcat, [[β, γ, k], vec(A)]), 
        reduce(vcat, [p[1:3], vec(adjacency_matrix(g))])
    )
)

@info "Building jump problem..."
jump_prob = JumpProblem(jump_sys, discrete_prob, DirectCR())

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