I am using Catalyst.jl to make JumpProblems of around 160k species and a fairly sparse interaction structure that includes what JumpProcesses.jl calls “constant-rate jump” (CRJ) reactions. This works, but JumpProblem construction is quite slow (and so is integrator initialization). I would like to understand why that is and what I could do about it.
As an (admittedly contrived) example, consider:
using Catalyst
using Graphs
using JumpProcesses
t0 = @elapsed begin
g = static_scale_free(n, 2n, 3.0, 3.0, seed = n)
t = default_t()
xs = map(vertices(g)) do i
name = Symbol("node_$i")
only(@species $name(t))
end
rate(i) = minimum(
[hill(xs[j], 1.0, 1 / j, 1.0) for j in neighbors(g, i)],
init = 1.0,
)
reactions = mapreduce(vcat, vertices(g)) do i
[
Reaction(rate(i), nothing, [xs[i]])
Reaction(1.0, [xs[i]], nothing)
]
end
system = complete(ReactionSystem(reactions, name = :name))
end
t1 = @timed JumpProblem(system, [s => 0 for s in unknowns(system)], (0.0, 1.0))
This builds a directed graph with n nodes, 2n edges and a power-law degree distribution, and then translates it into a reaction network such that the inbound edges are aggregated for each species (node) by a CRJ function (here map hill and reduce by minimum). Timing this construction for various sizes (on a smaller 2019 Intel laptop, no swap) gives:
(I have attached a Pluto notebook (113.9 KB) timing the
JumpProblem call, but not the integrator initialization. Apparently I am not allowed to upload the CSV file with the timings?)
Although this isn’t really a clean benchmark, it still shows what I observed in my actual application: Constructing the model specification is negligibly fast, but the JumpProblem call is very expensive. Without having a good mental model of what exactly it does internally, my expectation would be that it should take time and memory proportionally to the number of links (i.e. here, O(n)). But beyond a certain size, both time spent and memory allocated grow faster than the link count:
Whenever I send SIGUSR1, it is currently doing assemble_crj/compile_equational_affect (for JumpProblem construction) or concretize_affects! (for integrator initialization). Since the timings show that most of that is Julia compiling stuff, I suspect that model preparation actually compiles these many rate update functions one-by-one. (Edit: “affect” actually refers to the state updates, right?)
In my application (which is ~160k species, ~64k CRJ-rate reactions, and ~300k purely mass-action reactions, run on a passively cooled Ryzen 9 HX 370, so effectively laptop hardware) JumpProblem construction takes ~36h, and the first integrator intitialization ~14½h. But each subsequent integrator initialization, despite being produced from a deepcopyd JumpProblem (!), is fast enough that I am unsure whether the delay might actually just have been faster reaction propensities early on. And the whole process takes ~18GiB peak resident memory, most of which seems to be for SciML, let’s call it 16GiB.
Now my questions are:
- Am I using this incorrectly? It feels like I’m missing something here. (Why does creating an integrator mutate the
JumpProblem? Why isJumpProblemmutable in the first place, isn’t the whole point of having an integrator object to encapsulate the iteration and execution context?) - Can I reduce the memory footprint somehow? ~100KiB per species (resp. ~50KiB per specified graph edge) seems quite high since I’m not saving trajectories.
- Why is preparation so slow in absolute terms? Are these
rateupdate functions all individually optimized or something like that? (Edit: “affect” = updates) - Why is preparation time super-linear in the number of species + links (which determine the complexity and size of the
rateupdate functions)? (Edit: I guess the same consideration holds for the updates, i.e. the complexity/size of the state updater would depend on the number of outbound edges?) - Given that subsequent initializations are faster despite the preceding
JumpProblemdeepcopy: What part of the execution state here is global? Do the rate functions all enter Julia’s method table? - If so, will they be evicted at some point, or do I need to avoid processing a sequence of large SciML models in the same Julia process?
- Are there easy ways to accelerate model preparation?
E.g. by predefining/registering the rate functions somewhere, perhaps once for each arity, since they are so similar?(Edit: That probably does not make sense for the affects…) Otherwise, could their compilation be multi-threaded? (At that point I guess they would be independent of each other, right?)