Partial pooling across groups in a composed model

A multi-group epidemic is a panel, one shared infection process seen by several groups, each reporting it at its own level. Stratify expresses this panel directly. A shared process is drawn once over time, and a per-group deviation is drawn once over the group axis. The two combine into one groups × time latent matrix. Split then observes every group through the same observation model, namespaced by group. See the Composable design page for how the two compose.

This page drives a per-group reporting level with a Hierarchy inside a Stratify. It simulates from that model and fits it end-to-end under NUTS, recovering the per-group levels. The group dimension threads from the data and the group prior is namespaced by the component, so the panel composes with no hand-orchestration.

The composed model

The shared epidemic is a DirectInfections process carrying a RandomWalk latent, observed with a PoissonError. Stratify puts a per-group level on that latent. Every group shares the same random walk, offset by its own log-level , so each group's infection curve is the shared curve scaled by . Those per-group levels are partially pooled with a Hierarchy, supplied as Stratify's across slot.

using ComposableTuringIDModels, Distributions, Turing, Random, Statistics
using Turing: returned
Random.seed!(77)

hierarchy = Hierarchy(; mean = Normal(0.0, 0.5), across = IID(Normal(0.0, 0.5)))
model = IDModel(
    DirectInfections(; Z = Stratify(RandomWalk(), hierarchy),
        initialisation = Normal(log(50.0), 0.2)),
    Split(PoissonError()))

IDModel
├─ infection: DirectInfections
│  └─ Z: Stratify
│     ├─ shared: RandomWalk
│     │  └─ ϵ_t: HierarchicalNormal
│     └─ across: Hierarchy
│        └─ across: IID
└─ observation: Split
   └─ template: PoissonError

The grouping dimension is not a field of any component. Passing as_turing_model(model, Y) reads n_groups and n_time from the shape of the data matrix Y, where rows are groups and columns are time. n_groups goes to the Hierarchy through Stratify's across slot, and n_time to the shared random walk, the same way a series length is passed to as_turing_model(latent, n).

The group prior carries its own innovations (the IID across process samples an ϵ_t), which under the prefix-off submodel convention would collide with the infection RandomWalk's own ϵ_t. Stratify prefixes its across slot automatically so the two never collide, and Split prefixes each group's observation with group<g>. Stratify's combine argument maps a group's level and the shared path onto that group's row of the latent matrix. By default this is additive,   . Swap combine for a different mapping the way Ascertainment swaps its transform.

Simulate

Passing an all-missing matrix makes the model a prior simulator. We simulate eight groups over 24 time steps. Only the relative group levels are identified. Adding a constant to every group's level and subtracting it from the shared path gives the same , so the two are confounded. The comparison below is between relative levels.

n_time, n_groups = 24, 8
Ymiss = Matrix{Union{Missing, Float64}}(missing, n_groups, n_time)
sim = as_turing_model(model, Ymiss)()
Ydata = Float64.(reduce(vcat,
    [permutedims(sim.generated_y_t[Symbol(:group, g)]) for g in 1:n_groups]))
group_level(Z) = vec(mean(Z; dims = 2)) .- mean(Z)
true_levels = group_level(sim.Z_t)
(n_time = n_time, n_groups = n_groups,
    true_levels = round.(true_levels, digits = 2))

(n_time = 24, n_groups = 8, true_levels = [0.18, -0.23, 0.48, -0.14, -1.09, 0.8, -0.27, 0.26])

The rows of Ydata differ only through the per-group level and Poisson noise.

Fit

Conditioning on the simulated counts and sampling with NUTS recovers the posterior end-to-end. We draw two chains in parallel with MCMCThreads() and differentiate with Mooncake, the recommended backend for this package (see Automatic differentiation backend).

using Mooncake
using ADTypes: AutoMooncake

posterior = as_turing_model(model, Ydata)
chain = sample(
    posterior, NUTS(0.85; adtype = AutoMooncake(; config = nothing)),
    MCMCThreads(), 300, 2; progress = false)

┌ Info: Found initial step size
└   ϵ = 0.0125
┌ Info: Found initial step size
└   ϵ = 1.220703125e-5

The relative per-group levels are recovered per draw from the generated Z_t with returned, then compared with the simulated truth.

level_draws = reduce(hcat,
    [group_level(g.Z_t) for g in vec(returned(posterior, chain))])
post_mean = vec(mean(level_draws; dims = 2))
(true_levels = round.(true_levels, digits = 2),
    posterior_means = round.(post_mean, digits = 2),
    correlation = round(cor(true_levels, post_mean), digits = 3))

(true_levels = [0.18, -0.23, 0.48, -0.14, -1.09, 0.8, -0.27, 0.26], posterior_means = [0.2, -0.16, 0.49, -0.11, -1.14, 0.79, -0.33, 0.25], correlation = 0.997)

The posterior per-group levels line up with the simulated truth.

using CairoMakie
qs = [quantile(level_draws[g, :], [0.1, 0.5, 0.9]) for g in 1:n_groups]
lo = getindex.(qs, 1)
md = getindex.(qs, 2)
hi = getindex.(qs, 3)

fig = Figure(; size = (620, 460))
ax = Axis(fig[1, 1]; xlabel = "True group level ℓ_g",
    ylabel = "Posterior level ℓ_g")
lims = (minimum(true_levels) - 0.4, maximum(true_levels) + 0.4)
lines!(ax, [lims...], [lims...]; color = :grey, linestyle = :dash)
rangebars!(ax, true_levels, lo, hi; color = :seagreen, whiskerwidth = 10)
scatter!(ax, true_levels, md; color = :seagreen, markersize = 12)
fig

Each group's credible interval covers the   line, so the partially pooled per-group levels are recovered inside a full composed panel. Swapping across = RandomWalk() in the Hierarchy would instead pool neighbouring groups (correlated ordered strata). Swapping Stratify's across slot for a bare IID or a Distribution would drop the shared level for independent per-group levels, each with no other change.

When the groups are genuinely separate infection processes rather than one shared curve, several distinct regions each with its own latent, see CombineInfections instead. That is described on the Multiple observation streams page.