This tutorial demonstrates how to construct and fit a large panel model to weekly measles case data from 1,422 cities and towns in England and Wales (Korevaar, Metcalf, and Grenfell 2020) using the pypomp package.
Partially Observed Markov Process (POMP) models are powerful tools for studying infectious disease transmission dynamics. When data are available for multiple distinct units (e.g., cities, districts, patients), we can model them jointly as a panel model. In a panel model, some parameters values might be shared across all units (e.g., the recovery rate from a disease), while others are unit-specific (e.g., district-specific reporting rates).
In this tutorial, we will use the UK Measles model (001b in pypomp) (He, Ionides, and King 2010), which is fit to weekly case data from all 1,422 cities and towns. We will go through the setup and code to construct and fit the full 1,422-city panel model.
First, we import the required libraries.
import jax
import jax.numpy as jnp
import numpy as np
import pandas as pd
import plotnine as p9
import pypomp as pp
from pypomp.types import *
# The 1,422 units (cities and towns) in the dataset
UNITS = pp.models.UKMeasles.units()
print("JAX version:", jax.__version__)
print("pypomp version:", pp.__version__)
print("Total available units:", len(UNITS))JAX version: 0.11.2
pypomp version: 1.0.8
Total available units: 1422In pypomp, a panel model is managed using the PanelPomp class. A PanelPomp object wraps a collection of individual Pomp objects (one for each unit/district).
We will fit Model 001b of the UK Measles model family built into pypomp.
Here, we initialize our random seed:
# Set random seed
MAIN_SEED = 631409
key = jax.random.key(MAIN_SEED)
np.random.seed(MAIN_SEED)Next, we define the search box for sampling initial parameter values. This bounding box defines the minimum and maximum constraints for each parameter during initial random searches, and it will be the same for each unit.
measles_box = {
"R0": (10.0, 60.0), # Basic reproduction number
"sigma": (25.0, 100.0), # Incubation rate
"gamma": (25.0, 320.0), # Recovery rate
"iota": (0.004, 3.0), # Importation rate
"rho": (0.1, 0.9), # Reporting rate
"sigmaSE": (0.04, 0.1), # Environmental noise
"psi": (0.05, 3.0), # Overdispersion parameter
"cohort": (0.1, 0.7), # Cohort effect
"amplitude": (0.1, 0.6), # Seasonal amplitude
"S_0": (0.01, 0.07), # Initial susceptible fraction
"E_0": (0.000004, 0.0001), # Initial exposed fraction
"I_0": (0.000003, 0.001), # Initial infectious fraction
"R_0": (0.9, 0.99), # Initial recovered fraction
}We now sample initial parameter replicates for each of our units.
The division of parameters into shared (the value is common to all units) and unit-specific (the value varies independently across units) is not fixed by the model itself. In pypomp, you have complete control over which parameters are shared and which are unit-specific when initializing a PanelParameters object.
We will configure the model such that R0, sigma, gamma, sigmaSE, cohort, and amplitude have shared values across units. The shared_names parameter in pp.PanelPomp.sample_params() allows us to specify which parameters should be shared globally across all units in the resulting PanelParameters object.
NREPS_FITR = 3 # Number of parallel parameter searches (chains)
# Sample parameters for the panel with specified shared names
key, subkey = jax.random.split(key)
initial_params = pp.PanelPomp.sample_params(
measles_box,
n=NREPS_FITR,
units=UNITS,
key=subkey,
shared_names=["R0", "sigma", "gamma", "sigmaSE", "cohort", "amplitude"],
)Instead of constructing the Pomp objects from scratch, we will use the premade model available in pypomp.models.UKMeasles. UKMeasles.panel_pomp() builds a Pomp object for each unit and bundles them together into a single PanelPomp object. The same could be done by hand by passing a dictionary of Pomp objects, such as those made by UKMeasles.pomp(), to pp.PanelPomp.
panel_measles_obj = pp.models.UKMeasles.panel_pomp(
units=UNITS,
theta=initial_params,
model="001b",
clean=True,
)We would like to call attention to the transition process function for this model:
def rproc(
X_: StateDict,
theta_: ParamDict,
key: RNGKey,
covars: CovarDict,
t: TimeFloat,
dt: StepSizeFloat,
):
S, E, I, R, W, C = X_["S"], X_["E"], X_["I"], X_["R"], X_["W"], X_["C"]
R0 = theta_["R0"]
sigma = theta_["sigma"]
gamma = theta_["gamma"]
iota = theta_["iota"]
sigmaSE = theta_["sigmaSE"]
cohort = theta_["cohort"]
amplitude = theta_["amplitude"]
pop = covars["pop"]
birthrate = covars["birthrate"]
mu = 0.02
t_mod = t - jnp.floor(t)
is_cohort_time = jnp.abs(t_mod - 251.0 / 365.0) < 0.5 * dt
br = jnp.where(
is_cohort_time,
cohort * birthrate / dt + (1 - cohort) * birthrate,
(1 - cohort) * birthrate,
)
# term-time seasonality
t_days = t_mod * 365.25
in_term_time = (
((t_days >= 7) & (t_days <= 100))
| ((t_days >= 115) & (t_days <= 199))
| ((t_days >= 252) & (t_days <= 300))
| ((t_days >= 308) & (t_days <= 356))
)
seas = jnp.where(in_term_time, 1.0 + amplitude * 0.2411 / 0.7589, 1 - amplitude)
# transmission rate
beta = R0 * seas * (1.0 - jnp.exp(-(gamma + mu) * dt)) / dt
# expected force of infection
foi = beta * (I + iota) / pop
# white noise (extrademographic stochasticity)
keys = jax.random.split(key, 3)
dw = pp.random.fast_gamma(keys[0], dt / sigmaSE**2) * sigmaSE**2
rate = jnp.array([foi * dw / dt, mu, sigma, mu, gamma, mu])
# Poisson births
births = pp.random.fast_poisson(keys[1], br * dt)
# transitions between classes
rt_final = jnp.zeros((3, 3))
rate_pairs = jnp.array([[rate[0], rate[1]], [rate[2], rate[3]], [rate[4], rate[5]]])
populations = jnp.array([S, E, I])
rate_sums = jnp.sum(rate_pairs, axis=1)
p0_values = jnp.exp(-rate_sums * dt)
rt_final = (
rt_final.at[:, 0:2]
.set(jnp.einsum("ij,i,i->ij", rate_pairs, 1 / rate_sums, 1 - p0_values))
.at[:, 2]
.set(p0_values)
)
transitions = pp.random.fast_multinomial(keys[2], populations, rt_final)
trans_S = transitions[0]
trans_E = transitions[1]
trans_I = transitions[2]
S = S + births - trans_S[0] - trans_S[1]
E = E + trans_S[0] - trans_E[0] - trans_E[1]
I = I + trans_E[0] - trans_I[0] - trans_I[1]
R = pop - S - E - I
W = W + (dw - dt) / sigmaSE
C = C + trans_I[0]
return {"S": S, "E": E, "I": I, "R": R, "W": W, "C": C}Note that this function draws from the gamma, Poisson, and multinomial distributions, and that it does not use the corresponding jax.random functions. This is because these jax.random functions, among others, use sampling algorithms with extensive branching logic, which must be executed sequentially on GPUs, severely impacting performance. In contrast, the samplers in pypomp.random utilize approximate inverse cumulative distribution functions (CDFs) that avoid most of this branching logic (Giles 2016; Giles and Beentjes 2024; Temme 1992). Using the standard JAX samplers can result in execution speeds up to 25 times slower, potentially making the code slower than its R pomp counterpart. The functions jax.random.normal and jax.random.uniform are among the exceptions that are highly optimized; they are perfectly fine to use.
The following table includes execution times for a given POMP task that involved fitting a model to the London time series with IF2 and then evaluating it with the particle filter. Without using the inverse polynomial approximations, pypomp takes even longer on the task than pomp does despite using a GPU.
| Package | Method | Device | Task Time (minutes) |
|---|---|---|---|
pomp |
Standard | CPU | 44 |
pypomp |
Standard | V100 GPU | 180 |
pypomp |
Poly approximation | V100 GPU | 7 |
Despite using approximations, these samplers are quite accurate, as indicated by our quant tests.
We perform parameter estimation using the Marginalized Panel Iterated Filtering (MPIF) algorithm on the panel model (Wheeler, Abkemeier, and Ionides 2025).
First, we set the random walk standard deviations for the parameters. The same set will be used for each unit.
COOLING_RATE = 0.5
DEFAULT_SD = 0.02
DEFAULT_IVP_SD = DEFAULT_SD * 12
RW_SD = pp.RWSigma(
sigmas={
"R0": DEFAULT_SD * 0.25,
"sigma": DEFAULT_SD * 0.25,
"gamma": DEFAULT_SD * 0.5,
"iota": DEFAULT_SD,
"rho": DEFAULT_SD * 0.5,
"sigmaSE": DEFAULT_SD,
"psi": DEFAULT_SD * 0.25,
"cohort": DEFAULT_SD * 0.5,
"amplitude": DEFAULT_SD * 0.5,
"S_0": DEFAULT_IVP_SD,
"E_0": DEFAULT_IVP_SD,
"I_0": DEFAULT_IVP_SD,
"R_0": DEFAULT_IVP_SD,
},
init_names=["S_0", "E_0", "I_0", "R_0"],
).geometric_cooling(a=COOLING_RATE)As usual, we define run levels to control the computational cost of the estimation.
RUN_LEVEL = 1
NP_FITR = (2, 500, 1000, 10000)[RUN_LEVEL - 1] # Number of particles for IF2
NFITR = (2, 10, 100, 200)[RUN_LEVEL - 1] # Number of IF2 iterations
NP_EVAL = (2, 1000, 1000, 5000)[RUN_LEVEL - 1] # Number of particles for evaluation
NREPS_EVAL = (2, 5, 24, 36)[RUN_LEVEL - 1] # Replicates for evaluationWe are now ready to fit the panel model. We run the IF2 estimation algorithm first:
key, subkey = jax.random.split(key)
panel_measles_obj.mif(
rw_sd=RW_SD,
M=NFITR,
J=NP_FITR,
key=subkey,
)After running the estimation, we run a particle filter on our parameter sets to evaluate their log-likelihoods:
panel_measles_obj.pfilter(J=NP_EVAL, reps=NREPS_EVAL)The panel log-likelihood is the sum of the unit log-likelihoods, and each unit’s log-likelihood depends only on the shared parameters and that unit’s own parameters. We can therefore try to improve on the best search by combining the best components of different searches. mix_and_match() takes the shared parameters from the search with the highest panel log-likelihood and, for each unit, the unit-specific parameters from the search with the highest log-likelihood for that unit. If each search is fixed at the same shared parameter values, the sum of the best unit-specific log-likelihoods is an accurate estimate of the mix-and-matched panel model log-likelihood.
Choosing the best of several noisy estimates biases the resulting log-likelihood upward, so we keep only this combined parameter set with prune() and evaluate it again with the particle filter:
panel_measles_obj.mix_and_match()
panel_measles_obj.prune(n=1, refill=False)
panel_measles_obj.pfilter(J=NP_EVAL, reps=NREPS_EVAL)Fitting a panel model of this scale (1,422 units with large numbers of particles and iterations) requires significant computing resources. Because pypomp is written in Python using JAX, it compiles down to highly parallelized operations that can execute concurrently across hundreds of thousands of particles and method replicates on modern GPUs.
Below, we detail the execution performance and results of a full Level 4 panel model search run on an NVIDIA RTX6000 Pro Blackwell GPU. Pending a fit using shared parameters, we show the results of this run, which use only unit-specific parameters.
The run directory results/2026-10-05_134449_u1422_L4_d001_reps36/ contains a copy of the exact script (script_snapshot.py), the locked Python environment (uv.lock.gz), and a record of the settings and software versions (metadata.json). That script relies on a helper package, pymeasles, which is not part of this repository. The script big_measles.py in this directory runs the same search without it.
These timings and the results in the next section come from a run with pypomp 1.0.8.
The recorded execution times are as follows:
| Stage | Algorithm | Config Details | Time (seconds) | Time (hours) |
|---|---|---|---|---|
| Estimation | MPIF | \(J=10000, M=200, 36 \text{ chains}\) | \(127,866.61 \text{ s}\) | \(35.52 \text{ hours}\) |
| Initial Eval | PFilter | \(J=5000, 36 \text{ replicates, } 36 \text{ chains}\) | \(8,933.12 \text{ s}\) | \(2.48 \text{ hours}\) |
| Post-Pruning Eval | PFilter | \(J=5000, 36 \text{ replicates, } 1 \text{ chain (pruned)}\) | \(338.17 \text{ s}\) | \(5.64 \text{ minutes}\) |
| Total | \(137,137.90 \text{ s}\) | \(38.09 \text{ hours}\) |
This computation would have taken about 4.8 months in the R package panelPomp on a 36-core CPU node! By using a capable GPU, we completed the same computation about 91 times faster. We could complete it even faster by splitting the replications across additional GPUs.
The level 4 run saved the output of traces() and results() as Parquet files. We saved these rather than pickling the PanelPomp object, because pickled pypomp objects generally cannot be loaded by other versions of pypomp. The script summarize_results.py condenses these files into the small CSV files read below. Note that the level 4 run treated all parameters as unit-specific.
RESULTS_DIR = "results/2026-10-05_134449_u1422_L4_d001_reps36/summary"
panel_trace = pd.read_csv(f"{RESULTS_DIR}/panel_loglik_trace.csv")
chain_evals = pd.read_csv(f"{RESULTS_DIR}/chain_evaluations.csv.gz")
final_estimates = pd.read_csv(f"{RESULTS_DIR}/final_estimates.csv.gz")
town_traces = pd.read_csv(f"{RESULTS_DIR}/he10_town_traces.csv.gz")The plot below shows the panel log-likelihood, summed over all 1,422 units, at each MPIF iteration for each of the 36 searches. We omit the first nine iterations, whose log-likelihoods are far lower. The log-likelihood rises steeply at first and then keeps improving slowly. The searches track each other closely, and they are still improving after 200 iterations, which suggests that more iterations would help.
(
p9.ggplot(
panel_trace.query("iteration >= 10"),
p9.aes(x="iteration", y="logLik", color="factor(theta_idx)"),
)
+ p9.geom_line()
+ p9.theme_minimal()
+ p9.labs(x="Iteration", y="Panel log-likelihood", color="Search")
)
The table below compares the panel log-likelihood of the best single search with two versions of the mixed-and-matched parameter set. Adding up each unit’s best log-likelihood from the first round of particle filtering overstates the log-likelihood of the combined parameters, because each unit’s best estimate is partly the luck of the draw. Evaluating the combined parameters again with the particle filter removes this selection bias. Even so, mixing and matching improves the panel log-likelihood by more than 3,500 units over the best single search.
search_evals = chain_evals.query("unit == 'shared'")
unit_evals = chain_evals.query("unit != 'shared'")
best_search = search_evals.loc[search_evals["logLik"].idxmax()]
best_per_unit = unit_evals.loc[unit_evals.groupby("unit")["logLik"].idxmax()]
pd.DataFrame(
{
"Panel log-likelihood": [
best_search["logLik"],
best_per_unit["logLik"].sum(),
final_estimates["shared logLik"].iloc[0],
],
"SE": [best_search["se"], np.nan, final_estimates["shared logLik se"].iloc[0]],
},
index=[
"Best single search",
"Mix and match (sum of best unit log-likelihoods)",
"Mix and match (re-evaluated)",
],
).round(1)| Panel log-likelihood | SE | |
|---|---|---|
| Best single search | -1444159.1 | 21.7 |
| Mix and match (sum of best unit log-likelihoods) | -1440152.9 | NaN |
| Mix and match (re-evaluated) | -1440580.5 | 37.2 |
Every search supplies the best parameters for some units, so no single search dominates:
best_per_unit["theta_idx"].value_counts().sort_index().rename("Units").to_frame().T| theta_idx | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | ... | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Units | 34 | 42 | 41 | 45 | 45 | 33 | 38 | 41 | 35 | 40 | ... | 37 | 47 | 38 | 28 | 36 | 35 | 38 | 48 | 39 | 43 |
1 rows × 36 columns
The exported traces include every parameter for the 20 cities and towns studied by He, Ionides, and King (2010). Below, we show the traces of four parameters for London, the largest unit, and Halesworth, a town of about 2,000 people.
unit_traces_long = town_traces.query("unit in ['London', 'Halesworth']").melt(
id_vars=["unit", "theta_idx", "iteration"],
value_vars=["R0", "sigmaSE", "rho", "psi"],
)
(
p9.ggplot(
unit_traces_long,
p9.aes(x="iteration", y="value", color="factor(theta_idx)"),
)
+ p9.geom_line()
+ p9.facet_grid("variable ~ unit", scales="free_y")
+ p9.theme_minimal()
+ p9.labs(x="Iteration", y="Value", color="Search")
)
In London, all but one search agree closely on the reporting rate rho and the reporting overdispersion psi, and they estimate R0 between 12 and 58. Search 5 instead settles at an R0 of about 1.4 and a reporting rate above 0.9. Its London log-likelihood is 78 units below the best search, so it appears to be stuck near an inferior local maximum. This is one reason to run several searches from different starting points. The searches for Halesworth disagree more on every parameter, because the data from a small town carry less information.
Finally, we plot the final estimates of some parameters for all 1,422 units against their 1950 populations.
demog = pp.models.UKMeasles.subset()["demog"]
pop_1950 = demog.query("year == 1950").set_index("unit")["pop"]
df_plot = final_estimates.assign(pop_1950=final_estimates["unit"].map(pop_1950))
for param, log_y in [("R0", True), ("iota", True), ("psi", True), ("rho", False)]:
p = (
p9.ggplot(df_plot, p9.aes(x="pop_1950", y=param))
+ p9.geom_point(alpha=0.3, size=1)
+ p9.scale_x_log10()
+ (p9.scale_y_log10() if log_y else p9.scale_y_continuous())
+ p9.theme_minimal()
+ p9.labs(x="1950 population", y="Estimate", title=param)
)
p.show()
Estimates for small towns vary widely, reflecting how little information each small town’s data carry. Even so, some parameters show clear trends with population size: larger towns have higher estimated importation rates (iota) and less overdispersed reporting (psi). Estimates of R0 and rho are more consistent among large towns than among small ones. Patterns like these can guide which parameters to share across units or to model as functions of covariates. For example, model 002 in pypomp makes iota a function of 1950 population size.