post = function(theta){
  exp(-0.5*(prod(theta)^2+sum(theta^2)-8*sum(theta)))
}
gradient = function(theta){
  c(-theta[1]*theta[2]^2+theta[1]-4,
    -theta[2]*theta[1]^2+theta[2]-4)
}

leapfrog = function(theta,p,eps,M){
  p     = p     + (eps/2)*gradient(theta)
  theta = theta + eps*iM%*%p
  p     = p     +  (eps/2)*gradient(theta)
  return(list(theta=theta,p=p))
}

N = 200
th1 = seq(-3,7,length=N)
th2 = seq(-3,7,length=N)
f = matrix(0,N,N)
for (i in 1:N)
 for (j in 1:N)
   f[i,j] = post(c(th1[i],th2[j]))
   
contour(th1,th2,f,drawlabels=FALSE)



k  = 2

# SIR
M = 10000
thetas.sir = matrix(0,M,k)




# RW-MH
set.seed(32425)
sd.th     = 0.1
burnin = 1000
N      = 20000
niter  = burnin + N
thetas.rwmh = matrix(0,niter,k)
theta     = c(0,0)

for (iter in 1:niter){
  theta.new = rnorm(2,theta,sd.th)
  accept = min(1,post(theta.new)/post(theta))
  if (runif(1)<accept){
    theta = theta.new
  }
  thetas.rwmh[iter,] = theta
}


# HMC in action!
set.seed(123456)
M      = diag(1,2)
iM     = solve(M)
tcM    = t(chol(M))
L      = 100
eps    = 0.001
theta  = c(0,0)
burnin = 1000
N      = 20000
niter  = burnin + N
thetas.hmc = matrix(0,niter,k)
for (iter in 1:niter){
  p      = tcM%*%rnorm(k)
  theta1 = theta
  p1     = p
  for (i in 1:L){
    run = leapfrog(theta1,p1,eps,M)
    p1  = run$p
    theta1 = run$theta
  }
  term1 = post(theta1)/post(theta)
  term2 = exp(-0.5*t(p1)%*%iM%*%p1)/exp(-0.5*t(p)%*%iM%*%p)
  accept = min(1,term1*term2)
  if (runif(1)<accept){
    theta = theta1
    p     = - p1
  }
  thetas.hmc[iter,]  = theta
}

par(mfrow=c(2,6))
for (i in 1:2){
  ts.plot(thetas.rwmh[,i])
  acf(thetas.rwmh[,i])
  hist(thetas.rwmh[,i])
}
for (i in 1:2){
  ts.plot(thetas.hmc[,i])
  acf(thetas.hmc[,i])
  hist(thetas.hmc[,i])
}


par(mfrow=c(1,2))
plot(thetas.rwmh,xlim=range(th1),ylim=range(th2),pch=16,col=grey(0.8))
contour(th1,th2,f,drawlabels=FALSE,add=TRUE)
plot(thetas.hmc,xlim=range(th1),ylim=range(th2),pch=16,col=grey(0.8))
contour(th1,th2,f,drawlabels=FALSE,add=TRUE)