"""
SIR model with seasonal forcing, translated from pomp::sir in R.
This version supports DPOP (Differentiable Particle Filter) by accumulating
process log-density in the state variable 'logw'.
"""
from typing import Any
import jax
import jax.numpy as jnp
import jax.scipy.special as jspecial
import numpy as np
import pandas as pd
from pypomp.core.pomp import Pomp
from pypomp.core.par_trans import ParTrans
from pypomp.types import ParamDict
from pypomp.models.ctmc_multinom import sample_and_log_prob
from pypomp.random.poisson import fast_poisson
from pypomp.random.gamma import fast_gamma
from pypomp.random.nbinom import fast_nbinomial
STATENAMES = ["S", "I", "R", "cases", "W", "logw"]
DEFAULT_THETA = {
"gamma": 26.0,
"mu": 0.02,
"iota": 0.01,
"beta1": 400.0,
"beta2": 480.0,
"beta3": 320.0,
"beta_sd": 0.001,
"rho": 0.6,
"k": 0.1,
"pop": 2100000.0,
"S_0": 26.0 / 400.0,
"I_0": 0.001,
"R_0": 1.0 - 26.0 / 400.0 - 0.001,
}
# Periodic B-spline basis (from pomp's C implementation)
def _bspline_eval(
x: Any,
knots: Any,
i: int,
degree: int,
deriv: int = 0,
) -> Any:
if deriv > degree:
return jnp.zeros_like(x)
elif deriv > 0:
i2 = i + 1
p2 = degree - 1
d2 = deriv - 1
y1 = _bspline_eval(x, knots, i, p2, d2)
y2 = _bspline_eval(x, knots, i2, p2, d2)
denom1 = knots[i + degree] - knots[i]
denom2 = knots[i2 + degree] - knots[i2]
a = jnp.where(jnp.abs(denom1) > 1e-10, degree / denom1, 0.0)
b = jnp.where(jnp.abs(denom2) > 1e-10, degree / denom2, 0.0)
return a * y1 - b * y2
else:
if degree > 0:
i2 = i + 1
p2 = degree - 1
y1 = _bspline_eval(x, knots, i, p2, 0)
y2 = _bspline_eval(x, knots, i2, p2, 0)
denom1 = knots[i + degree] - knots[i]
denom2 = knots[i2 + degree] - knots[i2]
a = jnp.array(
jnp.where(jnp.abs(denom1) > 1e-10, (x - knots[i]) / denom1, 0.0)
)
b = jnp.array(
jnp.where(
jnp.abs(denom2) > 1e-10,
(knots[i2 + degree] - x) / denom2,
0.0,
)
)
return a * y1 + b * y2
else:
return jnp.where((knots[i] <= x) & (x < knots[i + 1]), 1.0, 0.0)
def periodic_bspline_basis_eval(
x: Any,
period: float,
degree: int,
nbasis: int,
deriv: int = 0,
) -> jax.Array:
nknots = nbasis + 2 * degree + 1
shift = (degree - 1) // 2
dx = period / nbasis
knots = jnp.array([(k) * dx for k in range(-degree, nbasis + degree + 2)])
x_wrapped = x % period
x_wrapped = jnp.where(x_wrapped < 0, x_wrapped + period, x_wrapped)
yy = jnp.array(
[_bspline_eval(x_wrapped, knots, k, degree, deriv) for k in range(nknots)]
)
yy_adjusted = yy.at[:degree].add(yy[nbasis : nbasis + degree])
y = jnp.array([yy_adjusted[(shift + k) % nbasis] for k in range(nbasis)])
return y
def precompute_bspline_covars(times, t0, period=1.0, nbasis=3, degree=3):
t_min = t0
t_max = max(times) + 0.2
t_grid = np.arange(t_min, t_max + 0.01, 0.01)
basis_matrix = periodic_bspline_basis_eval(t_grid, period, degree, nbasis, deriv=0)
basis_data = {f"seas_{i + 1}": np.array(basis_matrix[i]) for i in range(nbasis)}
return pd.DataFrame(basis_data, index=pd.Index(t_grid, name="time"))
def to_est(theta: ParamDict) -> ParamDict:
SIR_0 = jnp.array([theta["S_0"], theta["I_0"], theta["R_0"]])
SIR_0 = SIR_0 / jnp.sum(SIR_0)
S_0_est, I_0_est, R_0_est = jnp.log(SIR_0)
return {
"gamma": jnp.log(theta["gamma"]),
"mu": jnp.log(theta["mu"]),
"iota": jnp.log(theta["iota"]),
"beta1": jnp.log(theta["beta1"]),
"beta2": jnp.log(theta["beta2"]),
"beta3": jnp.log(theta["beta3"]),
"beta_sd": jnp.log(theta["beta_sd"]),
"rho": jspecial.logit(theta["rho"]),
"k": jnp.log(theta["k"]),
"pop": jnp.log(theta["pop"]),
"S_0": S_0_est,
"I_0": I_0_est,
"R_0": R_0_est,
}
def from_est(theta: ParamDict) -> ParamDict:
SIR_0 = jnp.exp(jnp.array([theta["S_0"], theta["I_0"], theta["R_0"]]))
SIR_0 = SIR_0 / jnp.sum(SIR_0)
return {
"gamma": jnp.exp(theta["gamma"]),
"mu": jnp.exp(theta["mu"]),
"iota": jnp.exp(theta["iota"]),
"beta1": jnp.exp(theta["beta1"]),
"beta2": jnp.exp(theta["beta2"]),
"beta3": jnp.exp(theta["beta3"]),
"beta_sd": jnp.exp(theta["beta_sd"]),
"rho": jspecial.expit(theta["rho"]),
"k": jnp.exp(theta["k"]),
"pop": jnp.exp(theta["pop"]),
"S_0": SIR_0[0],
"I_0": SIR_0[1],
"R_0": SIR_0[2],
}
def rinit(theta_, key, covars=None, t0=None):
pop = theta_["pop"]
S_0 = theta_["S_0"]
I_0 = theta_["I_0"]
R_0 = theta_["R_0"]
m = pop / (S_0 + I_0 + R_0)
S = jnp.round(m * S_0)
I = jnp.round(m * I_0)
R = jnp.round(m * R_0)
return {"S": S, "I": I, "R": R, "cases": 0.0, "W": 0.0, "logw": 0.0}
def rproc(X_, theta_, key, covars, t, dt):
S, I, R = X_["S"], X_["I"], X_["R"]
cases, W, logw = X_["cases"], X_["W"], X_["logw"]
gamma = theta_["gamma"]
mu = theta_["mu"]
iota = theta_["iota"]
beta1, beta2, beta3 = theta_["beta1"], theta_["beta2"], theta_["beta3"]
beta_sd = theta_["beta_sd"]
pop = theta_["pop"]
seas_1, seas_2, seas_3 = (
covars["seas_1"],
covars["seas_2"],
covars["seas_3"],
)
beta = beta1 * seas_1 + beta2 * seas_2 + beta3 * seas_3
key, subkey = jax.random.split(key)
shape = dt / (beta_sd**2 + 1e-10)
dW = fast_gamma(subkey, shape) * (beta_sd**2)
rate_foi = (iota + beta * I * dW / dt) / pop
key, k1, k_trans = jax.random.split(key, 3)
births = fast_poisson(k1, mu * pop * dt)
S_int = jnp.maximum(jnp.round(S), 0.0)
I_int = jnp.maximum(jnp.round(I), 0.0)
R_int = jnp.maximum(jnp.round(R), 0.0)
N = jnp.stack([S_int, I_int, R_int])
rates = jnp.stack(
[jnp.array([rate_foi, mu]), jnp.array([gamma, mu]), jnp.array([mu, 0.0])]
)
keys_compartments = jax.random.split(k_trans, 3)
# Should manually vectorize sample_and_log_prob
trans, lps, _ = jax.vmap(sample_and_log_prob, in_axes=(0, 0, None, 0))(
N, rates, dt, keys_compartments
)
infections, deaths_S = trans[0, 0], trans[0, 1]
recoveries, deaths_I = trans[1, 0], trans[1, 1]
deaths_R = trans[2, 0]
logw_step = jnp.sum(lps)
logw_step = jnp.where(jnp.isfinite(logw_step), logw_step, 0.0)
S_new = S + births - infections - deaths_S
I_new = I + infections - recoveries - deaths_I
R_new = R + recoveries - deaths_R
cases_new = cases + recoveries
W_new = jnp.where(beta_sd > 0, W + (dW - dt) / beta_sd, W)
logw_new = logw + logw_step
return {
"S": S_new,
"I": I_new,
"R": R_new,
"cases": cases_new,
"W": W_new,
"logw": logw_new,
}
def dmeas(Y_, X_, theta_, covars=None, t=None):
reports = Y_["reports"]
cases = X_["cases"]
rho, k = theta_["rho"], theta_["k"]
mu = jnp.maximum(cases * rho, 1e-10)
size = 1.0 / k
reports_int = jnp.round(reports)
return (
jax.scipy.special.gammaln(reports_int + size)
- jax.scipy.special.gammaln(size)
- jax.scipy.special.gammaln(reports_int + 1)
+ size * jnp.log(size / (size + mu))
+ reports_int * jnp.log(mu / (size + mu))
)
def rmeas(X_, theta_, key, covars=None, t=None):
cases = X_["cases"]
rho, k = theta_["rho"], theta_["k"]
mu = jnp.maximum(cases * rho, 1e-10)
size = 1.0 / k
reports = fast_nbinomial(key, n=size, mu=mu)
return {"reports": reports}
[docs]
def sir(
gamma: float = 26.0,
mu: float = 0.02,
iota: float = 0.01,
beta1: float = 400.0,
beta2: float = 480.0,
beta3: float = 320.0,
beta_sd: float = 0.001,
rho: float = 0.6,
k: float = 0.1,
pop: float = 2100000.0,
S_0: float = 26.0 / 400.0,
I_0: float = 0.001,
R_0: float | None = None,
t0: float = 0.0,
times: np.ndarray | None = None,
seed: int = 329343545,
delta_t: float = 1 / 52 / 20,
) -> Pomp:
"""
Create a Pomp object for the SIR model with seasonal forcing.
Supports DPOP through the logw state variable.
"""
if R_0 is None:
R_0 = 1.0 - S_0 - I_0
if times is None:
times = np.arange(t0 + 1 / 52, t0 + 4 + 1 / 52, 1 / 52)
theta = {
"gamma": gamma,
"mu": mu,
"iota": iota,
"beta1": beta1,
"beta2": beta2,
"beta3": beta3,
"beta_sd": beta_sd,
"rho": rho,
"k": k,
"pop": pop,
"S_0": S_0,
"I_0": I_0,
"R_0": R_0,
}
covars = precompute_bspline_covars(times, t0)
par_trans = ParTrans(to_est=to_est, from_est=from_est)
ys_dummy = pd.DataFrame({"reports": np.zeros(len(times))}, index=pd.Index(times))
accumvars = ("cases", "logw")
from pypomp.core.parameters import PompParameters
sir_temp = Pomp(
ys=ys_dummy,
theta=PompParameters(theta),
statenames=STATENAMES,
t0=t0,
rinit=rinit,
rproc=rproc,
dmeas=dmeas,
rmeas=rmeas,
par_trans=par_trans,
nstep=20,
accumvars=accumvars,
covars=covars,
)
key = jax.random.key(seed)
sim_pomp = sir_temp.simulate(key=key, nsim=1, as_pomp=True)
return sim_pomp
def get_process_weight_index():
"""Return the index of logw in STATENAMES for DPOP."""
return STATENAMES.index("logw")