Controlling save behavior and performance of DiffEqBiological

I’m interested in the bistable Schlogl model, which, under resevoir assumptions on species A and B, can be represented inside of DiffEqBiological as

schlogl = @reaction_network begin
    c1 * a/ V, 2*S --> 3*S
    c2 / V^2, 3*S --> 2*S
    c3 * b * V, ∅ --> S
    c4, S --> ∅
    end c1 c2 c3 c4 a b V

I am picking the constants such that this problem is relativley metatstable, and I want to explore the equilibrium statistics of the problem, having run it until Tmax, a sufficiently large time. Thus, I am inclined to try the following code:

V = 25
a = 1
b = 2
c1 = 3 * 2
c2 = 0.6 * 6
c3 = 0.25
c4 = 2.95

params = (c1, c2, c3, c4, a, b, V);
Tmax = 10^3;


schlogl = @reaction_network begin
    c1 * a/ V, 2*S --> 3*S
    c2 / V^2, 3*S --> 2*S
    c3 * b * V, ∅ --> S
    c4, S --> ∅
end c1 c2 c3 c4 a b V

function run_schlogl(Sinit)
    dprob = DiscreteProblem(schlogl, Sinit, (0.0, Tmax), params);
    jprob = JumpProblem(dprob, Direct(), schlogl);
    jsol = solve(jprob, SSAStepper());
    return jsol.u[end];
end

and run run_schlogl on a variety of input data a sufficient number of times to sample the distribution. Notice that I am only making use of the terminal value in order to have independent samples.

Now, the problem I am encountering is that even though I only need the state at Tmax, it is recording many intermediate values, and I see that the memory usage of the problem scales linearly with Tmax:

@btime run_schlogl([5])
  4.768 ms (34676 allocations: 4.23 MiB)

Presumably, I could improve performance if I avoided all these allocation, but I do not see a way to do this; saveat does not seem to work here.

I also notice there are a lot of type instabilities when I use @code_warntype on this. Is that simply an issue with the larger package, or are there ways to clean that up?

do jprob = JumpProblem(dprob, Direct(), schlogl;save_positions = (false,false)). The jumps utilize the event handling so in just the same way, they save by default when firing so that discontinuities are appropriately handled when added to ODEs. But if you don’t need it, turn it off.

At the highest level, you have a type instability for some dynamic behavior in choosing defaults. However, the inner loop should be completely stable. If it isn’t, that is worth an issue.

Thanks, just a quick follow up question: Is there a way to reuse either the DiscreteProblem or the JumpProblem, such that if I change the initial condition, I don’t have to reconstruct these two data structures? In other words, is there a way to make these mutable?

remake(prob,u0=...).

That works for the DiscreteProblem structure with respect to u0, but it doesn’t seem to work with the JumpProblem:

dprob = DiscreteProblem(schlogl, [5], (0.0, Tmax), params);
jprob = JumpProblem(dprob, Direct(), schlogl,save_positions = (false,false));
jprob.prob
jsol=solve(remake(jprob, prob=remake(dprob, u0=[10])), SSAStepper())

generates a massive error that begins with:

AssertionError: length(jump_prob.jump_callback.discrete_callbacks) == 1

This, however, does work:

jprob.prob = remake(dprob, u0=Sinit);
jsol = solve(jprob, SSAStepper());

and saves a number of allocations.

That’s worth a bug report. I don’t think we have a test on remake + JumpProblem. Sorry about that.

added to the issue to github