Skip to content

Sampling with Turing

The log-density is PPL-neutral, so Turing is a layer on top of it. Loading DynamicPPL activates distribution_to_turing, which wraps the same problem as a DynamicPPL model: each estimated row becomes a named site drawn from its own prior, and the data likelihood is added from the distribution rebuilt at the draw.

This tutorial fits the delay distribution from Fitting a custom distribution with NUTS, reads the chain back with the same two verbs, and swaps in a different likelihood.

The distribution and its two protocol methods are repeated so this page runs on its own.

julia
using DistributionsInference, Distributions, Random
using DynamicPPL, Turing
using FlexiChains: VNChain

struct ToyDelay{T <: Real}
    shape::T
    scale::T
end

function Distributions.logpdf(d::ToyDelay, y::Real)
    return logpdf(Weibull(d.shape, d.scale), y)
end

function DistributionsInference.parameter_rows(d::ToyDelay)
    return [
        (name = :shape, value = d.shape,
            prior = LogNormal(log(2.0), 0.2), support = (0.0, Inf)),
        (name = :scale, value = d.scale, prior = nothing,
            support = (0.0, Inf))]
end

function DistributionsInference.reconstruct(d::ToyDelay, x::AbstractVector)
    return ToyDelay(x[1], oftype(x[1], d.scale))
end

delay = ToyDelay(2.0, 2.5)
data = [1.5, 2.0, 3.2, 1.8, 2.6]
5-element Vector{Float64}:
 1.5
 2.0
 3.2
 1.8
 2.6

The model

distribution_to_turing takes the distribution and its data, exactly as distribution_to_logdensity does. Its sites are the estimated rows, named as the protocol named them under the model's prefix (d by default, set with the prefix keyword). The fixed scale row has no prior, so it is not a site.

julia
model = distribution_to_turing(delay, data)
keys(VarInfo(model))
1-element Vector{AbstractPPL.VarName}:
 d.shape

It samples like any other Turing model. NUTS differentiates through reconstruct, which is why the fields of ToyDelay are typed loosely enough to carry a dual number.

julia
Random.seed!(1)
chain = sample(model, NUTS(), 500; chain_type = VNChain, progress = false)
╭─FlexiChain (500 iterations, 1 chain) ────────────────────────────────────────
 ↓ iter  = 251:750
 → chain = 1:1

 Parameters (1) ── AbstractPPL.VarName
  Float64  d.shape                                                            

 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                   
╰──────────────────────────────────────────────────────────────────────────────╯

Reading the chain back

The sites are keyed by the same dotted names the protocol declared, under the model's prefix, so inference_to_distribution and inference_to_distributions read a VNChain exactly as they read a hand-rolled chain.

julia
inference_to_distribution(delay, chain, mean).shape
2.3142742102297116

Keeping every draw gives a posterior-predictive summary.

julia
fits = inference_to_distributions(delay, chain)
quantile([mean(Weibull(d.shape, d.scale)) for d in fits], [0.025, 0.5, 0.975])
3-element Vector{Float64}:
 2.2140132279172713
 2.2177960614157413
 2.2448267537231392

The readback contract is the dotted names, so switching sampler leaves this code alone.

A different likelihood

distribution_to_turing takes the same loglik reducer as distribution_to_logdensity. Aggregated records, a delay and the number of cases reporting it, score through a weighted sum rather than one term per row, and NUTS samples the result with no other change.

julia
weighted_loglik(obj, records) = sum(w * logpdf(obj, y) for (y, w) in records)
counts = [(1.5, 12.0), (2.0, 30.0), (3.2, 8.0), (1.8, 21.0), (2.6, 15.0)]
Random.seed!(1)
weighted_chain = sample(
    distribution_to_turing(delay, counts; loglik = weighted_loglik),
    NUTS(), 500; chain_type = VNChain, progress = false)
inference_to_distribution(delay, weighted_chain, mean).shape
3.946036572608266

The survival reducer from Fitting a custom distribution goes through the same keyword: logccdf for a Weibull differentiates, so NUTS samples that one unchanged too.

Next