Measles London Model Speed Comparison

R pomp (CPU) vs. pypomp (GPU with fast samplers) vs. pypomp (GPU with JAX samplers)

Published

July 6, 2026

Introduction

This report compares the runtime speed and statistical accuracy of the London measles model using three different setups:

  1. R pomp running in parallel on a CPU (36 cores of a 3.30GHz AMD EPYC 9575F).
  2. pypomp running on a GPU (NVIDIA RTX 6000 Blackwell) using the default fast approximate inverse CDF samplers (fast_poisson, fast_gamma, fast_multinomial).
  3. pypomp running on a GPU (NVIDIA RTX 6000 Blackwell) using the standard built-in JAX samplers (replacing the fast samplers).

Both models were run under identical conditions:

  • IF2: 5,000 particles, 100 iterations, 36 replicates.
  • Pfilter: 5,000 particles, 36 replications per final parameter set.

Timing Comparison

The table below presents the execution runtimes and the speedup factors relative to R pomp. We should see that Pypomp, when run with the pypomp.random fast samplers on a strong GPU, is substantially faster than the other configurations.

Method MIF2 Time MIF2 Speedup Pfilter Time Pfilter Speedup Total Time Total Speedup Throughput (vs 1 CPU Core)
R pomp (CPU, 36 Cores) 1934.73 s (32.25 m) 1.00x 591.20 s (9.85 m) 1.00x 2525.93 s (42.10 m) 1.00x 36.00x
pypomp (GPU, Fast Samplers) 48.96 s (0.82 m) 39.51x 18.90 s (0.31 m) 31.29x 67.86 s (1.13 m) 37.22x 1340.03x
pypomp (GPU, JAX Samplers) 1745.66 s (29.09 m) 1.11x 629.03 s (10.48 m) 0.94x 2374.69 s (39.58 m) 1.06x 38.29x

Extrapolation to 1,422 Cities

Based on the runtimes measured for a single city (London) under each of the three setups, we can project the time required to scale the estimation across a larger dataset of 1,422 cities under two scenarios:

  1. Config 1 (Same configuration): Running IF2 with 5,000 particles, 100 iterations, 36 replicates per city, and Pfilter with 5,000 particles, 36 replications per parameter set.
  2. Config 2 (2x particles, 2x iterations): Running IF2 with 10,000 particles, 200 iterations (scaling IF2 computation by a factor of 4), 36 replicates per city, and Pfilter with 10,000 particles (scaling Pfilter computation by a factor of 2), 36 replications per parameter set.

Assuming linear scaling with respect to the number of cities, particles, and iterations:

Method Config 1: MIF2 Time Config 1: Pfilter Time Config 1: Total Time Config 2: MIF2 Time Config 2: Pfilter Time Config 2: Total Time
R pomp (CPU, 36 Cores) 31.84 days
(1.05 months)		 9.73 days
(0.32 months)		 41.57 days
(1.37 months)		 127.37 days
(4.18 months)		 19.46 days
(0.64 months)		 146.83 days
(4.82 months)
pypomp (GPU, Fast Samplers) 0.81 days
(0.03 months)		 0.31 days
(0.01 months)		 1.12 days
(0.04 months)		 3.22 days
(0.11 months)		 0.62 days
(0.02 months)		 3.85 days
(0.13 months)
pypomp (GPU, JAX Samplers) 28.73 days
(0.94 months)		 10.35 days
(0.34 months)		 39.08 days
(1.28 months)		 114.92 days
(3.78 months)		 20.71 days
(0.68 months)		 135.63 days
(4.46 months)

Note that, for the GPU, times could be cut to about 1/G of the recorded times if the replications were split across G GPUs.

Log-Likelihood Comparison

While the point of these runs was to benchmark the execution time of different approaches, and 36 replications isn’t really sufficient to accurately compare estimate distributions, we also include comparisons of the log-likelihood and parameter estimates in case anything interesting crops up.

Ideally, the estimated log-likelihood distributions are roughly similar.

Show Code 
# Extract logLiks
# For R, it is mean_logLik per replicate
r_logliks = r_results.groupby("replicate")["mean_logLik"].first().values

# For pypomp, the results_df contains logLik
py_fast_logliks = py_fast["results"]["logLik"].values
py_jax_logliks = py_jax["results"]["logLik"].values

loglik_df = pd.DataFrame(
    {
        "Method": (["R pomp"] * len(r_logliks))
        + (["pypomp (Fast)"] * len(py_fast_logliks))
        + (["pypomp (JAX)"] * len(py_jax_logliks)),
        "LogLik": np.concatenate([r_logliks, py_fast_logliks, py_jax_logliks]),
    }
)

p_loglik = (
    ggplot(loglik_df, aes(x="LogLik", fill="Method"))
    + geom_density(alpha=0.4)
    + theme_minimal()
    + labs(
        title="Comparison of Final Log-Likelihood Estimates across 36 Replicates",
        x="Log-Likelihood",
        y="Density",
    )
    + theme(legend_position="bottom")
    + xlim((-3950, None))
)
p_loglik

Show Code 
# Print summary table
summary_loglik = pd.DataFrame(
    {
        "Method": ["R pomp", "pypomp (Fast)", "pypomp (JAX)"],
        "Mean LogLik": [
            np.mean(r_logliks),
            np.mean(py_fast_logliks),
            np.mean(py_jax_logliks),
        ],
        "Max LogLik": [
            np.max(r_logliks),
            np.max(py_fast_logliks),
            np.max(py_jax_logliks),
        ],
        "Min LogLik": [
            np.min(r_logliks),
            np.min(py_fast_logliks),
            np.min(py_jax_logliks),
        ],
        "Std Dev": [np.std(r_logliks), np.std(py_fast_logliks), np.std(py_jax_logliks)],
    }
)

HTML(summary_loglik.to_html(classes="table table-striped table-hover", index=False))
Method Mean LogLik Max LogLik Min LogLik Std Dev
R pomp -3844.813834 -3827.900641 -3930.173213 22.901298
pypomp (Fast) -3834.208354 -3826.178438 -3852.534698 5.735832
pypomp (JAX) -3838.517351 -3825.317179 -3895.226480 13.523887

Parameter Estimate Comparisons

Here we display some of the estimated parameter distributions. Ideally, the estimated parameter distributions are roughly similar.

Show Code 
# Extract parameter names and values
# R results are in long format: unit, replicate, names (param), coef (value)
r_params = r_results[["replicate", "names", "coef"]].copy()
r_params["source"] = "R pomp"
r_params = r_params.rename(columns={"names": "param", "coef": "value"})

# pypomp results are in wide format
py_fast_df = py_fast["results"].copy()
py_fast_df["source"] = "pypomp (Fast)"
py_fast_df["replicate"] = py_fast_df.index + 1
py_fast_melt = py_fast_df.drop(
    columns=["logLik", "se", "theta_idx", "unit"], errors="ignore"
).melt(id_vars=["source", "replicate"], var_name="param", value_name="value")

py_jax_df = py_jax["results"].copy()
py_jax_df["source"] = "pypomp (JAX)"
py_jax_df["replicate"] = py_jax_df.index + 1
py_jax_melt = py_jax_df.drop(
    columns=["logLik", "se", "theta_idx", "unit"], errors="ignore"
).melt(id_vars=["source", "replicate"], var_name="param", value_name="value")

# Combine all results
combined_params = pd.concat([r_params, py_fast_melt, py_jax_melt], ignore_index=True)

# Select a subset of key parameters to display
key_params = ["R0", "rho", "sigmaSE", "amplitude", "gamma", "sigma", "psi", "cohort"]
combined_key_params = combined_params[combined_params["param"].isin(key_params)]

p_params = (
    ggplot(combined_key_params, aes(x="value", fill="source"))
    + geom_density(alpha=0.4)
    + facet_wrap("~param", scales="free", ncol=2)
    + theme_minimal()
    + labs(
        title="Parameter Estimate Distributions across 36 Replicates",
        x="Parameter Value",
        y="Density",
        fill="Method"
    )
    + theme(figure_size=(10, 12))
)
p_params

Conclusion

This speed comparison benchmarks the performance advantages of GPU-accelerated JAX-based parameter estimation:

  • pypomp with Fast Samplers achieves a massive speedup over R pomp due to custom GPU approximations.
  • pypomp with standard JAX samplers serves as a baseline to highlight the benefits of the optimized inverse CDF approximations.