SimpleExplicitTauLeaping solver #513

sivasathyaseeelan · 2025-08-09T23:48:51Z

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.

Project.toml

ChrisRackauckas · 2025-08-10T12:39:22Z

Rebase onto master

isaacsas · 2025-08-11T02:49:08Z

@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.

Project.toml

src/simple_regular_solve.jl

isaacsas

General comments:

Please add comments citing the formulas you are using from the original paper for the leap selection.
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.

src/simple_regular_solve.jl

test/regular_jumps.jl

src/simple_regular_solve.jl

isaacsas · 2025-08-23T17:42:17Z

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

src/simple_regular_solve.jl

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 · 2025-08-28T12:41:12Z

src/simple_regular_solve.jl

+function compute_gi(u, reaction_conditions, i)
+    # 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) 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 reaction_conditions to optimize checks for HOR = 2 or 3 with nu_ij >= 2.
+    max_hor = maximum(isempty(reaction_conditions[i]) ? 0 : [hor_j for (j, nu_ij, hor_j) in reaction_conditions[i]])
+    max_gi = 1
+    for (j, nu_ij, hor_j) in reaction_conditions[i]
+        if hor_j == max_hor
+            if hor_j == 1
+                max_gi = max(max_gi, 1)
+            elseif hor_j == 2
+                if nu_ij == 1
+                    max_gi = max(max_gi, 2)
+                elseif nu_ij == 2
+                    if u[i] > 1  # Ensure x_i - 1 > 0
+                        gi = 2 + 1 / (u[i] - 1)
+                        max_gi = max(max_gi, ceil(Int64, gi))
+                    end
+                end
+            elseif hor_j == 3
+                if nu_ij == 1
+                    max_gi = max(max_gi, 3)
+                elseif nu_ij == 2
+                    if u[i] > 1  # Ensure x_i - 1 > 0
+                        gi = 1.5 * (2 + 1 / (u[i] - 1))
+                        max_gi = max(max_gi, ceil(Int64, gi))
+                    end
+                elseif nu_ij == 3
+                    if u[i] > 2  # Ensure x_i - 2 > 0
+                        gi = 3 + 1 / (u[i] - 1) + 2 / (u[i] - 2)
+                        max_gi = max(max_gi, ceil(Int64, gi))
+                    end
+                end
+            end
+        end
+    end
+    return max_gi


This is getting better, but there are still some issues.

Just initialize max_gi = zero(typeof(t)) and get rid of all these ceil's and conversions to integers. WIth the correction formulas g_i can be a floating point type so I think it makes sense to just give it the same type as the time variable.

You are currently doing num_species * num_rxs work per timestep in computing all these g_i, when it should only be num_species work. You need to precalculate when initializing the solver hor_order, where hor_order[i] is the maximum substrate order of species i over all reactions. Then you use this to determine which g_i case you are in. For example, if HOR(i) == 2 and hor_order[i] == 1 you just set g_i = 2, but if HOR(i) == 2 and hor_order[i] == 2 you use the formula that involves u[i].

Now compute_gi is in O(num_species) time complexity

src/simple_regular_solve.jl

isaacsas · 2025-09-08T14:10:43Z

src/simple_regular_solve.jl

+    # HOR is the sum of stoichiometric coefficients of reactants in reaction j.
+    hor = zeros(Int, numjumps)
+    for j in 1:numjumps
+        order = sum(stoch for (spec_idx, stoch) in reactant_stoch[j]; init=0)


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 · 2025-09-08T14:12:42Z

src/simple_regular_solve.jl

+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.
+    hor = zeros(Int, numjumps)


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 · 2025-09-08T14:14:05Z

src/simple_regular_solve.jl

+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.


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 · 2025-09-08T14:14:37Z

src/simple_regular_solve.jl

+    max_hor = zeros(Int, numspecies)
+    max_stoich = zeros(Int, numspecies)


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

isaacsas · 2025-09-08T14:17:08Z

src/simple_regular_solve.jl

+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.


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 · 2025-09-08T14:53:31Z

src/simple_regular_solve.jl

+    max_hor, max_stoich = precompute_reaction_conditions(reactant_stoch, hor, length(u0), numjumps)
+
+    # Set up saveat_times
+    saveat_times = nothing


Suggested change

saveat_times = nothing

Not needed.

isaacsas · 2025-09-08T14:53:59Z

src/simple_regular_solve.jl

+
+    save_idx = 1
+
+    while t_current < t_end


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

isaacsas · 2025-09-08T14:55:57Z

src/simple_regular_solve.jl

+        if !isempty(saveat_times) && save_idx <= length(saveat_times) && t_current + tau > saveat_times[save_idx]
+            tau = saveat_times[save_idx] - t_current
+        end
+        counts .= pois_rand.(rng, max.(rate_cache * tau, 0.0))


Suggested change

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

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

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 · 2025-09-08T14:59:17Z

src/simple_regular_solve.jl

+        for i in eachindex(u_new)
+            u_new[i] = max(u_new[i], 0)
+        end
+        t_new = t_current + tau


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 · 2025-09-08T15:00:13Z

src/simple_regular_solve.jl

+                end
+            end
+        end
+        u_new = u_current + du


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

isaacsas reviewed Aug 10, 2025

View reviewed changes

Project.toml Outdated Show resolved Hide resolved

sivasathyaseeelan requested review from ChrisRackauckas and isaacsas August 10, 2025 01:05

sivasathyaseeelan force-pushed the adaptive-implicit branch from f266a9d to 5963f3e Compare August 10, 2025 13:10

isaacsas requested changes Aug 11, 2025

View reviewed changes

sivasathyaseeelan changed the title ~~SimpleAdaptiveTauLeaping and SimpleImplicitTauLeaping~~ SimpleAdaptiveTauLeaping solver Aug 12, 2025

sivasathyaseeelan requested a review from isaacsas August 13, 2025 23:33

ChrisRackauckas reviewed Aug 14, 2025

View reviewed changes