Measles in England and Wales — Spatial SEIR Model¶

import sys
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from laser.core import PropertySet
from laser.core.utils import grid
from laser.core.migration import gravity as gravity_model
from laser.generic.utils import ValuesMap
from laser.cohorts import Model, Campaign, Intervention
from laser.cohorts.vitaldynamics import NonDiseaseMortality, ConstantPopBirths
import laser.cohorts.SEIR as SEIR

data_dir = Path("data")
sys.path.insert(0, str(data_dir))
from EnglandWalesMeasles import data as engwal
distances = np.load(data_dir / "EnglandWalesDistances.npy")

print(f"Places: {len(engwal.placenames)}")
print(f"Years: {engwal.years[0]+1900} – {engwal.years[-1]+1900}")
print(f"Reporting periods: {len(engwal.reports)}")
print(f"Distances matrix: {distances.shape}")

placenames = engwal.placenames
n_places   = len(placenames)

# Per-city populations and birth rates across 1944-1964
populations_annual = np.array([engwal.places[p].population for p in placenames])  # (954, 21)
births_annual      = np.array([engwal.places[p].births     for p in placenames])  # (954, 21)

# Annual birth rate per capita per city
cbr_annual = births_annual / populations_annual   # (954, 21)  — dimensionless fraction per year

# Latitude and longitude for mapping
latitudes  = np.array([engwal.places[p].latitude  for p in placenames])
longitudes = np.array([engwal.places[p].longitude for p in placenames])

# Observed bi-weekly cases per city — (954, 546)
obs_cases  = np.array([engwal.places[p].cases for p in placenames])

print(f"Population range: {populations_annual[:,0].min():,} \u2013 {populations_annual[:,0].max():,}")
print(f"Mean annual CBR:  {cbr_annual.mean():.4f}")

# 1944 populations and birth rates
pop_1944 = populations_annual[:, 0].astype(np.float64)   # (954,)
cbr_1944 = cbr_annual[:, 0]                              # (954,) annual fraction

# Seed infection: use first bi-weekly report as initial I estimate
initial_cases = obs_cases[:, 0].astype(np.float64)  # (954,)

# SEIR initial conditions (approximate):
# R0 ≈ 15 → endemic S* ≈ N/R0
R0_approx = 15.0
S_frac = 1.0 / R0_approx         # ~6.7% susceptible at endemic equilibrium
I0 = np.maximum(initial_cases, 1).astype(np.int32)   # at least 1 per city
E0 = I0.copy()                    # assume similar number exposed
S0 = np.maximum((pop_1944 * S_frac).astype(np.int32) - I0 - E0, 0)
R0_state = (pop_1944.astype(np.int32) - S0 - E0 - I0)  # remainder immune

print("Initial conditions:")
print(f"  Total N:  {pop_1944.sum():,.0f}")
print(f"  Total S:  {S0.sum():,}")
print(f"  Total E:  {E0.sum():,}")
print(f"  Total I:  {I0.sum():,}")
print(f"  Total R:  {R0_state.sum():,}")

# Gravity model parameters (tunable)
k_gravity = 5e-4    # coupling strength
a_gravity = 0.0     # origin-population exponent (0 = no origin-size weighting)
b_gravity = 1.0     # destination-population exponent
c_gravity = 2.0     # distance decay exponent

# laser.core.migration.gravity computes:
#   network[i,j] = k * pop[i]^a * pop[j]^b / distance[i,j]^c
# (diagonal set to 1 internally to avoid div-by-zero, then zeroed in output)
gravity_matrix = gravity_model(
    pop_1944, distances,
    k=k_gravity, a=a_gravity, b=b_gravity, c=c_gravity,
)

print(f"Gravity matrix shape: {gravity_matrix.shape}")
print(f"Max off-diagonal entry: {gravity_matrix.max():.3e}")

# Simulation: 20 years starting 1944
nticks = 20 * 365   # 7300 days

# Measles SEIR parameters
sigma_val = 1.0 / 12.0   # E→I progression rate (12-day incubation)
gamma_val = 1.0 /  5.0   # I→R recovery rate (5-day infectious period)
beta0     = sigma_val * R0_approx  # baseline transmission rate

# Seasonal forcing: sinusoidal with peak in mid-January.
# Day 15 = Jan 15 peak (day 0 = Jan 1 of first year).
epsilon   = 0.3    # seasonal amplitude
t         = np.arange(nticks)
seasonal_1d = 1.0 + epsilon * np.cos(2.0 * np.pi * (t - 15) / 365.0)
seasonal_1d = np.maximum(seasonal_1d, 0.0)   # clip to non-negative
# Tile to (nticks, n_places) — same seasonality everywhere
seasonality = np.tile(seasonal_1d[:, None], (1, n_places))   # (7300, 954)

# Per-city daily mortality rate from 1944 data (convert annual fraction → daily)
r_mort_per_city = cbr_1944 / 365.0              # (954,)
r_mortality_2d  = np.tile(r_mort_per_city[None, :], (nticks, 1))  # (7300, 954)

print(f"beta\u2080 = {beta0:.4f}/day,  \u03c3 = {sigma_val:.4f}/day,  \u03b3 = {gamma_val:.4f}/day")
print(f"R\u2080 (approx) = {beta0 / gamma_val:.1f}")
print(f"Seasonal factor range: {seasonal_1d.min():.3f} \u2013 {seasonal_1d.max():.3f}")

# Scenario: 954 nodes
scenario        = grid(M=n_places, N=1)
scenario["S"]   = S0
scenario["E"]   = E0
scenario["I"]   = I0
scenario["R"]   = R0_state

params = PropertySet({"nticks": nticks})
model  = Model(scenario, params)

# Set gravity network BEFORE assigning components
model.network = gravity_matrix.astype(np.float32)

campaign       = Campaign(model, schedule)
betas          = ValuesMap.from_scalar(beta0, nticks, n_places)
r_recoveries   = ValuesMap.from_scalar(gamma_val, nticks, n_places)
r_progressions = ValuesMap.from_scalar(sigma_val, nticks, n_places)

model.components = [
    SEIR.Susceptible(model),
    SEIR.Exposed(model, r_progression=r_progressions),
    SEIR.Infectious(model, r_recovery=r_recoveries),
    SEIR.Recovered(model),
    NonDiseaseMortality(model, r_mortality=r_mortality_2d),
    ConstantPopBirths(model),
    campaign,
    SEIR.Transmission(model, beta=betas, seasonality=seasonality),
]

print(f"Running {nticks}-day SEIR model with {n_places} nodes…")
model.run()
print("Done.")