Running Pathfinder on Turing.jl models
This tutorial demonstrates how Turing can be used with Pathfinder.
We'll demonstrate with a regression example.
using AbstractMCMC, AdvancedHMC, DynamicPPL, FlexiChains, MCMCChains, Pathfinder, Random, Turing
Random.seed!(39)
@model function regress(x)
α ~ Normal()
β ~ Normal()
σ ~ truncated(Normal(); lower=0)
μ = α .+ β .* x
y ~ product_distribution(Normal.(μ, σ))
end
x = 0:0.1:10
true_params = (; α=1.5, β=2, σ=2)
# simulate data
y = rand(regress(x) | true_params)[@varname(y)]
model = regress(x) | (; y)
n_chains = 8For convenience, pathfinder and multipathfinder can take Turing models as inputs and produce FlexiChains.VNChain objects as outputs. To access this, we run Pathfinder normally; the chains representation of the draws is stored in draws_transformed, with the type defaulting to whatever Turing.sample itself defaults to.
Since Turing v0.45, the default chain_type is FlexiChains.VNChain, while previous versions returned MCMCChains.Chains. Pathfinder will return the same default chain type that your installed Turing version returns, but you can always specify the chain_type manually. This tutorial uses the new default (FlexiChains) throughout; see Using MCMCChains below if you still want MCMCChains.Chains.
result_single = pathfinder(model; ndraws=1_000)Single-path Pathfinder result
tries: 1
draws: 1000
fit iteration: 11 (total: 31)
fit ELBO: -213.6 ± 0.15
fit distribution: MvNormal{Float64, Pathfinder.WoodburyPDMat{Float64, LinearAlgebra.Diagonal{Float64, Vector{Float64}}, Matrix{Float64}, Matrix{Float64}, Pathfinder.WoodburyPDFactorization{Float64, LinearAlgebra.Diagonal{Float64, Vector{Float64}}, LinearAlgebra.QRCompactWYQ{Float64, Matrix{Float64}, Matrix{Float64}}, LinearAlgebra.UpperTriangular{Float64, Matrix{Float64}}}}, Vector{Float64}}(
dim: 3
μ: [1.650895282497865, 1.9311755822776513, 0.5801240686597915]
Σ: [0.11095087434004708 -0.016539481756590615 -0.00152692711421274; -0.016539481756590625 0.003408329785186465 0.00020558810663449505; -0.00152692711421274 0.00020558810663449226 0.004600836738960978]
)
result_single.draws_transformed╭─FlexiChain (1000 iterations, 1 chain) ───────────────────────────────────────────────────────────╮
│ ↓ iter = 1:1000 │
│ → chain = 1:1 │
│ │
│ Parameters (3) ── VarName │
│ Float64 α, β, σ │
│ │
│ Extras (3) │
│ Float64 logprior, loglikelihood, logjoint │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯Note that while Turing's sample methods default to initializing parameters from the prior with InitFromPrior, Pathfinder defaults to uniformly sampling them in the range [-2, 2] in unconstrained space (equivalent to Turing's InitFromUniform(-2, 2)). To use Turing's default in Pathfinder, specify init_sampler=InitFromPrior().
result_multi = multipathfinder(model, 1_000; nruns=n_chains, init_sampler=InitFromPrior())Multi-path Pathfinder result
runs: 8
draws: 1000
Pareto shape diagnostic: 0.29 (good)The Pareto shape diagnostic indicates that it is likely safe to use these draws to compute posterior estimates.
chns_pf = result_multi.draws_transformed
summarystats(chns_pf)╭─FlexiSummary (9 statistics) ─────────────────────────────────────────────────────────────────────╮
│ iter collapsed │
│ chain collapsed │
│ ↓ stat = [mean, std, mcse, ess_bulk, ess_tail, rhat, q5, q50, q95] │
│ │
│ Parameters (3) ── VarName │
│ Float64 α, β, σ │
│ │
│ Extras (3) │
│ Float64 logprior, loglikelihood, logjoint │
│ │
│ Summary │
│ param mean std mcse ess_bulk ess_tail rhat q5 q50 q95 │
│ α 1.6426 0.3392 0.0110 963.1483 944.3411 1.0009 1.0356 1.6398 2.1807 │
│ β 1.9313 0.0571 0.0019 906.3761 870.5023 0.9996 1.8425 1.9305 2.0247 │
│ σ 1.8209 0.1213 0.0039 948.5482 973.3333 0.9990 1.6167 1.8212 2.0218 │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯We can also use these draws to initialize MCMC sampling with InitFromParams. resample with replace=false selects n_chains distinct draws (importance-weighted) from the candidate pool; see Resampling for details.
init_result = resample(result_multi, n_chains; replace=false)
params = AbstractMCMC.to_samples(DynamicPPL.ParamsWithStats, init_result.draws_transformed, model)
initial_params = [InitFromParams(p.params) for p in vec(params)]chns = sample(model, Turing.NUTS(), MCMCThreads(), 1_000, n_chains; initial_params, progress=false)
summarystats(chns)╭─FlexiSummary (9 statistics) ─────────────────────────────────────────────────────────────────────╮
│ iter collapsed │
│ chain collapsed │
│ ↓ stat = [mean, std, mcse, ess_bulk, ess_tail, rhat, q5, q50, q95] │
│ │
│ Parameters (3) ── VarName │
│ Float64 α, β, σ │
│ │
│ Extras (14) │
│ Float64 n_steps, is_accept, acceptance_rate, log_density, hamiltonian_energy, │
│ hamiltonian_energy_error, max_hamiltonian_energy_error, tree_depth, numerical_error, │
│ step_size, nom_step_size, logprior, loglikelihood, logjoint │
│ │
│ Summary │
│ param mean std mcse ess_bulk ess_tail rhat q5 q50 q95 │
│ α 1.6532 0.3371 0.0055 3696.0689 4315.6339 0.9999 1.1032 1.6509 2.2098 │
│ β 1.9304 0.0593 0.0010 3798.0841 4082.7859 0.9999 1.8333 1.9302 2.0272 │
│ σ 1.8158 0.1267 0.0018 5255.3855 4714.4861 1.0007 1.6178 1.8119 2.0307 │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯We can use Pathfinder's estimate of the metric and only perform enough warm-up to tune the step size.
inv_metric = result_multi.pathfinder_results[1].fit_distribution.Σ
metric = AdvancedHMC.RankUpdateEuclideanMetric(inv_metric)
kernel = HMCKernel(Trajectory{MultinomialTS}(Leapfrog(0.0), GeneralisedNoUTurn()))
adaptor = StepSizeAdaptor(0.8, 1.0) # adapt only the step size
nuts = AdvancedHMC.HMCSampler(kernel, metric, adaptor)
n_adapts = 50
n_draws = 1_000
chns = sample(
model,
externalsampler(nuts),
MCMCThreads(),
n_draws + n_adapts,
n_chains;
n_adapts,
initial_params,
progress=false,
)[iter=(n_adapts + 1):(n_adapts + n_draws)] # drop warm-up draws
summarystats(chns)╭─FlexiSummary (9 statistics) ─────────────────────────────────────────────────────────────────────╮
│ iter collapsed │
│ chain collapsed │
│ ↓ stat = [mean, std, mcse, ess_bulk, ess_tail, rhat, q5, q50, q95] │
│ │
│ Parameters (3) ── VarName │
│ Float64 α, β, σ │
│ │
│ Extras (14) │
│ Float64 n_steps, is_accept, acceptance_rate, log_density, hamiltonian_energy, │
│ hamiltonian_energy_error, max_hamiltonian_energy_error, tree_depth, numerical_error, │
│ step_size, nom_step_size, logprior, loglikelihood, logjoint │
│ │
│ Summary │
│ param mean std mcse ess_bulk ess_tail rhat q5 q50 q95 │
│ α 1.6435 0.3345 0.0034 9528.7153 6685.8156 1.0004 1.0972 1.6455 2.1919 │
│ β 1.9328 0.0593 0.0006 9328.2981 6659.0277 1.0011 1.8372 1.9318 2.0313 │
│ σ 1.8145 0.1269 0.0013 9945.6452 6009.3354 1.0008 1.6191 1.8074 2.0312 │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯See Initializing HMC with Pathfinder for further examples.
Using MCMCChains
To request transformed draws be returned as MCMCChains.Chains instead of the default, you may specify chain_type directly.
result_multi_mcmc = multipathfinder(
model, 1_000; nruns=n_chains, init_sampler=InitFromPrior(), chain_type=MCMCChains.Chains
)
chns_pf_mcmc = result_multi_mcmc.draws_transformedChains MCMC chain (1000×6×1 Array{Float64, 3}):
Iterations = 1:1:1000
Number of chains = 1
Samples per chain = 1000
parameters = α, β, σ
internals = logprior, loglikelihood, logjoint
Use `describe(chains)` for summary statistics and quantiles.
As before, we can use these draws to initialize MCMC sampling with InitFromParams:
init_result_mcmc = resample(result_multi_mcmc, n_chains; replace=false)
params_mcmcchains = AbstractMCMC.to_samples(
DynamicPPL.ParamsWithStats, init_result_mcmc.draws_transformed, model
)
initial_params_mcmcchains = [InitFromParams(p.params) for p in vec(params_mcmcchains)]