Markov Chain Monte Carlo

General MCMC utilities

import plot
LogProb : Type = Float
def runChain {a} (initialize: Key -> a) (step: Key -> a -> a) (numSamples: Nat) (k:Key) : Fin numSamples => a = [k1, k2] = split_key k with_state (initialize k1) \s. for i:(Fin numSamples). x = step (ixkey k2 i) (get s) s := x x
def propose {a} (logDensity : a -> LogProb) (cur : a) (proposal : a) (k : Key) : a = accept = logDensity proposal > (logDensity cur + log (rand k)) select accept proposal cur
def meanAndCovariance {n d} (xs:n=>d=>Float) : (d=>Float & d=>d=>Float) = xsMean : d=>Float = (for i. sum for j. xs.j.i) / n_to_f (size n) xsCov : d=>d=>Float = (for i i'. sum for j. (xs.j.i' - xsMean.i') * (xs.j.i - xsMean.i ) ) / (n_to_f (size n) - 1) (xsMean, xsCov)

Metropolis-Hastings implementation

MHParams : Type = Float -- step size
def mhStep {d} [Ix d] (stepSize: MHParams) (logProb: (d=>Float) -> LogProb) (k:Key) (x:d=>Float) : d=>Float = [k1, k2] = split_key k proposal = x + stepSize .* randn_vec k1 propose logProb x proposal k2

HMC implementation

HMCParams : Type = (Nat & Float) -- leapfrog steps, step size
def leapfrogIntegrate {a} [VSpace a] ((nsteps, dt): HMCParams) (logProb: a -> LogProb) ((x, p): (a & a)) : (a & a) = x = x + (0.5 * dt) .* p (x, p) = apply_n nsteps (x, p) \(xOld, pOld). pNew = pOld + dt .* grad logProb xOld xNew = xOld + dt .* pNew (xNew, pNew) p = p + (0.5 * dt) .* grad logProb x (x, p)
def hmcStep {d} [Ix d] (params: HMCParams) (logProb: (d=>Float) -> LogProb) (k:Key) (x:d=>Float) : d=>Float = hamiltonian = \(x, p). logProb x - 0.5 * vdot p p [k1, k2] = split_key k p = randn_vec k1 proposal = leapfrogIntegrate params logProb (x, p) fst $ propose hamiltonian (x, p) proposal k2

Test it out

Generate samples from a multivariate normal distribution N([1.5, 2.5], [[1., 0.], [0., 0.05]]).

def myLogProb (x:(Fin 2)=>Float) : LogProb = x' = x - [1.5, 2.5] neg $ 0.5 * inner x' [[1.,0.],[0.,20.]] x'
numSamples : Nat = if dex_test_mode () then 1000 else 10000
k0 = new_key 1
mhParams = 0.1
mhSamples = runChain randn_vec (mhStep mhParams myLogProb) numSamples k0
:p meanAndCovariance mhSamples
([2.082871, 2.504329], [[1.205363, 0.009561], [0.009561, 0.049768]])
:html show_plot $ y_plot $ slice (map head mhSamples) 0 (Fin 1000)
hmcParams = (10, 0.1)
hmcSamples = runChain randn_vec (hmcStep hmcParams myLogProb) numSamples k0
:p meanAndCovariance hmcSamples
([1.548751, 2.499404], [[1.021255, 0.003061], [0.003061, 0.051128]])
:html show_plot $ y_plot $ slice (map head hmcSamples) 0 (Fin 1000)