ComposableTuringIDModels.jl

Appearance

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 case study 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_priors = [truncated(Normal(0.8, 0.05), 0, 1),
        truncated(Normal(0.1, 0.05), 0, 1)],
    init_priors = [Normal(0.0, 0.2), Normal(0.0, 0.2)],
    ϵ_t = HierarchicalNormal(std_prior = 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 [7], 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(gen_distribution = Gamma(6.5, 0.62);
    rt = latent, initialisation_prior = 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 [7].

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(renewal.gen_int;
    rt = FixedIntercept(fixed_logR), initialisation_prior = 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_prior = 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", "case-studies", "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.9; adtype = AutoMooncake(; config = nothing)),
    MCMCThreads(), 500, 2; progress = false)
┌ Warning: Only a single thread available: MCMC chains are not sampled in parallel
└ @ AbstractMCMC ~/.julia/packages/AbstractMCMC/C1aKp/src/sample.jl:544
┌ Info: Found initial step size
└   ϵ = 0.003125
┌ Info: Found initial step size
└   ϵ = 0.0015625
┌ Warning: There were 76 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 76 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 keep their flat component names (prefixing is disabled throughout the package). 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.0190  0.1949  0.0089  538.2791  710.0376  1.0009  …       │
│     ar_init[2]   0.0040  0.1695  0.0077  471.4991  762.3618  1.0097  …       │
│     damp_AR[1]   0.8180  0.0411  0.0026  257.1754  351.7164  1.0016  …       │
│     damp_AR[2]   0.0835  0.0384  0.0028  163.5054  116.2066  1.0048  …       │
│            std   0.4177  0.0437  0.0034  168.6978  250.8066  1.0229  …       │
│         ϵ_t[1]   0.3491  0.8981  0.0752  146.6286  428.4816  1.0015  …       │
│         ϵ_t[2]   1.0534  0.8358  0.0380  482.0088  519.2632  1.0162  …       │
│         ϵ_t[3]   1.1630  0.8693  0.0580  229.0616  667.6761  1.0032  …       │
│         ϵ_t[4]   1.2677  0.7619  0.0373  414.3989  553.0892  1.0134  …       │
│         ϵ_t[5]   2.4524  0.7702  0.0355  454.6820  593.7017  1.0021  …       │
│         ϵ_t[6]   1.6429  0.6962  0.0252  769.1307  607.3802  1.0058  …       │
│         ϵ_t[7]   0.8151  0.6025  0.0304  388.3976  555.2656  0.9991  …       │
│         ϵ_t[8]   0.6702  0.5175  0.0269  373.9920  303.0110  1.0051  …       │
│         ϵ_t[9]  -0.5696  0.4795  0.0232  418.2356  542.5845  1.0004  …       │
│        ϵ_t[10]  -2.2891  0.5501  0.0481  128.3031  395.4316  1.0074  …       │
│        ϵ_t[11]  -1.6915  0.4758  0.0420  119.1035  349.9149  1.0069  …       │
│        ϵ_t[12]   0.5292  0.4245  0.0262  318.2400  326.4057  1.0070  …       │
│        ϵ_t[13]   1.0827  0.4036  0.0292  228.5777  302.7830  1.0013  …       │
│        ϵ_t[14]   0.0193  0.3692  0.0234  244.3967  257.3761  1.0003  …       │
│        ϵ_t[15]   1.2264  0.4422  0.0395  123.4034  238.2643  1.0124  …       │
│        ϵ_t[16]  -1.0948  0.4090  0.0260  257.4322  435.9165  1.0015  …       │
│        ϵ_t[17]  -0.4379  0.3690  0.0216  287.5688  278.7597  1.0011  …       │
│        ϵ_t[18]  -0.7655  0.4105  0.0263  177.5683  309.2109  1.0062  …       │
│        ϵ_t[19]  -0.6802  0.3633  0.0216  291.5474  279.8767  1.0017  …       │
│        ϵ_t[20]  -0.5875  0.3674  0.0240  257.2188  159.0038  1.0030  …       │
│        ϵ_t[21]   0.2320  0.3324  0.0202  280.5701  462.2522  1.0012  …       │
│        ϵ_t[22]  -0.0444  0.3348  0.0186  348.1135  274.3257  1.0047  …       │
│        ϵ_t[23]  -0.5480  0.3639  0.0226  224.6069  414.0286  1.0084  …       │
│        ϵ_t[24]  -0.8776  0.3422  0.0130  687.0087  579.6001  1.0048  …       │
│        ϵ_t[25]  -1.3686  0.4038  0.0256  285.4126  247.6075  1.0018  …       │
│        ϵ_t[26]   1.0378  0.4609  0.0295  254.3289  273.5258  1.0016  …       │
│        ϵ_t[27]  -1.0605  0.4395  0.0235  362.7851  370.3438  1.0005  …       │
│        ϵ_t[28]  -0.0188  0.4343  0.0227  383.5133  394.5985  1.0017  …       │
│        ϵ_t[29]   0.1797  0.4394  0.0254  294.0063  416.9311  1.0008  …       │
│        ϵ_t[30]  -0.2995  0.4476  0.0218  423.7260  297.4228  1.0007  …       │
│        ϵ_t[31]   0.1998  0.4433  0.0216  381.7717  400.8702  1.0059  …       │
│        ϵ_t[32]   0.5553  0.4674  0.0427  123.9313  453.2743  1.0102  …       │
│        ϵ_t[33]   0.6466  0.4418  0.0219  396.9119  400.3757  1.0023  …       │
│        ϵ_t[34]   1.3274  0.4428  0.0491   90.2933   46.9219  1.0111  …       │
│   init_incide…  -0.0431  0.0971  0.0051  361.3629  618.5938  1.0037  …       │
│   cluster_fac…   0.0979  0.0645  0.0081   61.8622  105.1723  1.0002  …       │
╰──────────────────────────────────────────────────────────────────────────────╯

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.

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).