SimpleAdaptiveTauLeaping solver

[sivasathyaseeelan]

using JumpProcesses, DiffEqBase
using Test, LinearAlgebra, Statistics
using StableRNGs, Plots
rng = StableRNG(12345)

Nsims = 10

β = 0.1 / 1000.0
ν = 0.01
influx_rate = 1.0
p = (β, ν, influx_rate)

regular_rate = (out, u, p, t) -> begin
    out[1] = p[1] * u[1] * u[2]  # β*S*I (infection)
    out[2] = p[2] * u[2]         # ν*I (recovery)
    out[3] = p[3]                # influx_rate
end

regular_c = (dc, u, p, t, counts, mark) -> begin
    dc .= 0.0
    dc[1] = -counts[1] + counts[3]  # S: -infection + influx
        dc[2] = counts[1] - counts[2]   # I: +infection - recovery
        dc[3] = counts[2]               # R: +recovery
end

u0 = [999.0, 10.0, 0.0]  # S, I, R
tspan = (0.0, 250.0)

prob_disc = DiscreteProblem(u0, tspan, p)
rj = RegularJump(regular_rate, regular_c, 3)
jump_prob = JumpProblem(prob_disc, Direct(), rj)

sol = solve(EnsembleProblem(jump_prob), SimpleAdaptiveTauLeaping(), EnsembleSerial(); trajectories = Nsims)
plot(sol)

SimpleAdaptiveTauLeaping Plot:

Checklist

• Appropriate tests were added
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Additional context

Add any other context about the problem here.

[ChrisRackauckas]

Rebase onto master

[isaacsas]

@sivasathyaseeelan can you make the correctness test more strict. Compare against Direct, which will probably require a smaller timestep for the leaping methods, and calculate the maximum error over time (say every second) instead of just at the final time.

I’m reading through the references to refresh my memory on how the adaptive stepping works, and will try to finish that tomorrow and give you some more feedback afterwards.

[isaacsas]

@sivasathyaseeelan, to speed up review I suggest splitting this into separate PRs for explicit vs. implicit. Let's get explicit merged and then we can work on implicit.

@ChrisRackauckas I think there are some design decisions to be made here for getting at the individual rate functions and stochiometry vectors. So it is just a question if you want to work on that as part of these or as a followup -- I'd defer to what you want, but I don't think these should be advertised until we get that settled and have more extensive testing.

[isaacsas]

General comments:

1. Please add comments citing the formulas you are using from the original paper for the leap selection.
2. Please revise with an eye towards making this non-allocating during the solution timestepping, see the examples I've highlighted but keep in mind they are not all the cases of extraneous allocations.

[isaacsas]

Why not just make this an integer? It should just be an integer right given that hor is integer?

Generally try not to assume Float64 types in your code. They should be inferred from the types of t or the type of the elements of u as appropriate.

[isaacsas]

As I mentioned in the other PR, I suggest switching to the input being MassActionJumps. Then you can get rid of most of this stoichiometry-related code. It will also have a big performance benefit since you won't need to recalculate nu every timestep (which is likely to be very expensive).

[isaacsas]

Once you've made the switch to MassActionJumps you can get the stoichiometry for the ith reaction from net_stoch[i]. It only needs to be extracted once as it doesn't change during a simulation.

[isaacsas]

Doesn't this need to be calculated based on the substrate stoichiometry, not the net stoichiometry? (The order of a reaction is determined by the substrates coefficients, i.e. for $m_1 S_1 + m_2 S_2 \to n_1 S_1 + n_2 S_2$ the order is $m_1 + m_2$.) If you use MassActionJumps you can get these from the reactant_stoch field.

[isaacsas]

Suggested change

	saveat_times = isnothing(saveat) ? Float64[] : saveat isa Number ? collect(range(tspan[1], tspan[2], step=saveat)) : collect(saveat)
	saveat_times = isnothing(saveat) ? Vector{typeof(t))() : saveat isa Number ? collect(range(tspan[1], tspan[2], step=saveat)) : collect(saveat)

This is too long to read / parse, expand it out with an outer if/else block. Try not to nest ternary operators as it gets hard to read the code.

[isaacsas]

Make these the current u and current t, and have separate vectors usave and tsave or such for the history. That is more consistent with SciML naming.

[isaacsas]

Why is this needed? Aren't you directly stepping to saveat times?

[isaacsas]

This isn't necessarily the correct value for $g_i$ though when hor[i] > 1. If hor[i] == 2 you need to precalculate if any reaction for which species i is a substrate has the substrate stoichoimetry of 2 or more (using reactant_stoch for example). Then you need to use the formulate for the updated g_i in case (II). Similar for if hor[i] ==3.

We don't need to support if hor[i] > 3, but the solver should probably error if it detects this is the case.

[isaacsas]

Avoid allocations. Can't these just be scalar variables within the loop over i?

[isaacsas]

These can just be written as dot products using views of nu.

[isaacsas]

This is too loose a bound. At least decrease it to 5%.

[ChrisRackauckas]

Suggested change

	max_order = max(max_order, Float64(hor[j]))
	max_order = max(max_order, float(hor[j]))

[ChrisRackauckas]

Suggested change

	struct SimpleAdaptiveTauLeaping <: DiffEqBase.DEAlgorithm
	epsilon::Float64 # Error control parameter
	struct SimpleAdaptiveTauLeaping{T <: AbstractFloat} <: DiffEqBase.DEAlgorithm
	epsilon::T # Error control parameter