Skip to content

Commit 79876e8

Browse files
Merge pull request #181 from sivasathyaseeelan/sciml-final-blog
Add SciML Fellowship development 2025 - JumpProcesses.jl blog post
2 parents ec675f9 + 6ea33b9 commit 79876e8

File tree

1 file changed

+311
-0
lines changed

1 file changed

+311
-0
lines changed
Lines changed: 311 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,311 @@
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+
![image](../../../../assets/gsoc2025pics/siva.jpeg)
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

Comments
 (0)