Skip to content

ethanluoyc/dhmc-jax

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

4 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

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.

About

Discontinuous Hamiltonian Monte Carlo in JAX

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published