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.89 (bad)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.6394 0.3496 0.0112 987.8035 962.5129 1.0024 0.9775 1.6398 2.2316 │
│ β 1.9325 0.0640 0.0020 1009.2928 864.3256 1.0008 1.8326 1.9306 2.0507 │
│ σ 1.8224 0.1229 0.0038 1072.2154 958.0427 0.9990 1.6287 1.8162 2.0283 │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯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.6485 0.3335 0.0054 3840.1329 4347.6686 1.0005 1.0913 1.6518 2.1919 │
│ β 1.9316 0.0582 0.0009 3759.6656 4241.0184 1.0006 1.8358 1.9311 2.0279 │
│ σ 1.8149 0.1247 0.0017 5188.7189 4968.5817 1.0003 1.6215 1.8082 2.0278 │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯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.6492 0.3354 0.0044 5720.8622 5802.7820 1.0000 1.0943 1.6481 2.1995 │
│ β 1.9309 0.0592 0.0009 4255.4731 5142.0801 1.0002 1.8330 1.9309 2.0278 │
│ σ 1.8154 0.1278 0.0015 6874.4542 5733.7980 1.0000 1.6176 1.8087 2.0345 │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯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)]