|
| 1 | +@def rss_pubdate = Date(2025,9,12) |
| 2 | +@def rss = """SciML Fellowship development 2025 - JumpProcesses.jl""" |
| 3 | +@def published = " 12 September 2025 " |
| 4 | +@def title = "SciML Fellowship development 2025 - JumpProcesses.jl" |
| 5 | +@def authors = """<a href="https://github.com/sivasathyaseeelan">Siva Sathyaseelan D N</a>""" |
| 6 | + |
| 7 | +# SciML Fellowship development 2025 - JumpProcesses.jl: Introducing vr_aggregator for VariableRateJumps, GPU-enhanced SimpleTauLeaping and Extending with τ-Leap Algorithms for Jump Process Simulation |
| 8 | + |
| 9 | + |
| 10 | + |
| 11 | + Hello! I’m [Siva Sathyaseelan D N](https://github.com/sivasathyaseeelan), a pre-final year B.Tech + M.Tech Industrial Chemistry student at IIT BHU, Varanasi, India. Over the summer of 2025, as part of the SciML Fellowship, I made significant enhancements to `JumpProcesses.jl`, a Julia package for simulating stochastic jump processes. My contributions focused on three key areas: the introduction of the `vr_aggregator` framework for handling variable rate jump processes, the development of a GPU-enhanced `SimpleTauLeaping` ensemble solver, and the extension of JumpProcesses.jl with new `τ-leaping algorithms` for efficient jump process simulations. These advancements improve performance, scalability, and usability for stochastic simulations, particularly in applications like chemical reaction networks, epidemiological models, and other systems with discrete state changes. |
| 12 | + |
| 13 | +# Variable Rate Aggregator (vr_aggregator) |
| 14 | +The `vr_aggregator` framework introduces a robust and flexible approach to handling variable rate jump processes, addressing limitations in previous implementations and enabling efficient simulations for systems with time-varying rates. |
| 15 | + |
| 16 | +## Key Contributions |
| 17 | +- **VR_Direct and VR_FRM ([PR #477](https://github.com/SciML/JumpProcesses.jl/pull/477), Merged)**: Introduced `vr_aggregator` with `VR_Direct` (Variable Rate Direct with Constant Bounds) and `VR_FRM` (Variable Rate Forward Rate Mode) for optimized variable rate handling. |
| 18 | +- **Optimized VR_DirectEventCache ([PR #486](https://github.com/SciML/JumpProcesses.jl/pull/486), Merged)**: Removed redundant `VariableRateJumps` structures for improved memory efficiency. |
| 19 | +- **VR_DirectFW Aggregator ([PR #488](https://github.com/SciML/JumpProcesses.jl/pull/488), Merged)**: Added forward-time optimized aggregator for frequent rate updates. |
| 20 | +- **Bug Fixes ([PR #489](https://github.com/SciML/JumpProcesses.jl/pull/489), Merged)**: Enhanced stability and reliability for edge cases. |
| 21 | +- **Benchmarking ([PR #1230](https://github.com/SciML/SciMLBenchmarks.jl/pull/1230), Open)**: Ongoing performance evaluations in `SciMLBenchmarks.jl`, showing significant improvements. |
| 22 | + |
| 23 | +## Practical Examples |
| 24 | + |
| 25 | +### Example 1: Variable Rate Jump Process with ODE |
| 26 | +This example demonstrates a simple ODE with two variable rate jumps, solved using different `vr_aggregator` methods (`VR_FRM`, `VR_Direct`, `VR_DirectFW`). |
| 27 | + |
| 28 | +```julia |
| 29 | +using JumpProcesses, DiffEqBase, StableRNGs |
| 30 | + |
| 31 | +# Set random seed for reproducibility |
| 32 | +rng = StableRNG(12345) |
| 33 | + |
| 34 | +# Define variable rate jump |
| 35 | +rate(u, p, t) = u[1] # State-dependent rate |
| 36 | +affect!(integrator) = (integrator.u[1] = integrator.u[1] / 2) # Halve the state |
| 37 | +jump = VariableRateJump(rate, affect!, interp_points=1000) |
| 38 | +jump2 = deepcopy(jump) |
| 39 | + |
| 40 | +# Define ODE |
| 41 | +f(du, u, p, t) = (du[1] = u[1]) # Exponential growth |
| 42 | +prob = ODEProblem(f, [0.2], (0.0, 10.0)) |
| 43 | + |
| 44 | +# Solve with different vr_aggregators |
| 45 | +jump_prob_vrfrm = JumpProblem(prob, jump, jump2; vr_aggregator=VR_FRM(), rng) |
| 46 | +sol_vrfrm = solve(jump_prob_vrfrm, Tsit5()) |
| 47 | + |
| 48 | +jump_prob_vrdirect = JumpProblem(prob, jump, jump2; vr_aggregator=VR_Direct(), rng) |
| 49 | +sol_vrdirect = solve(jump_prob_vrdirect, Tsit5()) |
| 50 | + |
| 51 | +jump_prob_vrdirectfw = JumpProblem(prob, jump, jump2; vr_aggregator=VR_DirectFW(), rng) |
| 52 | +sol_vrdirectfw = solve(jump_prob_vrdirectfw, Tsit5()) |
| 53 | +``` |
| 54 | + |
| 55 | +This example shows how to use different `vr_aggregator` methods to simulate a system where the state decreases by half at random times, driven by a state-dependent rate. |
| 56 | + |
| 57 | +### Example 2: Ensemble Simulation with Birth and Death Jumps |
| 58 | +This example runs an ensemble simulation of an ODE with birth and death jumps, computing the mean state across trajectories. |
| 59 | + |
| 60 | +```julia |
| 61 | +using JumpProcesses, DiffEqBase, StableRNGs, Statistics |
| 62 | + |
| 63 | +# Set random seed for reproducibility |
| 64 | +rng = StableRNG(12345) |
| 65 | + |
| 66 | +# Define ODE (no continuous dynamics) |
| 67 | +function ode_fxn(du, u, p, t) |
| 68 | + du .= 0 |
| 69 | + nothing |
| 70 | +end |
| 71 | + |
| 72 | +# Define birth jump: ∅ → X |
| 73 | +b_rate(u, p, t) = u[1] * p[1] # Birth rate proportional to state |
| 74 | +birth!(integrator) = (integrator.u[1] += 1; nothing) |
| 75 | +b_jump = VariableRateJump(b_rate, birth!) |
| 76 | + |
| 77 | +# Define death jump: X → ∅ |
| 78 | +d_rate(u, p, t) = u[1] * p[2] # Death rate proportional to state |
| 79 | +death!(integrator) = (integrator.u[1] -= 1; nothing) |
| 80 | +d_jump = VariableRateJump(d_rate, death!) |
| 81 | + |
| 82 | +# Set up problem |
| 83 | +u0 = [1.0] # Initial state |
| 84 | +tspan = (0.0, 4.0) |
| 85 | +p = [2.0, 1.0] # Birth and death rates |
| 86 | +prob = ODEProblem(ode_fxn, u0, tspan, p) |
| 87 | + |
| 88 | +# Run ensemble simulation |
| 89 | +Nsims = 1000 |
| 90 | +jump_prob = JumpProblem(prob, b_jump, d_jump; vr_aggregator=VR_FRM(), rng) |
| 91 | +ensemble = EnsembleProblem(jump_prob) |
| 92 | +sol = solve(ensemble, Tsit5(), trajectories=Nsims, save_everystep=false) |
| 93 | + |
| 94 | +# Compute mean state at end |
| 95 | +mean_state = mean(sol.u[i][1] for i in 1:Nsims) |
| 96 | +println("Mean state at t=$(tspan[2]): $mean_state") |
| 97 | +``` |
| 98 | + |
| 99 | +This example models a birth-death process and computes the average state across 1000 trajectories, showcasing the use of `vr_aggregator` in ensemble simulations. |
| 100 | + |
| 101 | +# GPU-Enhanced SimpleTauLeaping Ensemble Solver |
| 102 | + |
| 103 | +The `SimpleTauLeaping` solver leverages GPU parallelism to accelerate large-scale ensemble simulations by approximating multiple jump events in a single time step, significantly reducing computational cost for systems with high event rates. |
| 104 | + |
| 105 | +## Key Contributions |
| 106 | +- **GPU Kernels for SimpleTauLeaping ([PR #490](https://github.com/SciML/JumpProcesses.jl/pull/490), Merged)**: Implemented parallelized leap steps for high performance on NVIDIA GPUs. |
| 107 | +- **GPU-Compatible Poisson Sampling ([PR #512](https://github.com/SciML/JumpProcesses.jl/pull/512), Merged; [PR #54](https://github.com/SciML/PoissonRandom.jl/pull/54), Merged)**: Optimized Poisson random number generation for GPU-based jump events, enhancing simulation efficiency. |
| 108 | + |
| 109 | +GPU enhancements parallelize these steps across thousands of trajectories, achieving up to 10x speedups for large ensemble simulations on NVIDIA GPUs. |
| 110 | + |
| 111 | +## Practical Examples |
| 112 | + |
| 113 | +### Example 1: SIR Model with Influx |
| 114 | + |
| 115 | +This example simulates an SIR epidemiological model with an influx of susceptible individuals, comparing GPU and serial implementations of `SimpleTauLeaping`. |
| 116 | + |
| 117 | +```julia |
| 118 | +using JumpProcesses, DiffEqBase, StableRNGs, Statistics, CUDA |
| 119 | + |
| 120 | +# Set random seed for reproducibility |
| 121 | +rng = StableRNG(12345) |
| 122 | + |
| 123 | +# Define SIR model parameters |
| 124 | +β = 0.1 / 1000.0 # Infection rate |
| 125 | +ν = 0.01 # Recovery rate |
| 126 | +influx_rate = 1.0 # Influx rate |
| 127 | +p = (β, ν, influx_rate) |
| 128 | + |
| 129 | +# Define jump rates |
| 130 | +regular_rate = (out, u, p, t) -> begin |
| 131 | + out[1] = p[1] * u[1] * u[2] # β*S*I (infection) |
| 132 | + out[2] = p[2] * u[2] # ν*I (recovery) |
| 133 | + out[3] = p[3] # influx_rate |
| 134 | +end |
| 135 | + |
| 136 | +# Define state changes |
| 137 | +regular_c = (dc, u, p, t, counts, mark) -> begin |
| 138 | + dc .= 0.0 |
| 139 | + dc[1] = -counts[1] + counts[3] # S: -infection + influx |
| 140 | + dc[2] = counts[1] - counts[2] # I: +infection - recovery |
| 141 | + dc[3] = counts[2] # R: +recovery |
| 142 | +end |
| 143 | + |
| 144 | +# Set up problem |
| 145 | +u0 = [999.0, 10.0, 0.0] # S, I, R |
| 146 | +tspan = (0.0, 250.0) |
| 147 | +prob_disc = DiscreteProblem(u0, tspan, p) |
| 148 | +rj = RegularJump(regular_rate, regular_c, 3) |
| 149 | +jump_prob = JumpProblem(prob_disc, PureLeaping(), rj) |
| 150 | + |
| 151 | +# Run ensemble simulations |
| 152 | +Nsims = 100_000 |
| 153 | +sol_gpu = solve(EnsembleProblem(jump_prob), SimpleTauLeaping(), |
| 154 | + EnsembleGPUKernel(CUDABackend()); trajectories=Nsims, dt=1.0) |
| 155 | +mean_gpu = mean(sol.u[i][1, end] for i in 1:Nsims) |
| 156 | + |
| 157 | +sol_serial = solve(EnsembleProblem(jump_prob), SimpleTauLeaping(), |
| 158 | + EnsembleSerial(); trajectories=Nsims, dt=1.0) |
| 159 | +mean_serial = mean(sol.u[i][1, end] for i in 1:Nsims) |
| 160 | +println("GPU mean (S at t=$(tspan[2])): $mean_gpu, Serial mean: $mean_serial") |
| 161 | +``` |
| 162 | + |
| 163 | +### Example 2: SEIR Model with Exposed Compartment |
| 164 | +This example simulates an SEIR epidemiological model with an exposed compartment, using `SimpleTauLeaping` with GPU acceleration. |
| 165 | + |
| 166 | +```julia |
| 167 | +using JumpProcesses, DiffEqBase, StableRNGs, Statistics, CUDA |
| 168 | + |
| 169 | +# Set random seed for reproducibility |
| 170 | +rng = StableRNG(12345) |
| 171 | + |
| 172 | +# Define SEIR model parameters |
| 173 | +β = 0.3 / 1000.0 # Infection rate |
| 174 | +σ = 0.2 # Progression rate |
| 175 | +ν = 0.01 # Recovery rate |
| 176 | +p = (β, σ, ν) |
| 177 | + |
| 178 | +# Define jump rates |
| 179 | +regular_rate = (out, u, p, t) -> begin |
| 180 | + out[1] = p[1] * u[1] * u[3] # β*S*I (infection) |
| 181 | + out[2] = p[2] * u[2] # σ*E (progression) |
| 182 | + out[3] = p[3] * u[3] # ν*I (recovery) |
| 183 | +end |
| 184 | + |
| 185 | +# Define state changes |
| 186 | +regular_c = (dc, u, p, t, counts, mark) -> begin |
| 187 | + dc .= 0.0 |
| 188 | + dc[1] = -counts[1] # S: -infection |
| 189 | + dc[2] = counts[1] - counts[2] # E: +infection - progression |
| 190 | + dc[3] = counts[2] - counts[3] # I: +progression - recovery |
| 191 | + dc[4] = counts[3] # R: +recovery |
| 192 | +end |
| 193 | + |
| 194 | +# Set up problem |
| 195 | +u0 = [999.0, 0.0, 10.0, 0.0] # S, E, I, R |
| 196 | +tspan = (0.0, 250.0) |
| 197 | +prob_disc = DiscreteProblem(u0, tspan, p) |
| 198 | +rj = RegularJump(regular_rate, regular_c, 3) |
| 199 | +jump_prob = JumpProblem(prob_disc, PureLeaping(), rj; rng=StableRNG(12345)) |
| 200 | + |
| 201 | +# Run ensemble simulations |
| 202 | +Nsims = 100_000 |
| 203 | +sol_gpu = solve(EnsembleProblem(jump_prob), SimpleTauLeaping(), |
| 204 | + EnsembleGPUKernel(CUDABackend()); trajectories=Nsims, dt=1.0) |
| 205 | +mean_gpu = mean(sol.u[i][end, end] for i in 1:Nsims) |
| 206 | + |
| 207 | +sol_serial = solve(EnsembleProblem(jump_prob), SimpleTauLeaping(), |
| 208 | + EnsembleSerial(); trajectories=Nsims, dt=1.0) |
| 209 | +mean_serial = mean(sol.u[i][end, end] for i in 1:Nsims) |
| 210 | +println("GPU mean (R at t=$(tspan[2])): $mean_gpu, Serial mean: $mean_serial") |
| 211 | +``` |
| 212 | + |
| 213 | +# Extending JumpProcesses.jl with τ-Leap Algorithms |
| 214 | + |
| 215 | +The introduction of new `τ-leaping algorithms` extends JumpProcesses.jl to provide flexible and efficient methods for simulating jump processes, catering to diverse system dynamics, from non-stiff to highly stiff systems. |
| 216 | + |
| 217 | +## Key Contributions |
| 218 | +- **SimpleExplicitTauLeaping ([PR #513](https://github.com/SciML/JumpProcesses.jl/pull/513), Open)**: An explicit solver for non-stiff systems, offering fast simulations where stability is not a concern. |
| 219 | +- **SimpleImplicitTauLeaping ([PR #500](https://github.com/SciML/JumpProcesses.jl/pull/500), Open)**: An implicit solver for stiff jump processes, ensuring stability in systems with large rate disparities. |
| 220 | +- **SimpleAdaptiveTauLeaping ([PR #524](https://github.com/SciML/JumpProcesses.jl/pull/524), Open)**: An adaptive solver that dynamically adjusts leap sizes to balance accuracy and performance. |
| 221 | + |
| 222 | +## Practical Examples |
| 223 | + |
| 224 | +### Example 1: Birth-Death Process with Explicit τ-Leaping |
| 225 | +This example simulates a non-stiff birth-death process using `SimpleExplicitTauLeaping`, suitable for systems with balanced rates. |
| 226 | + |
| 227 | +```julia |
| 228 | +using JumpProcesses, DiffEqBase, StableRNGs, Plots |
| 229 | + |
| 230 | +# Set random seed for reproducibility |
| 231 | +rng = StableRNG(12345) |
| 232 | + |
| 233 | +# Define birth-death system |
| 234 | +function define_birth_death_system() |
| 235 | + c = (0.5, 0.5) # Birth and death rates |
| 236 | + # Reaction 1: Birth (∅ → X) |
| 237 | + reactant_stoich1 = [] # No reactants consumed |
| 238 | + net_stoich1 = [Pair(1, 1)] # X +1 |
| 239 | + # Reaction 2: Death (X → ∅) |
| 240 | + reactant_stoich2 = [Pair(1, 1)] # X consumed |
| 241 | + net_stoich2 = [Pair(1, -1)] # X -1 |
| 242 | + jumps = MassActionJump([c[1], c[2]], [reactant_stoich1, reactant_stoich2], |
| 243 | + [net_stoich1, net_stoich2]) |
| 244 | + return jumps |
| 245 | +end |
| 246 | + |
| 247 | +# Set up problem |
| 248 | +u0 = [10] # Initial population |
| 249 | +tspan = (0.0, 10.0) |
| 250 | +prob = DiscreteProblem(u0, tspan, nothing) |
| 251 | +jump_prob = JumpProblem(prob, PureLeaping(), define_birth_death_system()) |
| 252 | + |
| 253 | +# Solve with SimpleExplicitTauLeaping |
| 254 | +sol_explicit = solve(jump_prob, SimpleExplicitTauLeaping(); saveat=0.01) |
| 255 | +plot(sol_explicit, label="Population", title="Explicit Tau-Leaping: Birth-Death Process") |
| 256 | +``` |
| 257 | + |
| 258 | +### Example 2: Stiff System with Implicit and Adaptive τ-Leaping |
| 259 | +This example simulates a stiff chemical reaction system from Cao et al. (2007) using `SimpleImplicitTauLeaping` and `SimpleAdaptiveTauLeaping`. |
| 260 | + |
| 261 | +```julia |
| 262 | +using JumpProcesses, DiffEqBase, StableRNGs, Plots |
| 263 | + |
| 264 | +# Set random seed for reproducibility |
| 265 | +rng = StableRNG(12345) |
| 266 | + |
| 267 | +# Define stiff system (Cao et al., 2007) |
| 268 | +function define_stiff_system() |
| 269 | + c = (1000.0, 1000.0, 1.0) # Rate constants |
| 270 | + # Reaction 1: S1 → S2 |
| 271 | + reactant_stoich1 = [Pair(1, 1)] # S1 consumed |
| 272 | + net_stoich1 = [Pair(1, -1), Pair(2, 1)] # S1 -1, S2 +1 |
| 273 | + # Reaction 2: S2 → S1 |
| 274 | + reactant_stoich2 = [Pair(2, 1)] # S2 consumed |
| 275 | + net_stoich2 = [Pair(1, 1), Pair(2, -1)] # S1 +1, S2 -1 |
| 276 | + # Reaction 3: S2 → S3 |
| 277 | + reactant_stoich3 = [Pair(2, 1)] # S2 consumed |
| 278 | + net_stoich3 = [Pair(2, -1), Pair(3, 1)] # S2 -1, S3 +1 |
| 279 | + jumps = MassActionJump([c[1], c[2], c[3]], [reactant_stoich1, reactant_stoich2, reactant_stoich3], |
| 280 | + [net_stoich1, net_stoich2, net_stoich3]) |
| 281 | + return jumps |
| 282 | +end |
| 283 | + |
| 284 | +# Set up problem |
| 285 | +u0 = [100, 0, 0] # S1=100, S2=0, S3=0 |
| 286 | +tspan = (0.0, 5.0) |
| 287 | +prob = DiscreteProblem(u0, tspan, nothing) |
| 288 | +jump_prob = JumpProblem(prob, PureLeaping(), define_stiff_system()) |
| 289 | + |
| 290 | +# Solve with SimpleImplicitTauLeaping |
| 291 | +sol_implicit = solve(jump_prob, SimpleImplicitTauLeaping(); saveat=0.01) |
| 292 | +plot(sol_implicit, label=["S1" "S2" "S3"], title="Implicit Tau-Leaping") |
| 293 | + |
| 294 | +# Solve with SimpleAdaptiveTauLeaping |
| 295 | +sol_adaptive = solve(jump_prob, SimpleAdaptiveTauLeaping(); saveat=0.001) |
| 296 | +plot(sol_adaptive, label=["S1" "S2" "S3"], title="Adaptive Tau-Leaping") |
| 297 | +``` |
| 298 | + |
| 299 | +# Future Work Includes |
| 300 | + |
| 301 | +These contributions lay the foundation for further advancements in JumpProcesses.jl: |
| 302 | + |
| 303 | +- **Merging τ-Leap Algorithms**: Finalize and merge SimpleExplicitTauLeaping, SimpleImplicitTauLeaping, and SimpleAdaptiveTauLeaping to make them available to users. |
| 304 | +- **Benchmarking**: Complete performance evaluations for τ-leaping algorithms in SciMLBenchmarks.jl to quantify their benefits. |
| 305 | +- **Documentation**: Develop comprehensive documentation, including usage guides and performance tips for τ-leaping algorithms. |
| 306 | +- **GPU-Enhanced SSAStepper**: Extend GPU support to the SSAStepper for exact stochastic ensemble simulations, complementing the approximate τ-leaping methods. |
| 307 | + |
| 308 | +These developments enhance the SciML ecosystem’s capabilities for stochastic modeling, providing researchers with powerful tools for simulating complex systems. |
| 309 | + |
| 310 | +# Acknowledgments |
| 311 | +Thank you [Professor Dr. Samuel A. Isaacson](https://github.com/isaacsas) and [Dr. Christopher Rackauckas](https://github.com/ChrisRackauckas), for your incredible guidance and support throughout the SciML Fellowship program. Your expertise, encouragement, and constructive feedback during this fellowship is invaluable in helping me navigate the complexities of the project. Your mentorship not only enhanced my technical skills but also inspired me to grow as a contributor to the open-source community. Thank you for your time, patience, and dedication to my learning journey! |
0 commit comments