Skip to content

Fitting a composed distribution

A tree built with ComposedDistributions is fittable as it stands. Loading both packages activates an extension that maps the tree's own params_table onto this package's row schema, so every verb from Fitting a custom distribution works on the tree.

This tutorial fits a two-event pathway through both routes, then fits a partially pooled tree.

julia
using DistributionsInference, Distributions, Random
using ComposedDistributions
using ComposedDistributions: compose, uncertain, pool, event

tree = compose((
    onset_admit = uncertain(Gamma(2.0, 1.0); shape = LogNormal(log(2.0), 0.2)),
    admit_death = LogNormal(0.5, 0.4)))
Parallel (2 branches)
├─ onset_admit: uncertain(Distributions.Gamma{Float64}(α=2.0, θ=1.0); shape = Distributions.LogNormal{Float64}(μ=0.6931471805599453, σ=0.2))
└─ admit_death: Distributions.LogNormal{Float64}(μ=0.5, σ=0.4)

What the tree declares

parameter_rows reports every parameter in the tree, keyed by edge.parameter. The onset_admit shape carries the prior attached by uncertain, so it is estimated; everything else has no prior and stays at its value.

julia
parameter_rows(tree)
4-element Vector{NamedTuple{(:name, :value, :prior, :support)}}:
 (name = Symbol("onset_admit.shape"), value = 2.0, prior = Distributions.LogNormal{Float64}(μ=0.6931471805599453, σ=0.2), support = (0.0, Inf))
 (name = Symbol("onset_admit.scale"), value = 1.0, prior = nothing, support = (0.0, Inf))
 (name = Symbol("admit_death.mu"), value = 0.5, prior = nothing, support = (0.0, Inf))
 (name = Symbol("admit_death.sigma"), value = 0.4, prior = nothing, support = (0.0, Inf))

One row of the four carries a prior, so the fit has one parameter.

julia
DistributionsInference.flat_dimension(tree)
1

Fitting it

A tree simulates the records it scores, so the data here are 200 draws from the tree itself, each a named delay per event. A fit should land near the shape those draws were generated at, 2.

julia
rng = Xoshiro(1)
tree_data = [rand(rng, tree) for _ in 1:200]
tree_data[1]
(onset_admit = 2.417469545474557, admit_death = 4.405247894848771)

From here the calls are the ones the hand-written distribution used, pointed at tree. distribution_to_advancedmh samples on the unconstrained scale, so a tree's own dotted row names come back on the chain with no manual construction and no hand-written support guard.

julia
using AdvancedMH, Bijectors
using LinearAlgebra: I

dim = DistributionsInference.flat_dimension(tree)
sampler = RWMH(MvNormal(zeros(dim), 0.05^2 * I))

Random.seed!(1)
chain = distribution_to_advancedmh(tree, tree_data, sampler, 2000; burnin = 1000)
fitted = inference_to_distribution(tree, chain, mean)
Parallel (2 branches)
├─ onset_admit: Distributions.Gamma{Float64}(α=2.126007203289852, θ=1.0)
└─ admit_death: Distributions.LogNormal{Float64}(μ=0.5, σ=0.4)

The fit comes back as a tree, so its nodes are reachable by name.

julia
event(fitted, :onset_admit)
Distributions.Gamma{Float64}(α=2.126007203289852, θ=1.0)

distribution_to_turing samples the same model over a tree, one site per estimated row.

julia
using DynamicPPL, Turing

Random.seed!(1)
turing_chain = distribution_to_turing(tree, tree_data, NUTS(), 500;
    progress = false)
event(inference_to_distribution(tree, turing_chain, mean), :onset_admit)
Distributions.Gamma{Float64}(α=2.1342935145014366, θ=1.0)

Partial pooling

pool ties a parameter across branches through a shared population distribution, so three districts share what they know about their delay without being forced to agree.

julia
population = uncertain(LogNormal(log(2.0), 0.3);
    mu = Normal(log(2.0), 0.2),
    sigma = truncated(Normal(0.0, 0.3); lower = 0.0))
pooled = compose((
    north = uncertain(Gamma(2.0, 1.0); shape = pool(:district, population)),
    east = uncertain(Gamma(2.0, 1.0); shape = pool(:district, population)),
    south = uncertain(Gamma(2.0, 1.0); shape = pool(:district, population))))

[row.name for row in DistributionsInference.estimated_rows(pooled)]
5-element Vector{Symbol}:
 Symbol("district.mu")
 Symbol("district.sigma")
 Symbol("north.shape.z")
 Symbol("east.shape.z")
 Symbol("south.shape.z")

A location-scale population is reparameterised non-centred, so what is estimated is the population's two hyperparameters and one standard normal offset per district. The estimation boundary moved and the fitting code is unchanged.

julia
pooled_data = [rand(rng, pooled) for _ in 1:200]
Random.seed!(1)
pooled_chain = distribution_to_turing(pooled, pooled_data, NUTS(0.9), 500;
    progress = false, initial_params = InitFromPrior())
╭─FlexiChain (500 iterations, 1 chain) ────────────────────────────────────────
 ↓ iter  = 251:750
 → chain = 1:1

 Parameters (5) ── AbstractPPL.VarName
  Float64  d.district.mu, d.district.sigma, d.north.shape.z, d.east.shape.z,  
           d.south.shape.z                                                    

 Extras (14)
  Int64    n_steps, tree_depth                                                
  Bool     is_accept, numerical_error                                         
  Float64  acceptance_rate, log_density, hamiltonian_energy,                  
           hamiltonian_energy_error, max_hamiltonian_energy_error, step_size, 
           nom_step_size, logprior, loglikelihood, logjoint                   
╰──────────────────────────────────────────────────────────────────────────────╯

The readback puts the offsets back through the population, so a district's shape comes out rather than its offset.

julia
event(inference_to_distribution(pooled, pooled_chain, mean), :north)
Distributions.Gamma{Float64}(α=1.9782060890084385, θ=1.0)

InitFromPrior above is Turing's own. The default initialisation draws uniformly on the unconstrained scale, which for a hierarchical shape can start the chain at exp(16) and underflow the likelihood.

The one tree Turing refuses

A centred pool scores its members against the reconstructed population rather than against a fixed prior of their own, and DynamicPPL has no sampling path for that yet. distribution_to_turing says so instead of mis-scoring the model.

julia
centred_pool() = pool(:region, LogNormal(log(2.0), 0.3); noncentred = false)
centred = compose((
    north = uncertain(Gamma(2.0, 1.0); shape = centred_pool()),
    south = uncertain(Gamma(2.0, 1.0); shape = centred_pool())))

centred_data = [rand(rng, centred) for _ in 1:200]

try
    distribution_to_turing(centred, centred_data)
catch err
    println(sprint(showerror, err))
end
ArgumentError: distribution_to_turing does not support estimated parameter(s) [Symbol("north.shape"), Symbol("south.shape")] with no fixed `~` prior (scored instead through `extra_logprior`, an object-dependent prior term whose sampling path does not exist yet in DynamicPPL). Sample with `distribution_to_logdensity(obj, data)` + LogDensityProblemsAD (the LogDensityProblems extension) instead.

The log-density route has no such gap. A centred row's extra_logprior term has no per-row prior of its own, so distribution_to_advancedmh cannot build its unconstrained transform from it either — the same limitation distribution_to_turing has, for the same reason. distribution_to_logdensity plus a gradient-free sampler driven by hand, on the constrained scale directly, has no such gap; the trade is the hand-written -Inf guard a random-walk proposal needs there, since nothing stops it stepping shape negative.

julia
centred_prob = distribution_to_logdensity(centred, centred_data)
centred_model = AdvancedMH.DensityModel() do x
    any(<=(0), x) ? -Inf : DistributionsInference.logdensity(centred_prob, x)
end
centred_sampler = RWMH(MvNormal(zeros(2), 0.05^2 * I))
centred_transitions = sample(Xoshiro(1), centred_model, centred_sampler, 2000;
    param_names = ["north.shape", "south.shape"], progress = false)
centred_draws = [t.params for t in centred_transitions][1001:end]
centred_chain = DistributionsInference.draws_to_chain(centred, centred_draws)
event(inference_to_distribution(centred, centred_chain, mean), :north)
Distributions.Gamma{Float64}(α=1.907757181953803, θ=1.0)

Next

  • ComposedDistributions' verb map covers compose, uncertain, pool and update, the verbs that built the trees here.

  • Public API lists the protocol this extension implements on a tree's behalf.