ComposedDistributions.jl

Appearance

Multi-strata trees and parameter uncertainty ​

Introduction ​

A composed tree is stationary by default: every leaf is a fixed distribution. Real delays are often non-stationary — a delay shortens over a wave, or differs by region — and some parameters are not known but estimated. ComposedDistributions handles both by generalising the leaf: a varying leaf reads an observed covariate, and an uncertain leaf carries distribution-valued parameters.

This tutorial builds a natural-history tree that varies by region and calendar time, resolves it per stratum, ties a delay across strata, and then makes a parameter uncertain so it can be estimated. It builds on Composing distributions and the time-, strata-, and covariate-varying reference.

julia
using ComposedDistributions
using Distributions
using Random

Varying by region and time ​

We build an onset-to-admission delay that grows with calendar time and an admission-to-death delay that differs by region. Both are ordinary leaves, so they drop into compose unchanged: a time-varying leaf reads the default :time covariate, and a region-varying leaf names its covariate and gives a reference for when none is supplied.

julia
onset_admit = varying(t -> Gamma(2.0, 1.0 + 0.02t))

admit_death = varying(
    r -> r === :north ? LogNormal(0.5, 0.4) : LogNormal(0.8, 0.3);
    covariate = :region, reference = LogNormal(0.5, 0.4))

template = compose((onset_admit = onset_admit, admit_death = admit_death))
Parallel (2 branches)
├─ onset_admit: Varying(time -> Distributions.Gamma{Float64}(α=2.0, θ=1.0))
└─ admit_death: Varying(region -> Distributions.LogNormal{Float64}(μ=0.5, σ=0.4))

The template still carries varying leaves, so it is not yet ready to score.

julia
(template = has_varying(template),)
(template = true,)

Resolving per stratum ​

instantiate resolves a whole tree against a Context. We combine an observed time and region with with_covariates and resolve one concrete tree per stratum.

julia
north = instantiate(template,
    with_covariates(Context(region = :north); time = 5.0))

south = instantiate(template,
    with_covariates(Context(region = :south); time = 5.0))
Parallel (2 branches)
├─ onset_admit: Distributions.Gamma{Float64}(α=2.0, θ=1.1)
└─ admit_death: Distributions.LogNormal{Float64}(μ=0.8, σ=0.3)

Each resolved tree is concrete and ready to score.

julia
(north = has_varying(north), south = has_varying(south))
(north = false, south = false)

The region-varying admission delay differs between the two strata.

julia
(north = event(north, :admit_death), south = event(south, :admit_death))
(north = Distributions.LogNormal{Float64}(μ=0.5, σ=0.4), south = Distributions.LogNormal{Float64}(μ=0.8, σ=0.3))

Sharing a parameter across strata ​

Some parameters are common to every stratum, such as a reporting delay recorded the same way everywhere. shared tags a leaf so its occurrences are one free parameter, and the tag survives instantiate, so the leaf is identical in every resolved stratum.

julia
report = shared(:report, Gamma(1.5, 1.0))

tied_template = compose((onset_admit = onset_admit,
    admit_death = admit_death, onset_report = report))

north_tied = instantiate(tied_template,
    with_covariates(Context(region = :north); time = 5.0))

south_tied = instantiate(tied_template,
    with_covariates(Context(region = :south); time = 5.0))
Parallel (3 branches)
├─ onset_admit: Distributions.Gamma{Float64}(α=2.0, θ=1.1)
├─ admit_death: Distributions.LogNormal{Float64}(μ=0.8, σ=0.3)
└─ onset_report: shared(:report, Distributions.Gamma{Float64}(α=1.5, θ=1.0))

The tied reporting delay is the same leaf in both strata, while the admission delay still differs.

julia
(north_report = event(north_tied, :onset_report),
    south_report = event(south_tied, :onset_report))
(north_report = shared(:report, Distributions.Gamma{Float64}(α=1.5, θ=1.0)), south_report = shared(:report, Distributions.Gamma{Float64}(α=1.5, θ=1.0)))

params_table inventories the tied leaf once, under its tag, rather than once per stratum.

julia
unique(params_table(north_tied).edge)
3-element Vector{Symbol}:
 :onset_admit
 :admit_death
 :report

Adding parameter uncertainty ​

A stratum resolves the observed covariates, but some parameters are still unknown and estimated. An uncertain leaf declares that directly: a parameter given a distribution is drawn from it rather than fixed, and that distribution is the parameter's prior. Here the admission-to-death location mu is uncertain.

julia
est_template = compose((
    onset_admit = varying(t -> Gamma(2.0, 1.0 + 0.02t)),
    admit_death = uncertain(LogNormal(0.5, 0.4); mu = Normal(0.5, 0.2))))
Parallel (2 branches)
├─ onset_admit: Varying(time -> Distributions.Gamma{Float64}(α=2.0, θ=1.0))
└─ admit_death: uncertain(Distributions.LogNormal{Float64}(μ=0.5, σ=0.4); mu = Distributions.Normal{Float64}(μ=0.5, σ=0.2))

Resolving the stratum leaves the uncertain leaf in place: instantiate fills the observed covariates, and the uncertain parameter is resolved later by a fit.

julia
resolved = instantiate(est_template, Context(time = 5.0))

(has_varying = has_varying(resolved), has_uncertain = has_uncertain(resolved))
(has_varying = false, has_uncertain = true)

params_table carries the uncertain parameter's prior on its prior column, so build_priors picks it up with no separate override.

julia
tbl = params_table(resolved)
(edge = tbl.edge, param = tbl.param, prior = tbl.prior)
(edge = [:onset_admit, :onset_admit, :admit_death, :admit_death], param = [:shape, :scale, :mu, :sigma], prior = Any[nothing, nothing, Distributions.Normal{Float64}(μ=0.5, σ=0.2), nothing])

Only rand reports the marginal: it draws the uncertain parameter from its prior, rebuilds the leaf, then draws the record.

julia
rand(Xoshiro(1), resolved)
(onset_admit = 1.7349142774838144, admit_death = 3.7487635604187353)

Every other query — logpdf, mean, and the rest — silently uses the template value while a leaf is still uncertain, so guard a scoring or fitting loop with has_uncertain. update collapses the uncertain leaf to a concrete one, and the guard then passes.

julia
fitted = update(resolved, (onset_admit = (shape = 2.0, scale = 1.1),
    admit_death = (mu = 0.6, sigma = 0.4)))

(before = has_uncertain(resolved), after = has_uncertain(fitted))
(before = true, after = false)

The collapsed tree is fully concrete, so logpdf scores a record.

julia
logpdf(fitted, rand(Xoshiro(2), fitted))
-1.9820517919151672

Uncertain-first estimation ​

The uncertain surface is the direct way to say "this parameter is estimated, with this prior": the spec is the prior and the declaration in one, and every parameter without a spec stays fixed. The estimation layer keys off exactly these specs. flatten / unflatten / flat_dimension and as_logdensity target the spec'd parameters only, so a tree with no uncertain leaves estimates nothing (a pure likelihood at the fixed tree), and the flat table is a derived view.

julia
flat = ComposedDistributions.flat_dimension(resolved)
(estimated_parameters = flat,)
(estimated_parameters = 1,)

update is the verb that moves the estimation boundary. A distribution in a parameter slot makes just that parameter uncertain (a partial update); update(tree, param_priors(tree)) promotes every free parameter to uncertain with support-derived default priors, the explicit estimate-everything path.

julia
promoted = update(resolved, param_priors(resolved))
(before = ComposedDistributions.flat_dimension(resolved),
    after = ComposedDistributions.flat_dimension(promoted))
(before = 1, after = 4)

Partial pooling across strata — estimating region-specific parameters that shrink towards a shared mean, rather than the fully independent (per-stratum) or fully tied (shared) extremes shown here — is designed in issue #23 and is not yet a built verb. For now a parameter is either independent per stratum or tied across all strata.

Summary ​

Where next ​