Hamiltonian Monte Carlo

Implementations of various Hamiltonian dynamics based Markov chain Monte Carlo (MCMC) samplers in Python. A modular design is used to as far as possible allowing mixing and matching elements of different proposed extensions to the original Hybrid Monte Carlo algorithm proposed in Duane et al. (1987).

Implemented methods

Samplers

Static integration time with Metropolis sampling (Duane et al., 1987)
Random integration time with Metropolis sampling (Mackenzie, 1989)
Correlated momentum updates in Metropolis samplers (Horowitz, 1991)
Dynamic integration time with multinomial sampling (Betancourt, 2017)

Hamiltonian systems

Euclidean-metric systems - isotropic, diagonal and dense metrics
Riemannian-metric systems (Girolami and Calderhead, 2011) inluding the log density Hessian based SoftAbs metric (Betancourt, 2013)
Euclidean-metric systems subject to holonomic constraints (Hartmann and Schütte, 2005; Brubaker, Salzmann and Urtasun, 2012; Lelièvre, Rousset and Stoltz, 2018) and for inference in differentiable generative models when conditioning on observed outputs (Graham and Storkey, 2017a)

Numerical integrators

Explicit leapfrog for separable Hamiltonian systems
Implicit leapfrog for non-separable Hamiltonian systems
Geodesic leapfrog for constrained Hamiltonian systems
'Split' leapfrog for Hamiltonian systems with an analytically tractable component for which the exact flow can be solved (Shahbaba et al., 2014)

Installation

To install and use the package the minimal requirements are a Python 3.6+ environment with NumPy (tested with v1.15.0) and SciPy (tested with v1.1.0) installed.

From a local clone of the repository run python setup.py install to install the package in the current Python environment.

Optional dependencies

Autograd: if available will be used to automatically compute the required derivatives of the model functions (providing they are specified using functions from the autograd.numpy and autograd.scipy interfaces).
tqdm: if available a simple progress bar will be shown during sampling.
Arviz: if available outputs of a sampling run can be returned in an arviz.InferenceData container object, allowing straightforward use of the extensive Arviz visualisation and diagnostic functionality.
multiprocess and dill: if available multiprocess.Pool will be used in preference to the in-built mutiprocessing.Pool for parallelisation as multiprocess supports serialisation ( via dill) of a much wider range of types, including of Autograd generated functions.
RandomGen: if available the Xorshift1024 random number generator will be used when running multiple chains in parallel, with the jump method of the object used to reproducibly generate independent substreams.

Example usage

A simple complete example of using the package to sample from a multivariate Gaussian distribution with randomly generated parameters is given below. Here an isotropic Euclidean metric Hamiltonian system is used (corresponding to a isotropic covariance Gaussian marginal distribution on the momenta) with the dynamic integration time HMC implementation described in Betancourt (2017), which is a extension of the NUTS algorithm (Hoffman and Gelman, 2014).

import hmc
import autograd.numpy as np

# Generate random precision and mean parameters for a Gaussian
n_dim = 50
rng = np.random.RandomState(seed=1234)
rnd_eigvec, _ = np.linalg.qr(rng.normal(size=(n_dim, n_dim)))
rnd_eigval = np.exp(rng.normal(size=n_dim) * 2)
prec = (rnd_eigvec / rnd_eigval) @ rnd_eigvec.T
mean = rng.normal(size=n_dim)

# Deine potential energy (negative log density) for the Gaussian target
# distribution (gradient will be automatically calculated using autograd)
def pot_energy(pos):
    pos_minus_mean = pos - mean
    return 0.5 * pos_minus_mean @ prec @ pos_minus_mean

# Specify Hamiltonian system with isotropic Gaussian kinetic energy
system = hmc.systems.EuclideanMetricSystem(pot_energy)

# Hamiltonian is separable therefore use explicit leapfrog integrator
integrator = hmc.integrators.LeapfrogIntegrator(system, step_size=0.15)

# Use dynamic integration-time HMC implementation with multinomial 
# sampling from trajectories
sampler = hmc.samplers.DynamicMultinomialHMC(system, integrator, rng)

# Sample an initial position from zero-mean isotropic Gaussian
init_pos = rng.normal(size=n_dim)

# Sample a Markov chain with 1000 transitions
chains, chain_stats = sampler.sample_chain(1000, init_pos)

# Print RMSE in mean estimate
mean_rmse = np.mean((chains['pos'].mean(0) - mean)**2)**0.5
print(f'Mean estimate RMSE: {mean_rmse}')

# Print average acceptance probability
mean_accept_prob = chain_stats['accept_prob'].mean()
print(f'Mean accept prob: {mean_accept_prob:0.2f}')

References