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

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.

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
   

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.