ComposableTuringIDModels.jl

Appearance

3

Renewal model with negative-binomial reporting ​

The renewal equation is the workhorse of real-time epidemic estimation: it expresses new infections as a function of past infections weighted by the generation interval, scaled by a time-varying reproduction number [1]. Mishra et al. [2] showed that this renewal construction follows from an age-dependent branching process and pairs naturally with a negative-binomial observation model to give a Bayesian hierarchical model for reported case counts.

This tutorial builds that model from two composed parts — a Renewal infection process that carries an autoregressive latent process for , and a NegativeBinomialError observation model — and fits it to the test-confirmed COVID-19 cases from South Korea that Mishra et al. [2] analysed. The latent process is folded into the renewal model rather than supplied as a separate top-level component: the reproduction number is the renewal model's own parameter process.

The model ​

is the discretised generation interval, the autoregressive damping, the innovation standard deviation, and the observation overdispersion.

Components ​

The latent process is a second-order autoregressive model on with a HierarchicalNormal innovation term, matching Mishra et al. [2]. Strong autocorrelation in the reproduction number is encoded by a first damping prior concentrated near one (  on ) with a weaker second lag. This process is the renewal model's reproduction-number process — it is folded into the infection model below rather than composed separately.

julia
using ComposableTuringIDModels, Distributions, Random, Turing, Mooncake
using ADTypes: AutoMooncake
Random.seed!(1234)

latent = AR(
    damp = [truncated(Normal(0.8, 0.05), 0, 1),
        truncated(Normal(0.1, 0.05), 0, 1)],
    init = [Normal(0.0, 0.2), Normal(0.0, 0.2)],
    ϵ_t = HierarchicalNormal(std = HalfNormal(0.1)))
AR
└─ ϵ_t: HierarchicalNormal

The infection process needs a discrete generation interval. Renewal takes a continuous distribution and discretises it with double interval censoring [9], using CensoredDistributions.jl. Following Mishra et al. [2] we use a serial interval as a proxy for the generation interval. Renewal is the only infection model that carries a generation interval, because it is the only one that uses one; it couples that interval to the latent process (its rt slot) and a prior for the initial infections.

julia
renewal = Renewal(; generation_time = Gamma(6.5, 0.62),
    rt = latent, initialisation = Normal(log(1.0), 0.1))
renewal.gen_int
8-element Vector{Float64}:
 0.026663134095601098
 0.14059778064943768
 0.2502660305615845
 0.24789569560506872
 0.1731751163417783
 0.09635404000022221
 0.045734375752163825
 0.01931382699414364

The stored gen_int is a probability vector — the continuous serial interval binned into daily weights that sum to one. Double interval censoring is not the same as evaluating the continuous density at integer days: it accounts for both the primary and secondary events falling anywhere within their days, which shifts and spreads the mass relative to the underlying [9].

julia
sum(renewal.gen_int), length(renewal.gen_int)
(0.9999999999999999, 8)

The infection process in isolation ​

Because the renewal model is a model in its own right, it can be exercised on its own — without an observation model — and we can isolate the contribution of the renewal equation by pinning its reproduction-number process to a known trajectory. With the latent folded in, the way to do that is to build a renewal model whose rt slot is a deterministic FixedIntercept latent, giving a constant , and to fix the initial-infections parameter. The same as_turing_model call that composes into the full model then runs the infection process standalone, returning its infections I_t and the internal latent draw Z_t.

julia
fixed_logR = log(1.4)
renewal_fixed = Renewal(; generation_time = renewal.gen_int,
    rt = FixedIntercept(fixed_logR), initialisation = Normal())
demo = fix(as_turing_model(renewal_fixed, 60), (init_incidence = 0.0,))()
(constant_Rt = round(exp(first(demo.Z_t)), digits = 2),
    grows = demo.I_t[end] > demo.I_t[1])
(constant_Rt = 1.4, grows = true)

A constant   grows incidence; a path that declined through zero would instead produce the textbook turn-over (incidence growing, decelerating as  , and falling once  ). Driving the renewal model with a richer fixed path is just a matter of swapping the FixedIntercept latent for a deterministic latent of the desired shape. Nothing here is conditioned on data; the component is inspected in isolation before it is assembled into the full model with its sampled process.

Reported cases are overdispersed counts of the latent infections. The prior is placed on the cluster factor , which is roughly the coefficient of variation of the observation noise and so easier to reason about a priori.

julia
obs = NegativeBinomialError(cluster_factor = HalfNormal(0.1))

IDModel assembles the two parts — the renewal infection process (which already carries the latent process) and the observation model — into one composed model.

julia
model = IDModel(renewal, obs)
IDModel
├─ infection: Renewal
│  └─ rt: AR
│     └─ ϵ_t: HierarchicalNormal
└─ observation: NegativeBinomialError

Before fitting, the composed model is also a prior simulator: passing missing observations makes as_turing_model return generated quantities — the reported cases generated_y_t, the latent infections I_t, and the latent process Z_t — instead of conditioning on data. That is the mechanism used for the prior checks above; here we go straight to real data.

The data ​

Mishra et al. [2] fit this model to daily test-confirmed COVID-19 cases in South Korea over the first wave of 2020. The series is stored with the docs and read with CSV/DataFrames.

julia
using CSV, DataFrames
datapath = joinpath(pkgdir(ComposableTuringIDModels),
    "docs", "src", "tutorials", "data", "south_korea_data.csv")
south_korea = CSV.read(datapath, DataFrame)
first(south_korea, 5)
5×4 DataFrame
RowColumn1datecases_newdeaths_new
Int64DateInt64Int64
112019-12-3100
222020-01-0100
332020-01-0200
442020-01-0300
552020-01-0400

We fit the growth-and-decline window of the first wave, matching the span used by Mishra et al. [2], and take the reported cases over it as the observed series.

julia
tspan = (45, 80)
y_obs = south_korea.cases_new[first(tspan):last(tspan)]
n = length(y_obs)
(n = n, total_cases = sum(y_obs),
    from = south_korea.date[first(tspan)], to = south_korea.date[last(tspan)])
(n = 36, total_cases = 8537, from = Dates.Date("2020-02-13"), to = Dates.Date("2020-03-19"))

Fit ​

Conditioning on the observed counts and sampling with NUTS recovers the posterior. We draw two chains in parallel with MCMCThreads() so the posterior is well resolved and the cross-chain diagnostic is available; the slightly raised target acceptance rate keeps the sampler stable on the hierarchical innovation scale. We differentiate with Mooncake, the recommended backend for this package (see Automatic differentiation backend).

julia
posterior = as_turing_model(model, y_obs, n)
chain = sample(
    posterior, NUTS(0.95; adtype = AutoMooncake(; config = nothing)),
    MCMCThreads(), 250, 2; progress = false)
┌ Info: Found initial step size
└   ϵ = 0.0015625
┌ Info: Found initial step size
└   ϵ = 0.003125
┌ Warning: There were 13 divergent transitions. Consider reparameterising your model or using a smaller step size. For adaptive samplers such as NUTS and HMCDA, consider increasing `target_accept`.
└ @ Turing.Inference ~/.julia/packages/Turing/WteH7/src/mcmc/hmc.jl:483
┌ Warning: There were 13 divergent transitions. Consider reparameterising your model or using a smaller step size. For adaptive samplers such as NUTS and HMCDA, consider increasing `target_accept`.
└ @ Turing.Inference ~/.julia/packages/Turing/WteH7/src/mcmc/hmc.jl:483

Sampling returns a chain whose parameters are namespaced by the component slot that samples them, so a prior's inner variables never collide across the model. sample returns a FlexiChains chain, which summarystats summarises directly — no conversion step — giving point estimates and their uncertainty alongside the effective sample size and convergence diagnostic. The autoregressive damping (damp_AR[1]), the innovation scale (std), and the observation cluster factor (cluster_factor) are all identified from the observed South Korean series:

julia
using MCMCChains
summarystats(chain)
╭─FlexiSummary (9 statistics) ─────────────────────────────────────────────────╮
│   iter    collapsed                                                          │
│   chain   collapsed                                                          │
│ ↓ stat  = [mean, std, mcse, ess_bulk, ess_tail, rhat, q5, q50, q95]          │
│                                                                              │
│ Parameters (41) ── AbstractPPL.VarName                                       │
│  Float64  ar_init[1], ar_init[2], damp_AR[1], damp_AR[2], std, ϵ_t[1],       │
│           ϵ_t[2], ϵ_t[3], ϵ_t[4], ϵ_t[5], ϵ_t[6], ϵ_t[7], ϵ_t[8], ϵ_t[9],    │
│           ϵ_t[10], ϵ_t[11], ϵ_t[12], ϵ_t[13], ϵ_t[14], ϵ_t[15], ϵ_t[16],     │
│           ϵ_t[17], ϵ_t[18], ϵ_t[19], ϵ_t[20], ϵ_t[21], ϵ_t[22], ϵ_t[23],     │
│           ϵ_t[24], ϵ_t[25], ϵ_t[26], ϵ_t[27], ϵ_t[28], ϵ_t[29], ϵ_t[30],     │
│           ϵ_t[31], ϵ_t[32], ϵ_t[33], ϵ_t[34], init_incidence,                │
│           cluster_factor                                                     │
│                                                                              │
│ Extras (14)                                                                  │
│  Float64  n_steps, is_accept, acceptance_rate, log_density,                  │
│           hamiltonian_energy, hamiltonian_energy_error,                      │
│           max_hamiltonian_energy_error, tree_depth, numerical_error,         │
│           step_size, nom_step_size, logprior, loglikelihood, logjoint        │
│                                                                              │
│ Summary                                                                      │
│          param     mean     std    mcse   ess_bulk  ess_tail    rhat  …      │
│     ar_init[1]   0.0297  0.1908  0.0057  1111.3174  276.4387  0.9989  …      │
│     ar_init[2]   0.0123  0.1716  0.0073   556.8629  427.6042  1.0111  …      │
│     damp_AR[1]   0.8135  0.0392  0.0025   234.1703  315.0399  1.0011  …      │
│     damp_AR[2]   0.0835  0.0417  0.0038   104.7721   61.3586  1.0099  …      │
│            std   0.4212  0.0483  0.0045   104.9005  224.1787  1.0292  …      │
│         ϵ_t[1]   0.3364  0.8700  0.0953    85.0932  196.0370  1.0032  …      │
│         ϵ_t[2]   0.9486  0.8566  0.0824   108.1035  178.7927  1.0397  …      │
│         ϵ_t[3]   1.2017  0.9175  0.0855   117.0090  165.6471  0.9999  …      │
│         ϵ_t[4]   1.3340  0.7813  0.0665   139.4304  212.5937  1.0019  …      │
│         ϵ_t[5]   2.4518  0.8176  0.0764   113.2875  190.8894  1.0071  …      │
│         ϵ_t[6]   1.6282  0.6631  0.0457   209.6050  265.9878  1.0054  …      │
│         ϵ_t[7]   0.8097  0.6396  0.0376   295.7441  338.0755  0.9992  …      │
│         ϵ_t[8]   0.7212  0.5366  0.0405   173.6624  338.5443  1.0018  …      │
│         ϵ_t[9]  -0.5671  0.5269  0.0347   224.5854  265.8754  1.0027  …      │
│        ϵ_t[10]  -2.2367  0.5244  0.0410   169.6790  186.3983  0.9985  …      │
│        ϵ_t[11]  -1.6688  0.4624  0.0258   327.2451  225.8668  1.0252  …      │
│        ϵ_t[12]   0.5386  0.4119  0.0241   302.7772  281.8821  1.0118  …      │
│        ϵ_t[13]   1.1110  0.3786  0.0201   367.2707  328.4725  1.0068  …      │
│        ϵ_t[14]   0.0230  0.3917  0.0211   357.2496  248.6253  1.0199  …      │
│        ϵ_t[15]   1.2152  0.4095  0.0354   143.7912  198.6995  1.0218  …      │
│        ϵ_t[16]  -1.0884  0.3760  0.0222   284.5610  248.7430  1.0081  …      │
│        ϵ_t[17]  -0.4086  0.3660  0.0209   324.8088  143.5720  1.0590  …      │
│        ϵ_t[18]  -0.7452  0.3670  0.0178   426.1923  372.4143  1.0228  …      │
│        ϵ_t[19]  -0.7088  0.3629  0.0186   417.5723  226.7155  1.0147  …      │
│        ϵ_t[20]  -0.5585  0.3317  0.0140   625.2704  330.8633  1.0070  …      │
│        ϵ_t[21]   0.2215  0.3761  0.0204   370.9979  177.0416  1.0184  …      │
│        ϵ_t[22]  -0.0541  0.3461  0.0172   406.1427  232.0172  1.0146  …      │
│        ϵ_t[23]  -0.5239  0.3690  0.0163   491.2709  296.4179  0.9994  …      │
│        ϵ_t[24]  -0.9147  0.3822  0.0154   661.0467  284.6925  0.9981  …      │
│        ϵ_t[25]  -1.3184  0.4067  0.0305   202.3158  158.5904  1.0052  …      │
│        ϵ_t[26]   0.9799  0.4056  0.0323   183.1909  203.4374  1.0015  …      │
│        ϵ_t[27]  -1.0886  0.4138  0.0333   157.8559  295.8284  1.0172  …      │
│        ϵ_t[28]   0.0182  0.4410  0.0292   250.2402  176.5282  1.0111  …      │
│        ϵ_t[29]   0.1566  0.4328  0.0256   296.6093  183.5589  0.9998  …      │
│        ϵ_t[30]  -0.2882  0.4400  0.0293   237.5254  252.1282  1.0026  …      │
│        ϵ_t[31]   0.1585  0.4310  0.0218   390.5525  408.3394  1.0109  …      │
│        ϵ_t[32]   0.5752  0.4975  0.0259   366.8800  256.9410  1.0151  …      │
│        ϵ_t[33]   0.6351  0.4671  0.0218   475.0790  224.2082  1.0052  …      │
│        ϵ_t[34]   1.2842  0.3794  0.0174   459.3705  342.7612  1.0111  …      │
│   init_incide…  -0.0326  0.0894  0.0038   571.0613  304.6170  0.9989  …      │
│   cluster_fac…   0.0963  0.0611  0.0101    34.9483   38.6355  1.0638  …      │
╰──────────────────────────────────────────────────────────────────────────────╯

Prior versus posterior ​

Before reading the trajectories it is worth asking what the data taught us. Sampling the same model with Prior — ignoring the observations — gives a prior draw over the same parameters, and overlaying it on the posterior shows which parameters moved. We load a Makie backend and PairPlots.jl; the FlexiChains PairPlots extension turns a chain (subset to a few keys with chain[[...]]) into a PairPlots.Series, so prior and posterior overlay on one corner plot.

julia
using CairoMakie, PairPlots

prior_chain = sample(posterior, Prior(), 1000; progress = false)
pp_keys = [@varname(damp_AR), @varname(std),
    @varname(cluster_factor), @varname(init_incidence)]
pairplot(
    PairPlots.Series(chain[pp_keys]; label = "posterior"),
    PairPlots.Series(prior_chain[pp_keys]; label = "prior"))

The innovation scale (std) is sharply updated away from its prior — the data are informative about how much wiggles — while the autoregressive damping (damp_AR), the cluster factor and the initial infections stay closer to their priors on this short window.

Posterior trajectories ​

The reproduction number   is a generated quantity rather than a sampled parameter. generated_observables re-runs the fitted model over the chain to recover the latent and infection trajectories per draw. The reported counts are scored element-wise, so their posterior predictive distribution — fresh counts drawn under each posterior parameter set — comes from predict on the same model with the observations set to missing.

A couple of small helpers reduce the per-draw trajectories to credible bands and draw a median line with 50% and 95% ribbons.

Stack the per-draw into an band, draw the posterior-predictive from the unconditioned model, and plot both against the observed series:

julia
gens = vec(generated_observables(posterior, y_obs, chain).generated)
Rt = credible_bands(reduce(hcat, (exp.(g.Z_t) for g in gens)))

pred = predict(as_turing_model(model, fill(missing, n), n), chain)
yt = predictive_bands(pred, n)

fig = Figure(; size = (760, 620))
ax1 = Axis(fig[1, 1]; ylabel = "Reproduction number Rₜ")
ci_ribbon!(ax1, 1:size(Rt, 1), Rt; color = :purple, label = "posterior median")
hlines!(ax1, [1.0]; color = :grey, linestyle = :dash)
axislegend(ax1; position = :rt)
ax2 = Axis(fig[2, 1]; xlabel = "Day", ylabel = "Reported cases")
ci_ribbon!(ax2, 1:size(yt, 1), yt; color = :teal,
    label = "posterior predictive")
scatter!(ax2, 1:n, y_obs; color = :black, markersize = 7, label = "observed")
axislegend(ax2; position = :lt)
fig

The posterior-predictive band tracks the observed South Korean series closely, and the path recovers the first-wave turn-over: an early rise well above one, a fall through   as the wave peaks, and a decline below one as cases drop.

Forecasting the next weeks ​

The same fitted model forecasts out of sample in one call. Because the latent AR process is non-centred, forecast carries each posterior draw forward — holding the fitted parameters and the in-sample path fixed, and continuing the process over the horizon with fresh prior innovations — then draws the future reported cases:

julia
h = 14
fc = forecast(model, y_obs, chain, h)
size(fc)
(250, 2)

The returned chain carries the predicted over    . We can then plot these forecasts:

julia
# Multi-level CI band quantiles: 90%, 60%, 30% + median
FC_CI_QS = [0.05, 0.2, 0.35, 0.5, 0.65, 0.8, 0.95]

# Draw three nested CI ribbons with the median line
function multi_ci_ribbon!(ax, ts, bands; color, label)
    keep = findall(!ismissing, view(bands, :, 4))
    x, b = ts[keep], Float64.(bands[keep, :])
    # 90% CI (cols 1, 7)
    band!(ax, x, b[:, 1], b[:, 7]; color = (color, 0.1))
    # 60% CI (cols 2, 6)
    band!(ax, x, b[:, 2], b[:, 6]; color = (color, 0.25))
    # 30% CI (cols 3, 5)
    band!(ax, x, b[:, 3], b[:, 5]; color = (color, 0.5))
    # median (col 4)
    lines!(ax, x, b[:, 4]; color = color, linewidth = 2, label = label)
end

# Multi-level credible bands for the forecast
fc_bands = credible_bands(reduce(vcat,
    (permutedims(vec(fc[@varname(y_t[i])])) for i in (n + 1):(n + h)));
    qs = FC_CI_QS)

# Sample 100 random forecast trajectories
n_draws = length(vec(fc[@varname(y_t[n + 1])]))
ntraj = min(100, n_draws)
sample_idx = rand(1:n_draws, ntraj)
trajectories = reduce(hcat,
    [[vec(fc[@varname(y_t[i])])[idx] for i in (n + 1):(n + h)] for idx in sample_idx])

fig_fc = Figure(; size = (760, 360))
axf = Axis(fig_fc[1, 1]; xlabel = "Day", ylabel = "Reported cases",
    yscale = log10)
scatter!(axf, 1:n, y_obs; color = :black, markersize = 7, label = "observed")
# Faint individual forecast trajectories
for idx in 1:ntraj
    lines!(axf, (n + 1):(n + h), max.(trajectories[:, idx], 1);
        color = (:teal, 0.08), linewidth = 0.5)
end
# Multi-level CI bands (offset to keep log10 scale safe at zero)
multi_ci_ribbon!(axf, (n + 1):(n + h), max.(fc_bands, 1); color = :teal,
    label = "forecast")
vlines!(axf, [n + 0.5]; color = :grey, linestyle = :dash)
axislegend(axf; position = :lt)
fig_fc

Swap a component ​

Because the parts share one interface, an alternative observation assumption is a one-line change. Swapping the negative-binomial reporting for a PoissonError leaves the renewal infection process — and its latent process — untouched:

julia
poisson_model = IDModel(renewal, PoissonError())
length(rand(as_turing_model(poisson_model, fill(missing, n), n)))
41

References ​

  1. A. Cori, N. M. Ferguson, C. Fraser and S. Cauchemez. A new framework and software to estimate time-varying reproduction numbers during epidemics. American Journal of Epidemiology 178, 1505–1512 (2013).

  2. S. Mishra, T. Berah, T. A. Mellan, H. J. Unwin, M. A. Vollmer, K. V. Parag, A. Gandy, S. Flaxman and S. Bhatt. On the derivation of the renewal equation from an age-dependent branching process: an epidemic modelling perspective, arXiv preprint arXiv:2006.16487 (2020).

  3. K. Charniga, S. W. Park, A. R. Akhmetzhanov, A. Cori, J. Dushoff, S. Funk and others. Best practices for estimating and reporting epidemiological delay distributions of infectious diseases. PLoS Computational Biology 20, e1012520 (2024).