AdvancedHMC.jl
AdvancedHMC.jl provides a robust, modular and efficient implementation of advanced HMC algorithms. An illustrative example for AdvancedHMC's usage is given below. AdvancedHMC.jl is part of Turing.jl, a probabilistic programming library in Julia. If you are interested in using AdvancedHMC.jl through a probabilistic programming language, please check it out!
Interfaces
IMP.hmc
: an experimental Python module for the Integrative Modeling Platform, which uses AdvancedHMC in its backend to sample protein structures.
NEWS
 We presented a paper for AdvancedHMC.jl at AABI 2019 in Vancouver, Canada. (abs, pdf, OpenReview)
 We presented a poster for AdvancedHMC.jl at StanCon 2019 in Cambridge, UK. (pdf)
API CHANGES
 [v0.2.22] Three functions are renamed.
Preconditioner(metric::AbstractMetric)
>MassMatrixAdaptor(metric)
andNesterovDualAveraging(δ, integrator::AbstractIntegrator)
>StepSizeAdaptor(δ, integrator)
find_good_eps
>find_good_stepsize
 [v0.2.15]
n_adapts
is no longer needed to constructStanHMCAdaptor
; the old constructor is deprecated.  [v0.2.8] Two Hamiltonian trajectory sampling methods are renamed to avoid a name clash with Distributions.
Multinomial
>MultinomialTS
Slice
>SliceTS
 [v0.2.0] The gradient function passed to
Hamiltonian
is supposed to return a valuegradient tuple now.
A minimal example  sampling from a multivariate Gaussian using NUTS
In this section we demonstrate a minimal example of sampling from a multivariate Gaussian (10 dimensional) using the no Uturn sampler (NUTS). Below we describe the major components of the Hamiltonian system which are essential to sample using this approach:
 Metric: In many sampling problems the sample space is usually associated with a metric, that allows us to measure the distance between any two points, and other similar quantities. In the example in this section, we use a special metric called the Euclidean Metric, represented with a
D × D
matrix from which we can compute distances.
Further details about the Metric component
The Euclidean metric is also known as the mass matrix in the physical perspective. For available metrics refer Hamiltonian mass matrix. Leapfrog integration: Leapfrog integration is a secondorder numerical method for integrating differential equations (In this case they are, equations of motion for the relative position of one particle with respect to the other). The order of this integration signifies its rate of convergence. Any alogrithm with a finite time step size will have numerical errors and the order is related to this error. For a secondorder algorithm, this error scales as the second power of the time step, hence, the name secondorder. Highorder intergrators are usually complex to code and have a limited region of convergence, hence they do not allow arbitrarily large time steps. A secondorder integrator is suitable for our purpose, hence we opt for the leapfrog integrator. It is called
leapfrog
due to the ways this algorithm is written, where the positions and velocities of particlesleap over
each other.
About the leapfrog integration scheme
Suppose Proposal for trajectories (static or dynamic): Different types of proposals can be used, which maybe static or dynamic. At each iteration of any variant of the HMC algorithm there are two main steps  the first step changes the momentum and the second step may change both the position and the momentum of a particle.
More about the proposals
In the classical HMC approach, during the first step, new values for the momentum variables are randomly drawn from their Gaussian distribution, independently of the current values of the position variables. Whereas, during the second step, a Metropolis update is performed, using Hamiltonian dynamics to provide a new state. For available proposals refer Proposal.using AdvancedHMC, ForwardDiff
using LogDensityProblems
using LinearAlgebra
# Define the target distribution using the `LogDensityProblem` interface
struct LogTargetDensity
dim::Int
end
LogDensityProblems.logdensity(p::LogTargetDensity, θ) = sum(abs2, θ) / 2 # standard multivariate normal
LogDensityProblems.dimension(p::LogTargetDensity) = p.dim
LogDensityProblems.capabilities(::Type{LogTargetDensity}) = LogDensityProblems.LogDensityOrder{0}()
# Choose parameter dimensionality and initial parameter value
D = 10; initial_θ = rand(D)
ℓπ = LogTargetDensity(D)
# Set the number of samples to draw and warmup iterations
n_samples, n_adapts = 2_000, 1_000
# Define a Hamiltonian system
metric = DiagEuclideanMetric(D)
hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff)
# Define a leapfrog solver, with initial step size chosen heuristically
initial_ϵ = find_good_stepsize(hamiltonian, initial_θ)
integrator = Leapfrog(initial_ϵ)
# Define an HMC sampler, with the following components
#  multinomial sampling scheme,
#  generalised NoUTurn criteria, and
#  windowed adaption for stepsize and diagonal mass matrix
proposal = NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator)
adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator))
# Run the sampler to draw samples from the specified Gaussian, where
#  `samples` will store the samples
#  `stats` will store diagnostic statistics for each sample
samples, stats = sample(hamiltonian, proposal, initial_θ, n_samples, adaptor, n_adapts; progress=true)
Parallel sampling
AdvancedHMC enables parallel sampling (either distributed or multithread) via Julia's parallel computing functions. It also supports vectorized sampling for static HMC and has been discussed in more detail in the documentation here.
The below example utilizes the @threads
macro to sample 4 chains across 4 threads.
# Ensure that julia was launched with appropriate number of threads
println(Threads.nthreads())
# Number of chains to sample
nchains = 4
# Cache to store the chains
chains = Vector{Any}(undef, nchains)
# The `samples` from each parallel chain is stored in the `chains` vector
# Adjust the `verbose` flag as per need
Threads.@threads for i in 1:nchains
samples, stats = sample(hamiltonian, proposal, initial_θ, n_samples, adaptor, n_adapts; verbose=false)
chains[i] = samples
end
GPU Sampling with CUDA
There is experimental support for running static HMC on the GPU using CUDA.
To do so the user needs to have CUDA.jl installed, ensure the logdensity of the Hamiltonian
can be executed on the GPU and that the initial points are a CuArray
.
A small working example can be found at test/cuda.jl
.
API and supported HMC algorithms
An important design goal of AdvancedHMC.jl is modularity; we would like to support algorithmic research on HMC. This modularity means that different HMC variants can be easily constructed by composing various components, such as preconditioning metric (i.e. mass matrix), leapfrog integrators, trajectories (static or dynamic), and adaption schemes etc. The minimal example above can be modified to suit particular inference problems by picking components from the list below.
metric
)
Hamiltonian mass matrix ( Unit metric:
UnitEuclideanMetric(dim)
 Diagonal metric:
DiagEuclideanMetric(dim)
 Dense metric:
DenseEuclideanMetric(dim)
where dim
is the dimensionality of the sampling space.
integrator
)
Integrator ( Ordinary leapfrog integrator:
Leapfrog(ϵ)
 Jittered leapfrog integrator with jitter rate
n
:JitteredLeapfrog(ϵ, n)
 Tempered leapfrog integrator with tempering rate
a
:TemperedLeapfrog(ϵ, a)
where ϵ
is the step size of leapfrog integration.
proposal
)
Proposal ( Static HMC with a fixed number of steps (
n_steps
) (Neal, R. M. (2011)):StaticTrajectory(integrator, n_steps)
 HMC with a fixed total trajectory length (
trajectory_length
) (Neal, R. M. (2011)):HMCDA(integrator, trajectory_length)
 Original NUTS with slice sampling (Hoffman, M. D., & Gelman, A. (2014)):
NUTS{SliceTS,ClassicNoUTurn}(integrator)
 Generalised NUTS with slice sampling (Betancourt, M. (2017)):
NUTS{SliceTS,GeneralisedNoUTurn}(integrator)
 Original NUTS with multinomial sampling (Betancourt, M. (2017)):
NUTS{MultinomialTS,ClassicNoUTurn}(integrator)
 Generalised NUTS with multinomial sampling (Betancourt, M. (2017)):
NUTS{MultinomialTS,GeneralisedNoUTurn}(integrator)
adaptor
)
Adaptor ( Adapt the mass matrix
metric
of the Hamiltonian dynamics:mma = MassMatrixAdaptor(metric)
 This is lowered to
UnitMassMatrix
,WelfordVar
orWelfordCov
based on the type of the mass matrixmetric
 This is lowered to
 Adapt the step size of the leapfrog integrator
integrator
:ssa = StepSizeAdaptor(δ, integrator)
 It uses Nesterov's dual averaging with
δ
as the target acceptance rate.
 It uses Nesterov's dual averaging with
 Combine the two above naively:
NaiveHMCAdaptor(mma, ssa)
 Combine the first two using Stan's windowed adaptation:
StanHMCAdaptor(mma, ssa)
Gradients
AdvancedHMC
supports both ADbased (Zygote
, Tracker
and ForwardDiff
) and userspecified gradients. In order to use userspecified gradients, please replace ForwardDiff
with ℓπ_grad
in the Hamiltonian
constructor, where the gradient function ℓπ_grad
should return a tuple containing both the logposterior and its gradient.
All the combinations are tested in this file except from using tempered leapfrog integrator together with adaptation, which we found unstable empirically.
sample
function signature in detail
The function sample(
rng::Union{AbstractRNG, AbstractVector{<:AbstractRNG}},
h::Hamiltonian,
κ::HMCKernel,
θ::AbstractVector{<:AbstractFloat},
n_samples::Int,
adaptor::AbstractAdaptor=NoAdaptation(),
n_adapts::Int=min(div(n_samples, 10), 1_000);
drop_warmup=false,
verbose::Bool=true,
progress::Bool=false,
)
Draw n_samples
samples using the proposal κ
under the Hamiltonian system h
 The randomness is controlled by
rng
. If
rng
is not provided,GLOBAL_RNG
will be used.
 If
 The initial point is given by
θ
.  The adaptor is set by
adaptor
, for which the default is no adaptation. It will perform
n_adapts
steps of adaptation, for which the default is1_000
or 10% ofn_samples
, whichever is lower.
 It will perform

drop_warmup
specifies whether to drop samples. 
verbose
controls the verbosity. 
progress
controls whether to show the progress meter or not.
Note that the function signature of the sample
function exported by AdvancedHMC.jl
differs from the sample
function used by Turing.jl
. We refer to the documentation of Turing.jl
for more details on the latter.
Citing AdvancedHMC.jl
If you use AdvancedHMC.jl for your own research, please consider citing the following publication:
Kai Xu, Hong Ge, Will Tebbutt, Mohamed Tarek, Martin Trapp, Zoubin Ghahramani: "AdvancedHMC.jl: A robust, modular and efficient implementation of advanced HMC algorithms.", Symposium on Advances in Approximate Bayesian Inference, 2020. (abs, pdf)
with the following BibTeX entry:
@inproceedings{xu2020advancedhmc,
title={AdvancedHMC. jl: A robust, modular and efficient implementation of advanced HMC algorithms},
author={Xu, Kai and Ge, Hong and Tebbutt, Will and Tarek, Mohamed and Trapp, Martin and Ghahramani, Zoubin},
booktitle={Symposium on Advances in Approximate Bayesian Inference},
pages={110},
year={2020},
organization={PMLR}
}
If you using AdvancedHMC.jl directly through Turing.jl, please consider citing the following publication:
Hong Ge, Kai Xu, and Zoubin Ghahramani: "Turing: a language for flexible probabilistic inference.", International Conference on Artificial Intelligence and Statistics, 2018. (abs, pdf)
with the following BibTeX entry:
@inproceedings{ge2018turing,
title={Turing: A language for flexible probabilistic inference},
author={Ge, Hong and Xu, Kai and Ghahramani, Zoubin},
booktitle={International Conference on Artificial Intelligence and Statistics},
pages={16821690},
year={2018},
organization={PMLR}
}
References

Neal, R. M. (2011). MCMC using Hamiltonian dynamics. Handbook of Markov chain Monte Carlo, 2(11), 2. (arXiv)

Betancourt, M. (2017). A Conceptual Introduction to Hamiltonian Monte Carlo. arXiv preprint arXiv:1701.02434.

Girolami, M., & Calderhead, B. (2011). Riemann manifold Langevin and Hamiltonian Monte Carlo methods. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 73(2), 123214. (arXiv)

Betancourt, M. J., Byrne, S., & Girolami, M. (2014). Optimizing the integrator step size for Hamiltonian Monte Carlo. arXiv preprint arXiv:1411.6669.

Betancourt, M. (2016). Identifying the optimal integration time in Hamiltonian Monte Carlo. arXiv preprint arXiv:1601.00225.

Hoffman, M. D., & Gelman, A. (2014). The NoUTurn Sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research, 15(1), 15931623. (arXiv)