ethanluoyc/dhmc-jax

Discontinuous Hamiltonian Monte Carlo in JAX

Discontinuous Hamiltonian Monte Carlo

Discontinuous Hamiltonian Monte Carlo (DHMC) (https://arxiv.org/abs/1705.08510) in JAX.

Description

The repo includes implementation of a few MCMC algorithms in JAX. The design follows largely from FunMCMC.

notebooks/HMC.ipynb includes a Jupyter notebook that demonstrates (D)HMC on a few problems.

Installation

In order to set up the necessary environment:

1. create an environment dhmc with the help of [conda],

   conda env create -f environment.yaml
   

2. activate the new environment with

   conda activate dhmc
   

TODOs

1. Allow for custom mass in the kinetic energy.
2. Implement operator splitting version of the integrator.
3. Allow for more tracing in the intermediate states.