ComposableTuringIDModels.jl

Appearance

3

Reporting delays and day-of-week effects ​

Real surveillance data is rarely a clean count of infections on the day they occur. Cases are reported after a delay — an incubation period followed by a reporting lag — and the number reported depends on the day of the week. Tools for real-time estimation such as those of Abbott et al. [3] build these features into the observation model so that the latent infection signal is estimated free of reporting artefacts.

This tutorial keeps the renewal infection core of the previous example but replaces the simple observation model with a layered one: infections are convolved through two delay distributions and then modulated by a day-of-week reporting pattern. It also shows the latent process as an ARIMA-style differenced process broadcast to a weekly timescale, and assembles everything with IDProblem. The model follows the configuration of the EpiNow2 package [3] and is fit to daily confirmed COVID-19 cases from Italy's first wave in 2020.

The model ​

where is the incubation-period pmf, the reporting-delay pmf, and a day-of-week reporting multiplier.

A weekly latent process ​

The latent process is an ARIMA(2,1,1): an AR/MA combination (arma) wrapped in a DiffLatentModel to difference it once. Differencing makes the level a random walk rather than mean-reverting, which suits a reproduction number that can drift.

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

arma21 = arma(
    init = [Normal(0, 0.2), Normal(0, 0.2)],
    damp = [truncated(Normal(0.1, 0.2), 0, 1), truncated(Normal(0.1, 0.05), 0, 1)],
    θ = [truncated(Normal(0.0, 0.2), -1, 1)],
    ϵ_t = HierarchicalNormal(std = HalfNormal(0.1)))

arima211 = DiffLatentModel(; model = arma21, init = [Normal(0.3, 0.3)])

broadcast_weekly makes the process piecewise-constant by week: a new value is drawn each week and held for seven days. This models as changing weekly rather than daily, which both regularises the estimate and cuts the number of latent parameters.

julia
weekly_latent = broadcast_weekly(arima211)

The infection process ​

As before, a Renewal process driven by a discretised generation interval. Here we use a generation time. The weekly process built above is folded into the renewal model's rt slot.

julia
renewal = Renewal(; generation_time = Gamma(1.4, 1 / 0.38),
    rt = weekly_latent, initialisation = Normal(log(1.0), 1.0))

A layered observation model ​

We start from the NegativeBinomialError link and build outward. ascertainment_dayofweek wraps it with a partially pooled day-of-week multiplier, so reporting can be systematically higher or lower on particular weekdays.

julia
negbin = NegativeBinomialError(cluster_factor = HalfNormal(0.1))
dayofweek_negbin = ascertainment_dayofweek(
    negbin; latent_model = HierarchicalNormal(std = HalfNormal(1.0)))

LatentDelay convolves the expected observations with a delay distribution (discretised by double interval censoring). Two layers compose sequentially: a fixed incubation period from infection to symptom onset, then a reporting delay from onset to report whose parameters are inferred. The reporting delay is an UncertainDelay: its LogNormal log-scale mean and standard deviation carry priors, so the delay is rediscretised each draw and estimated jointly with the reproduction number rather than fixed from external data.

julia
incubation = LogNormal(1.6, 0.42)   # infection -> symptom onset (fixed)
reporting = UncertainDelay(         # symptom onset -> report (inferred)
    LogNormal, [Normal(0.58, 0.3), truncated(Normal(0.47, 0.2), 0, Inf)];
    D = 8.0)

observation = LatentDelay(LatentDelay(dayofweek_negbin, incubation), reporting)

That single observation object now carries, from the inside out: a negative binomial link, a day-of-week ascertainment modifier, a fixed incubation-delay convolution, and an inferred reporting-delay convolution — assembled entirely by composition. The reporting-delay parameters flow through the same priors seam as every other parameter, so inferring the delay needs no change to the rest of the model.

The data ​

We fit the model to the daily confirmed COVID-19 cases from Italy's first wave (the example series shipped with the EpiNow2 package), stored with the docs.

julia
using CSV, DataFrames
datapath = joinpath(pkgdir(ComposableTuringIDModels),
    "docs", "src", "tutorials", "data", "italy_data.csv")
italy = CSV.read(datapath, DataFrame)
n = 42
y_obs = italy.confirm[1:n]
(n = n, total_cases = sum(y_obs), from = italy.date[1], to = italy.date[n])
(n = 42, total_cases = 115239, from = Dates.Date("2020-02-22"), to = Dates.Date("2020-04-03"))

Assemble and fit ​

IDProblem ties the latent, infection, and observation models to a time span. Its as_turing_model method takes data as a named tuple with a y_t field (passing missing values would instead simulate from the prior).

julia
problem = IDProblem(
    infection = renewal,
    observation_model = observation,
    tspan = (1, n))

Fitting conditions on the observed reports, differentiating with the recommended Mooncake backend (see Automatic differentiation backend). We draw two chains in parallel with MCMCThreads(), which gives a cross-chain :

julia
posterior = as_turing_model(problem, (y_t = y_obs,))
chain = sample(
    posterior, NUTS(0.95; adtype = AutoMooncake(; config = nothing)),
    MCMCThreads(), 250, 2; progress = false)
┌ Info: Found initial step size
└   ϵ = 0.003125
┌ Info: Found initial step size
└   ϵ = 0.003125

sample returns a FlexiChains chain, which summarystats summarises directly — no conversion step. The day-of-week scale (DayofWeek.std), the negative-binomial overdispersion (cluster_factor) and the inferred reporting-delay parameters (delay.θ, the LogNormal log-mean and log-sd) appear alongside the latent-process parameters:

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 (22) ── AbstractPPL.VarName                                       │
│  Float64  latent_init[1], ar_init[1], ar_init[2], damp_AR[1], damp_AR[2], θ, │
│           std, ϵ_t[1], ϵ_t[2], ϵ_t[3], init_incidence, delay.θ[1],           │
│           delay.θ[2], DayofWeek.std, DayofWeek.ϵ_t[1], DayofWeek.ϵ_t[2],     │
│           DayofWeek.ϵ_t[3], DayofWeek.ϵ_t[4], DayofWeek.ϵ_t[5],              │
│           DayofWeek.ϵ_t[6], DayofWeek.ϵ_t[7], 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  …       │
│   latent_init…   1.0338  0.1751  0.0159  126.7528  169.0446  1.0223  …       │
│     ar_init[1]  -0.0462  0.1731  0.0066  720.7868  410.3861  1.0017  …       │
│     ar_init[2]  -0.5382  0.1433  0.0128  129.4321  210.3263  1.0130  …       │
│     damp_AR[1]   0.3686  0.1512  0.0168   81.7719  112.4166  0.9999  …       │
│     damp_AR[2]   0.1128  0.0485  0.0036  180.8593  120.1818  1.0130  …       │
│              θ   0.0127  0.2076  0.0098  457.5837  340.6761  1.0026  …       │
│            std   0.1210  0.0672  0.0063  104.1615  165.5919  1.0054  …       │
│         ϵ_t[1]  -0.7253  0.8009  0.0768  110.8607  184.0332  0.9986  …       │
│         ϵ_t[2]  -0.6863  0.7444  0.0617  149.2828  183.4489  1.0008  …       │
│         ϵ_t[3]   0.0507  0.9802  0.1563   39.1821   97.9795  1.0226  …       │
│   init_incide…   3.1234  0.7877  0.0860   82.8104  108.8785  1.0234  …       │
│     delay.θ[1]   0.4717  0.3053  0.0103  845.5640  178.2790  1.0115  …       │
│     delay.θ[2]   0.4625  0.1939  0.0111  284.9462  133.3817  0.9989  …       │
│   DayofWeek.s…   0.1233  0.0634  0.0066  115.2592  158.5534  1.0035  …       │
│   DayofWeek.ϵ…   0.0591  0.6095  0.0559  124.4470  131.4159  1.0074  …       │
│   DayofWeek.ϵ…   0.7138  0.5532  0.0418  171.6197  249.9701  0.9991  …       │
│   DayofWeek.ϵ…   0.4738  0.6176  0.0490  160.5283  210.6864  1.0036  …       │
│   DayofWeek.ϵ…   1.0627  0.6248  0.0476  172.4484  241.1844  1.0064  …       │
│   DayofWeek.ϵ…  -0.2955  0.6746  0.0592  132.0423  169.0863  1.0088  …       │
│   DayofWeek.ϵ…  -0.7438  0.6490  0.0544  148.6868  184.7880  1.0000  …       │
│   DayofWeek.ϵ…  -0.9541  0.6636  0.0599  123.0045  259.7977  1.0005  …       │
│   cluster_fac…   0.1067  0.0209  0.0012  347.1404  264.6371  1.0059  …       │
╰──────────────────────────────────────────────────────────────────────────────╯

DayofWeek.std is the scale of the partially pooled weekday multipliers (its own block, namespaced because the ascertainment modifier introduces a named sub-process); cluster_factor is the negative-binomial overdispersion; delay.θ are the inferred reporting-delay parameters. The day-of-week effect, the reporting delay, and the weekly reproduction number were all estimated jointly — and any of them can be swapped, fixed, or removed by editing one line of the composition above.

Prior versus posterior ​

Sampling the same model with Prior gives a prior draw over the same parameters. Overlaying it on the posterior with PairPlots.jl — the FlexiChains extension turns each chain, subset to a few keys, into a PairPlots.Series — shows which parameters the six weeks of Italian data moved.

julia
using CairoMakie, PairPlots

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

The innovation scale (std) and the negative-binomial overdispersion (cluster_factor) tighten under the data, while the autoregressive damping (damp_AR) and moving-average (θ) coefficients of the ARIMA process stay close to their weakly informative priors.

Posterior trajectories ​

  and the infections are generated quantities recovered per draw with generated_observables; the reports are scored element-wise, so their posterior-predictive distribution comes from predict on the model with the observations set to missing. Two small helpers reduce the per-draw trajectories to credible bands.

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

pred = predict(as_turing_model(problem, (y_t = fill(missing, 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 = "Confirmed 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 weekly is piecewise-constant by construction, stepping down through one as the first wave turns over. The posterior-predictive band starts partway into the series — the two delay convolutions leave the earliest reference days without a fully supported expected count — and from there tracks the observed Italian reports, the layered observation model having absorbed the reporting pattern rather than the infection signal.

A time-varying reporting pattern ​

The day-of-week multiplier above is static: one weekly profile held fixed across the series. Reporting behaviour can itself drift — testing capacity changes, weekend effects strengthen or weaken — and the same composition expresses that. Because the ascertainment modifier takes any latent model, replacing the pooled HierarchicalNormal weekday effect with a BroadcastLatentModel over a process that evolves week to week turns the fixed profile into a time-varying one, at the cost of more latent parameters. The structural change is again local to the observation model; the infection and latent parts are untouched. We keep the static pattern here — it is identifiable from six weeks of data, where a fully time-varying weekday process would not be — and flag the richer variant rather than fit it.

The reporting delay can drift in the same way, through the same seam. An UncertainDelay parameter is a prior slot like any other, so replacing its constant log-mean prior with a process — a RandomWalk — makes the delay distribution itself time-varying: it is rediscretised at each time point and applied with a per-time convolution, while the log-scale spread keeps a constant prior.

julia
drifting = UncertainDelay(
    LogNormal, [RandomWalk(), truncated(Normal(0.47, 0.2), 0, Inf)]; D = 8.0)
tv_observation = LatentDelay(
    LatentDelay(dayofweek_negbin, incubation), drifting)

Nothing else changes: the infection process, the prior, and the fitting code are identical — only which prior fills the delay's log-mean slot. As with the weekday profile we flag rather than fit it here, since a delay that drifts day to day asks more of six weeks of data than they can answer.

References ​

  1. S. Abbott, J. Hellewell, R. N. Thompson, K. Sherratt, H. P. Gibbs, N. I. Bosse, J. D. Munday, S. Meakin, E. L. Doughty, J. Y. Chun and others. Estimating the time-varying reproduction number of SARS-CoV-2 using national and subnational case counts. Wellcome Open Research 5, 112 (2020).