SimpleExplicitTauLeaping solver by sivasathyaseeelan · Pull Request #513 · SciML/JumpProcesses.jl

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.

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.

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.

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

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

isaacsas · 2026-01-17T12:32:44Z

@sivasathyaseeelan thanks for coming back to this! Let me know when you are done updating.

isaacsas · 2026-01-26T15:37:46Z

@sivasathyaseeelan if tests pass now I'll merge. Thanks again for all your efforts on this one!

Let me know if you return to the implicit one and I can review it once you feel you are done when I have a chance.

isaacsas · 2026-01-26T18:39:37Z

@sivasathyaseeelan seems like the tau-leaping tests are stalling out on 1.10. Have you tested their locally?

sivasathyaseeelan · 2026-01-26T18:46:50Z

@sivasathyaseeelan seems like the tau-leaping tests are stalling out on 1.10. Have you tested their locally?

Yes I did, all tests were passing in my local machine

isaacsas · 2026-01-26T19:13:01Z

And was that in Julia 1.10?

isaacsas · 2026-01-26T19:13:13Z

(i.e. not 1.11 or 1.12)

isaacsas · 2026-01-26T19:13:48Z

Let's see if restarting CI fixes it.

sivasathyaseeelan · 2026-01-26T19:42:17Z

@isaacsas I think the CI takes too much time to run, https://github.com/SciML/JumpProcesses.jl/actions/runs/21111444469/job/60710583077 here it took more than 5 hrs

isaacsas · 2026-01-26T19:46:13Z

Right, but that seems to only be an issue in this PR from what I can see. For other PRs it is much less time.

That suggests there is an issue in this PR or the tests when applied to the new code on 1.10.

sivasathyaseeelan · 2026-01-26T19:50:56Z

Right, but that seems to only be an issue in this PR from what I can see. For other PRs it is much less time.

That suggests there is an issue in this PR or the tests when applied to the new code on 1.10.

How can I fix this

isaacsas · 2026-01-26T19:53:27Z

Fixing this requires some debugging work to see why it is occurring, which means figuring out in that test file what is taking so long on 1.10 vs. 1.11. The first step would be trying to run those tests line by line on Julia 1.10 and 1.11 to see if there is somewhere it hangs / takes forever (or use something like Claude Code to investigate the issue after explaining what is wrong).

isaacsas · 2026-01-26T19:54:34Z

If you aren't able to investigate this I can look into it when I have some time, but probably won't be till later in the week as I am busy with teaching and meetings the next three days.

sivasathyaseeelan · 2026-01-26T21:06:17Z

Thank you @isaacsas. I don't have Claude, But I will try looking into this. If I can't able to debug I will let you know

isaacsas · 2026-01-27T22:54:02Z

@sivasathyaseeelan seems there is a bug in the step rejection logic. See the summary below

Issue: Tau-Leaping Tests Taking 5+ Hours on CI

Problem

The "Tau Leaping Tests" are taking 267 minutes (4h 27m) on CI (Julia 1.10 LTS) but only ~4 seconds locally. This is causing CI failures.

Root Cause

There is a bug in the negative population handling logic in simple_explicit_tau_leaping_loop! at lines 244-248:

if any(<(0), u_new)
    # Halve tau to avoid negative populations, as per Cao et al. (2006), Section 3.3
    tau /= 2
    continue
end

Two Issues:

Halving is ineffective: When continue executes, control returns to line 217 where tau = compute_tau(...) recomputes tau from scratch, overwriting the halved value. The halving has no effect.
No iteration: According to Cao et al. (2005), when negative populations occur, the leap should be rejected and tau halved repeatedly in a loop until either:
- No negative populations occur, OR
- tau falls below dtmin

The current implementation relies on getting lucky with different Poisson samples on each retry (since tau is recomputed to the same value). In CI environments, this can cause many ineffective retries.

isaacsas · 2026-01-29T13:56:22Z

It looks like tests are now passing, which is great! But there is still an issue with the lts tests taking 330m. I'll dig into it when I have a chance, but please feel free to keep looking at it yourself. I don't think we can merge this until we get that straightened out -- that will be too much of a drain on CI time.

sivasathyaseeelan · 2026-02-03T20:01:34Z

It looks like tests are now passing, which is great! But there is still an issue with the lts tests taking 330m. I'll dig into it when I have a chance, but please feel free to keep looking at it yourself. I don't think we can merge this until we get that straightened out -- that will be too much of a drain on CI time.

@isaacsas I fixed the bug. The tests were not properly written and optimized, which was loading the system. Now Tau Leaping Tests are running below 50 sec!

isaacsas · 2026-02-14T14:07:11Z

@sivasathyaseeelan I'm going to merge this and start a new PR to improve the tests.

I know this was a long haul, thanks for your efforts on it!

isaacsas reviewed Aug 10, 2025

View reviewed changes

Comment thread Project.toml Outdated

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