Show Code
import os
import sys
sys.path.append("..")
import report_utils as ru
from IPython.display import display, HTML
display(HTML(f"<div style='margin-bottom: 20px;'>{ru.nav_bar('timing')}</div>"))Measles Runtime & Performance Benchmark
Execution Speed, Acceleration, and Scalability: R pomp vs. pypomp
Introduction
The London measles benchmark is on the computationally intensive end of the benchmark suite: 365 Euler time steps per year across 14 years of weekly data, with an rproc transition drawing 8 random variates per step from multiple distributions (1 Poisson, 1 gamma, 6 binomial). This benchmark gives us a sense of how pypomp performs when rproc is computationally intensive. Furthermore, it lets us examine how fast pypomp.random custom samplers are relative to stock jax.random in a realistic modeling scenario, instead of in a vacuum.
We benchmark 4 configurations to measure relative execution speeds and isolate the contribution of custom random variate algorithms:
| Configuration | What it isolates |
|---|---|
R pomp (36 CPU cores) |
The multi-threaded CPU reference baseline |
pypomp GPU (fast samplers) |
End-to-end GPU acceleration using pypomp.random |
pypomp GPU (stock JAX samplers) |
Acceleration under stock jax.random samplers |
pypomp CPU (36 cores) |
Multi-core host parallelization without a GPU |
Benchmark Settings & Environment
The table below summarizes algorithmic parameters, software environments, and compute hardware recorded in latest.json for each configuration.
Show Code
PLATFORMS = {
"R pomp (36 cores)": os.path.join("results", "R"),
"pypomp (GPU)": os.path.join("results", "gpu"),
"pypomp (GPU, JAX samplers)": os.path.join("results", "gpu_jax"),
"pypomp (CPU, 36 cores)": os.path.join("results", "cpu"),
}
runs = ru.load_timing_data(PLATFORMS)
display(HTML(ru.build_settings_comparison_html(runs, is_panel=False)))| Setting / Parameter | R pomp (36 cores) | pypomp (GPU) | pypomp (GPU, JAX samplers) | pypomp (CPU, 36 cores) |
|---|---|---|---|---|
| Algorithmic & Workload Settings | ||||
| Run Level | 4 | 4 | 4 | 4 |
| Starting Searches ($N_{starts}$) | 36 | 36 | 36 | 36 |
| IF2 Iterations ($N_{iter}$) | 100 | 100 | 100 | 100 |
| IF2 Particles ($N_p$) | 5,000 | 5,000 | 5,000 | 5,000 |
| Pfilter Particles ($N_{p,eval}$) | 5,000 | 5,000 | 5,000 | 5,000 |
| Evaluation Replicates ($N_{reps}$) | 36 | 36 | 36 | 36 |
| Unit(s) | London | London | London | London |
| Sampler Implementation | Compiled C snippets | fast (pypomp.random) | stock JAX (jax.random) | fast (pypomp.random) |
| Random Seed | — | 594709947 | 594709947 | 594709947 |
| Software & Environment | ||||
| Pomp Framework | pomp 6.3 | pypomp 1.0.0rc1 | pypomp 1.0.0rc1 | pypomp 1.0.0rc1 |
| Backend / Engine | R 4.4.0 | JAX 0.11.1 | JAX 0.11.1 | JAX 0.11.1 |
| Quant Git Commit | ea2b3d1 |
630cc23 |
630cc23 |
630cc23 |
| Run Timestamp | 2026-08-11 12:22:19 | 2026-09-02 19:38:21 | 2026-09-02 20:18:51 | 2026-09-02 21:21:46 |
| Hardware & Compute | ||||
| Compute Device | Intel(R) Xeon(R) Gold 6154 CPU @ 3.00GHz (36 cores) | NVIDIA RTX PRO 6000 Blackwell Server Edition (1 GPU) | NVIDIA RTX PRO 6000 Blackwell Server Edition (1 GPU) | Intel(R) Xeon(R) Gold 6154 CPU @ 3.00GHz (36 cores) |
| Slurm Partition | standard | gpu-rtx6000 | gpu-rtx6000 | standard |
| Slurm Job ID | 56947815 | 59687798 | 59687799 | 59687800 |
Show Code
missing = [(label, r["dir"]) for label, r in runs.items() if not r["available"]]
for label, path in missing:
display(HTML(f"<div class='alert alert-warning'><strong>Missing results for {label}:</strong> Expected at <code>{path}</code>.</div>"))Runtime & Throughput Comparison
The table below compares execution speed across platforms. Speedup and throughput calculations privilege the cold-start particle filter (pfilter_cold) to account for initial execution and JIT compilation overhead incurred in practice. Throughput evaluates parallel acceleration relative to a single R CPU core (\(Speedup \times N_{cores}\)).
Show Code
timing_df = ru.build_timing_comparison_df(runs, is_panel=False)
display(HTML(timing_df.to_html(classes="table table-striped table-hover", index=False)))| Configuration | IF2 (s) | IF2 Speedup | Pfilter (s) | Pfilter Speedup | Total (s) | Total Speedup | Throughput (vs 1 R CPU core) |
|---|---|---|---|---|---|---|---|
| R pomp (36 cores) | 2922.6s (48.71m) | 1.00x | 895.0s (14.92m) | 1.00x | 3817.6s (63.63m) | 1.00x | 36.00x |
| pypomp (GPU) | 47.6s (0.79m) | 61.35x | 19.5s (0.32m) | 45.91x | 67.1s (1.12m) | 56.86x | 2047.04x |
| pypomp (GPU, JAX samplers) | 1370.9s (22.85m) | 2.13x | 517.8s (8.63m) | 1.73x | 1888.7s (31.48m) | 2.02x | 72.77x |
| pypomp (CPU, 36 cores) | 4483.8s (74.73m) | 0.65x | 1544.4s (25.74m) | 0.58x | 6028.2s (100.47m) | 0.63x | 22.80x |
JIT Compilation & Overhead Breakdown
The first pfilter call in JAX incurs one-time JIT compilation overhead. Subsequent warm evaluations run without recompilation.
Show Code
cold_warm_df = ru.build_cold_vs_warm_df(runs)
display(HTML(cold_warm_df.to_html(classes="table table-striped table-hover", index=False)))| Configuration | Pfilter Cold (s) | Pfilter Warm (s) | Compilation Overhead (s) |
|---|---|---|---|
| R pomp (36 cores) | 895.03s | 895.03s | 0.00s |
| pypomp (GPU) | 19.50s | 14.29s | 5.21s |
| pypomp (GPU, JAX samplers) | 517.79s | 513.24s | 4.55s |
| pypomp (CPU, 36 cores) | 1544.38s | 1547.22s | — |
Performance Expectations & Takeaways
- GPU (Fast Samplers):
pypomp’s throughput on GPU should be thousands of times larger than Rpompwith 1 core. The fast-sampler GPU run should also substantially outperform the stock JAX sampler GPU run. - CPU: As it turns out, it’s difficult to quickly generate Poisson, gamma, and binomial random variates in JAX on a CPU. As such, the
pypomprun using 36 CPU cores is substantially slower than the Rpomprun using 36 CPU cores. This should be fine given that GPUs are far more preferable regardless. We leave this comparison here for completeness and to motivate a solution if anyone feels inclined.