h-emcee Documentation¶

h-emcee is an implementation of ensemble Hamiltonian Monte Carlo (HMC) samplers with affine invariance, based on Y. Chen (2025) and implemented in JAX. We also include affine-invariant schemes for tuning hyperparameters of HMC samplers (e.g., step size, number of leapfrog steps).

The philosophy and syntax of this package are meant as a minimal replacement for emcee. This is a pure-Python implementation designed for easy statistical inference (no graphical models needed).

Basic Example¶

All you need is access to the unnormalized log probability!

import jax
import jax.numpy as jnp
import hemcee

def log_prob(x):
   return -0.5 * jnp.sum(x ** 2)

key = jax.random.PRNGKey(0)
keys = jax.random.split(key, 2)

num_walkers, dim = 100, 5
inital_states = jax.random.normal(keys[0], shape=(num_walkers, dim))

sampler = hemcee.HamiltonianEnsembleSampler(num_walkers, dim, log_prob)
sampler.run_mcmc(keys[1], inital_states, 10000)

For a more through example, see Quick Start.

Navigating the documentation¶

User Guide

Tutorials