kazewong / flowMC

Normalizing-flow enhanced sampling package for probabilistic inference in Jax
https://flowmc.readthedocs.io/en/main/
MIT License
199 stars 23 forks source link

Acceptance ratio possibly calculated incorrectly for Hamiltonian Monte Carlo implementation #78

Closed matt-graham closed 1 year ago

matt-graham commented 1 year ago

Raising issue as part of JOSS review https://github.com/openjournals/joss-reviews/issues/5021

The logarithm of the acceptance ratio computed in the lines

https://github.com/kazewong/flowMC/blob/6585fc2b8fea4a9575459e62b7c3ad209f005f15/src/flowMC/sampler/HMC.py#L70-L73

appears to compute the acceptance ratio as the difference between the Hamiltonian H before the momentum resampling step in line 70 and the Hamiltonian for the proposed state pair at the final point in the simulated trajectory. The momentum resampling step is itself a Markov kernel which leaves the joint distribution on the positions and momenta invariant, in particular a Gibbs sampling like step which samples independently from the conditional distribution under the target of the momentum given the position (here equal to the marginal distribution as the momentum and position are independent). There is then a second Metropolis-Hastings based Markov kernel which simulates the Hamiltonian dynamics forward in time from the current position-momentum state pair to a new state pair, with the logarithm of the acceptance ratio then being the difference between the Hamiltonians as the start and end points of the simulated trajectory (providing a volume preserving, time-reversible integrator is used). I think therefore that the H value used to compute log_acc therefore needs to be recalculated between the current lines 70 and 71 that is adding a line

    H = self.potential(position) + self.kinetic(momentum, params) 

To avoid recomputing the potential energy function (which will remain the same before and after the momentum resampling) you could store pass through the value of the potential energy function rather than H as input argument and return value of hmc_kernel.

kazewong commented 1 year ago

@matt-graham We have updated the HMC sampler accordingly in #90.