ハイブリッドモンテカルロを実装してみた

恐ろしく久しぶりですが…。
PRMLの後ろの方にあるハイブリッドモンテカルロのところを読んでいて、実際どんな動きをするのかと思い実装してみました。
以下コードと結果。

# set any function (exp(-energy) should be relevant distribution)
energy <- function(z){
  z*(z-1)*(z-4)*(z-6)/12
}

# set derivative of function "energy" (make sure the function is vectorized)
dzenergy <- function(z){
  ((z-1)*(z-4)*(z-6) + z*(z-4)*(z-6) + z*(z-1)*(z-6) + z*(z-1)*(z-4))/12
}

#leapfrog integral
leapfrog <- function(z0, r0, dzenergy, epsilon, nstep){
  z <- z0
  r <- r0
  log <- c(z0,r0,0.5*(r%*%r)+energy(z))
  for(i in 1:nstep){
    r <- r - 0.5*epsilon*dzenergy(z)
    z <- z + epsilon*r
    r <- r - 0.5*epsilon*dzenergy(z)
    log <- rbind(log, c(z,r,0.5*(r%*%r)+energy(z)))
  }
  return(list(z,r,log))
}

#hybrid montecarlo sampling
hybridMontecarlo <- function(z0, energy, dzenergy, nloop, epsilon, nLFstep){
  z <- z0
  sample <- z
  log <- as.list(NULL)
  for(i in 1:nloop){
    # initialize
    r <- rnorm(length(z))
    H <- 0.5*(r%*%r) + energy(z)

    # leapfrog
    new <- leapfrog(z, r, dzenergy, epsilon, nLFstep)
    znew <- new[[1]]
    rnew <- new[[2]]
    log[[length(log) + 1]] <- new[[3]]

    # metropolis
    Hnew <- 0.5*(rnew%*%rnew) + energy(znew)
    dH <- Hnew - H
    if(dH < 0 || runif(1) < exp(-dH)){
      z <- znew
      sample <- rbind(sample, z)
    }
  }
  return(list(sample, log))
}

# only for 1 dimensional sampling
test <- function(z0, nloop, eps, nLFstep){
  s <- hybridMontecarlo(z0, energy, dzenergy, nloop, eps, nLFstep)

  X11()
  smin <- min(s[[1]])
  smax <- max(s[[1]])
  hist(s[[1]], xlim=c(smin,smax))
  par(new=T)
  plot(function(z)(exp(-energy(z))), xlim=c(smin, smax), col="red", ann=F, yaxt="n")
  
  X11()
  rmin <- min(s[[2]][[1]][,2])
  rmax <- max(s[[2]][[1]][,2])
  ptim <- seq(1,nloop,20)
  for(i in 1:length(ptim)){
    plot(s[[2]][[ptim[i]]][,1], s[[2]][[ptim[i]]][,2], xlim=c(-1,7), ylim=c(rmin,rmax), col=i, xlab="position", ylab="momentum")
    par(new=T)
  }

  return(s)
}

上の図はサンプル結果と真の分布、下の図はサンプル過程（リープフロッグ積分）での各粒子（サンプル点）の相空間中の動きです（ただし20回に1回だけ描画）。
ハイブリッドモンテカルロは、要するに

ギブス分布の空間座標についての周辺分布が欲しい分布となるようなポテンシャル関数を設計する。
- $V(x)=-\log p(x)$ となるポテンシャルを用意する（分配関数は定数なので無視できる（エネルギーの原点はどうでもいい））。
適当な運動量(普通はN(0,1)に従う)を持つ粒子群を生成し、上記のポテンシャルの下で物理系をシミュレートする（微分方程式を解く。普通は正準方程式をリープフロッグ積分で解く）。
適当なタイミングで系のスナップショット（アンサンブル）をとって、運動量について周辺化する（ヒストグラムをとる）と欲しい分布が得られる。

という考え方で（実際は積分誤差を補正するためのメトロポリス棄却なんかが入っています。コード参照。）、エルゴード性そのもの！という感じでハァハァしますね。
普通のMCMCでは提案分布からの乱数でステップを更新するのに対して、ハイブリッドモンテカルロではステップを求めるために微分方程式を解いて力学系をシミュレートする必要があるので（その分自己相関の低いステップが早く得られるのですが）、高次元ではなかなか厳しそうです。

メトロポリス棄却のところは、物理系のシミュレートは当然ながらエネルギーが保存するように行わなければならない（物理法則を破ってしまう）のに対し、リープフロッグ積分の数値誤差でどうしても保存しなくなってしまうため、そのような場合を棄却する効果を果たします。と言いつつ、エネルギーが減ったプロセスについてはアクセプトするようで、いまいちどういうものなのかよくわかっていません…。

テストコードのエネルギー関数はhttp://d.hatena.ne.jp/n_shuyo/20100609/hybrid_mcから引用させて頂きました。ありがとうございます。