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: 19)
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.6508952824878806, 1.9311755822668364, 0.5801240686071575]
Σ: [0.11095087456373508 -0.016539481905092007 -0.0015269282137262683; -0.016539481905092007 0.003408329839783224 0.00020558840423764497; -0.001526928213726264 0.00020558840423764454 0.004600836541139295]
)
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.62 (ok)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.6494 0.3125 0.0098 1015.4812 901.4612 1.0005 1.1273 1.6756 2.1899 │
│ β 1.9300 0.0551 0.0017 1007.6980 1023.0661 1.0012 1.8372 1.9307 2.0266 │
│ σ 1.8090 0.1265 0.0038 1116.1143 944.3411 0.9990 1.5988 1.8031 2.0357 │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯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.6431 0.3426 0.0056 3774.9233 3536.8047 1.0014 1.0777 1.6488 2.2006 │
│ β 1.9322 0.0600 0.0010 3759.1147 3993.0980 1.0019 1.8344 1.9312 2.0318 │
│ σ 1.8190 0.1295 0.0018 5245.6670 5036.1279 1.0006 1.6181 1.8120 2.0417 │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯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.6379 0.3339 0.0033 10228.6121 6607.2426 1.0000 1.0874 1.6411 2.1762 │
│ β 1.9330 0.0590 0.0006 10021.3751 5909.3915 1.0001 1.8360 1.9319 2.0312 │
│ σ 1.8134 0.1281 0.0013 9806.4848 6095.7680 1.0005 1.6168 1.8068 2.0362 │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯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)]