Understanding jump model preparation cost

Specific Domains Modelling & Simulations
, ,
1

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:

total-times
total-times1800×1350 71.2 KB

(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:

cost-per-link
cost-per-link1800×1350 138 KB

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 functions one-by-one.

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 is JumpProblem mutable in the first place, isn’t the whole point of having an integrator object to encapsulate the iteration and execution context?)
  • Why is preparation time super-linear in the number of species + links (which determine the complexity and size of the rate functions)?
  • 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 rate functions all individually optimized or something like that?
  • Given that subsequent initializations are faster despite the preceding JumpProblem deepcopy: 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? Otherwise, could their compilation be multi-threaded? (At that point I guess they would be independent of each other, right?)