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 (rows are groups, columns are time) and passes n_groups 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:   . After the model's exp transformation, each group's curve is the shared curve scaled by . Swap combine for a different mapping the way Ascertainment swaps its transform.

Simulate

Passing an all-missing matrix makes the model a prior simulator; the group dimension threads from its row count. We simulate eight groups over 24 time steps:

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]))
# Z_t[g, :] is the shared random walk plus group g's pooled level. Averaging
# over time and subtracting the grand mean cancels the shared component and
# leaves each group's level relative to the others. Only that relative level
# is identified: a constant added to every group's level and subtracted from
# the shared path gives the same Z_t, so the two are confounded.
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, data_size = size(Ydata),
    true_levels = round.(true_levels, digits = 2))

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

The shared infection curve is common to all groups; 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. n_groups again threads from the data matrix, so nothing about the group dimension is hard-coded in the components:

posterior = as_turing_model(model, Ydata)
chain = sample(posterior, NUTS(0.85; adtype = Turing.AutoForwardDiff()), 300;
    progress = false)
size(chain, 1)

300

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.19, -0.16, 0.49, -0.11, -1.13, 0.79, -0.33, 0.25], correlation = 0.998)

The posterior per-group levels line up with the simulated truth. A plot with 80% credible intervals against the   line makes the recovery visible:

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. Stratify supplied the panel structure, the Hierarchy supplied the per-group levels, and the group dimension threaded from the data with the group prior namespaced by the component. 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, say, each with its own latent — see CombineInfections instead, described on the Multiple observation streams page.