SimpleExplicitTauLeaping 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)

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

prob_disc = DiscreteProblem(u0, tspan, p)

# MassActionJump formulation for SimpleExplicitTauLeaping
reactant_stoich = [[1=>1, 2=>1], [2=>1], Pair{Int,Int}[]]
net_stoich = [[1=>-1, 2=>1], [2=>-1, 3=>1], [1=>1]]
param_idxs = [1, 2, 3]
maj = MassActionJump(reactant_stoich, net_stoich; param_idxs=param_idxs)
jump_prob_maj = JumpProblem(prob_disc, PureLeaping(), maj; rng=rng)

# Solve with SimpleExplicitTauLeaping
sol = solve(EnsembleProblem(jump_prob_maj), SimpleExplicitTauLeaping(), EnsembleSerial(); trajectories=Nsims)

plot(sol)

SimpleExplicitTauLeaping 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]

@sivasathyaseeelan this PR needs to be updated for the changes I just merged to master.

[isaacsas]

@sivasathyaseeelan getting better. I've only reviewed the tau-selection code -- once that seems to be the right algorithm / approach I'll give you a more thorough review.

[isaacsas]

Suggested change

	order = sum(stoch for (spec_idx, stoch) in reactant_stoch[j]; init=0)
	order = sum(stoch for (spec_idx, stoch) in reactant_stoch[j])

sum can figure this out and avoid type assumptions.

[isaacsas]

Suggested change

	hor = zeros(Int, numjumps)
	hor = zeros(Int, numjumps)

Instead of Int, extract the type from reactant_stoch via a parametric input to this function.

[isaacsas]

Suggested change

	function compute_hor(reactant_stoch, numjumps)
	# Compute the highest order of reaction (HOR) for each reaction j, as per Cao et al. (2006), Section IV.
	# HOR is the sum of stoichiometric coefficients of reactants in reaction j.
	# Compute the highest order of reaction (HOR) for each reaction j, as per Cao et al. (2006), Section IV.
	# HOR is the sum of stoichiometric coefficients of reactants in reaction j.
	function compute_hor(reactant_stoch, numjumps)

And please update other functions accordingly. Place these kind of design comments before the function, not right inside it.

[isaacsas]

Don't assume Int, just match the eltype(hor).

[isaacsas]

Suggested change

	function compute_gi(u, max_hor, max_stoich, i, t)
	# Compute g_i for species i to bound the relative change in propensity functions,
	# as per Cao et al. (2006), Section IV, equation (27).
	# g_i is determined by the highest order of reaction (HOR) and maximum stoichiometry (nu_ij) where species i is a reactant:
	# - HOR = 1 (first-order, e.g., S_i -> products): g_i = 1
	# - HOR = 2 (second-order):
	# - nu_ij = 1 (e.g., S_i + S_k -> products): g_i = 2
	# - nu_ij = 2 (e.g., 2S_i -> products): g_i = 2 + 1/(x_i - 1)
	# - HOR = 3 (third-order):
	# - nu_ij = 1 (e.g., S_i + S_k + S_m -> products): g_i = 3
	# - nu_ij = 2 (e.g., 2S_i + S_k -> products): g_i = (3/2) * (2 + 1/(x_i - 1))
	# - nu_ij = 3 (e.g., 3S_i -> products): g_i = 3 + 1/(x_i - 1) + 2/(x_i - 2)
	# Uses precomputed max_hor and max_stoich to reduce work to O(num_species) per timestep.
	# Compute g_i for species i to bound the relative change in propensity functions,
	# as per Cao et al. (2006), Section IV, equation (27).
	# g_i is determined by the highest order of reaction (HOR) and maximum stoichiometry (nu_ij) where species i is a reactant:
	# - HOR = 1 (first-order, e.g., S_i -> products): g_i = 1
	# - HOR = 2 (second-order):
	# - nu_ij = 1 (e.g., S_i + S_k -> products): g_i = 2
	# - nu_ij = 2 (e.g., 2S_i -> products): g_i = 2 + 1/(x_i - 1)
	# - HOR = 3 (third-order):
	# - nu_ij = 1 (e.g., S_i + S_k + S_m -> products): g_i = 3
	# - nu_ij = 2 (e.g., 2S_i + S_k -> products): g_i = (3/2) * (2 + 1/(x_i - 1))
	# - nu_ij = 3 (e.g., 3S_i -> products): g_i = 3 + 1/(x_i - 1) + 2/(x_i - 2)
	# Uses precomputed max_hor and max_stoich to reduce work to O(num_species) per timestep.
	function compute_gi(u, max_hor, max_stoich, i, t)

[isaacsas]

Suggested change

saveat_times = nothing

Not needed.

[isaacsas]

For type stability reasons, this while loop should probably be a separate function you call.

[isaacsas]

Suggested change

	counts .= pois_rand.(rng, max.(rate_cache * tau, 0.0))
	counts .= pois_rand.(rng, max.(rate_cache * tau, zero(tau)))

[isaacsas]

In fact, if a particular rate is <= 0 why not just set the count to zero directly and avoid the call to pois_rand for it? That seems a better approach than using max in this way.

[isaacsas]

Shouldn't the step rejection if any component of u_new is negative make this unneeded? i.e. at this point don't you know that all the entries in u_new are non-negative?

[isaacsas]

This is allocating... Have separate pre-declatred vectors for u_new and u_current and broadcast here.