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 os
import time
import jax
import numpy as np
import pandas as pd
import pypomp as pp
from datetime import datetime
# Load the list of 1,422 units (districts)
from units import UNITS
print("JAX version:", jax.__version__)
print("pypomp version:", pp.__version__)
print("Total available units:", len(UNITS))JAX version: 0.10.1
pypomp version: 0.4.6.0
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)
key, subkey = jax.random.split(key)
dummy_initial_params_list = pp.Pomp.sample_params(measles_box, NREPS_FITR, key=subkey)
# Sample parameters for the panel with specified shared names
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. We instantiate this Pomp class for each unit and bundle them together into a single PanelPomp object:
pomp_dict = {
unit: pp.models.UKMeasles.Pomp(
unit=[unit],
theta=dummy_initial_params_list,
model="001b",
clean=True,
)
for unit in UNITS
}
# Create the PanelPomp object
panel_measles_obj = pp.PanelPomp(
Pomp_dict=pomp_dict,
theta=initial_params,
)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.
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"],
)
COOLING_RATE = 0.5As 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,
a=COOLING_RATE,
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)
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.
The recorded execution times are as follows:
| Stage | Algorithm | Config Details | Time (seconds) | Time (hours) |
|---|---|---|---|---|
| Estimation | MPIF | \(J=10000, M=200, 12 \text{ chains}\) | \(113,321.28 \text{ s}\) | \(31.48 \text{ hours}\) |
| Initial Eval | PFilter | \(J=5000, 36 \text{ replicates, } 12 \text{ chains}\) | \(9,692.40 \text{ s}\) | \(2.69 \text{ hours}\) |
| Post-Pruning Eval | PFilter | \(J=5000, 36 \text{ replicates, } 1 \text{ chain (pruned)}\) | \(742.95 \text{ s}\) | \(12.38 \text{ minutes}\) |
| Total | \(123,756.63 \text{ s}\) | \(34.38 \text{ hours}\) |
This computation would have taken about 4.7 months in the R package pomp using CPU cores! By using a capable GPU, we compeleted the same computation about 98 times faster.