AdvancedHMC.jl

Documentation for AdvancedHMC.jl ↗

FlexiChains.to_nt_and_stats is overloaded for AdvancedHMC.Transition, so you can sample with AdvancedHMC into a FlexiChain{Symbol}. This is a slightly simplified version of the example in the AdvancedHMC README (here we use an analytical gradient rather than automatic differentiation):

using AdvancedHMC, AbstractMCMC
using LogDensityProblems
using FlexiChains: FlexiChain

# Set up AD-aware log-density function
struct LogTargetDensity
    dim::Int
end
LogDensityProblems.logdensity(::LogTargetDensity, θ) = -sum(abs2, θ) / 2
LogDensityProblems.logdensity_and_gradient(::LogTargetDensity, θ) = (-sum(abs2, θ) / 2, -θ)
LogDensityProblems.dimension(p::LogTargetDensity) = p.dim
function LogDensityProblems.capabilities(::Type{LogTargetDensity})
    return LogDensityProblems.LogDensityOrder{1}()
end

chn = AbstractMCMC.sample(
    LogTargetDensity(10),
    AdvancedHMC.NUTS(0.8),
    20;
    n_adapts=10,
    chain_type=FlexiChain{Symbol},
)

╭─FlexiChain (20 iterations, 1 chain) ─────────────────────────────────────────╮
│ ↓ iter  = 1:20                                                               │
│ → chain = 1:1                                                                │
│                                                                              │
│ Parameters (1) ── Symbol                                                     │
│  DimensionalData.DimVector{Float64, Tupl…  params (10,)                      │
│                                                                              │
│ Extras (12)                                                                  │
│  Int64    n_steps, tree_depth                                                │
│  Bool     is_accept, numerical_error, is_adapt                               │
│  Float64  acceptance_rate, log_density, hamiltonian_energy                   │
│           hamiltonian_energy_error, max_hamiltonian_energy_error, step_size  │
│           nom_step_size                                                      │
╰──────────────────────────────────────────────────────────────────────────────╯