replica exchange MCMC

0-replica-exchange-mcmc.R

library(rstan)
library(doParallel)

replica.exchange.mcmc <- function (inv_T, n_ex, stanmodel, data, par_list, init, iter, warmup) {
  n_rep <- length(inv_T)
  len <- iter - warmup
  n_param <- sum(unlist(lapply(par_list, prod))) + 2  # number of parameters included E and lp__
  ms_T1 <- matrix(0, len*n_ex, n_param)               # MCMC samples at inv_T=1
  idx_tbl <- matrix(0, n_ex, n_rep)                   # index table of (exchange time, replica)
  E_tbl   <- matrix(0, n_ex, n_rep)                   # E table along idx_tbl
  init_list <- rep(list(init), n_rep)

  idx4ex <- function (n_rep, e) if (e %% 2 == 0) 1:floor(n_rep/2) * 2 - 1 else 1:(floor(n_rep/2)-1) * 2

  for (e in seq_len(n_ex)) {
    fit_list <- foreach(r=seq_len(n_rep), .packages='rstan') %dopar% {
      data$Inv_T <- inv_T[r]
      sampling(
        stanmodel, data=data, pars=c(names(par_list), 'E'), init=list(init_list[[r]]),
        iter=iter, warmup=warmup, chains=1, seed=r, refresh=-1
      )
    }
    ms_T1[((e-1)*len+1):(e*len), ] <- extract(fit_list[[1]], permuted=FALSE, inc_warmup=FALSE)[,1,]

    # exchange replicas
    E <- sapply(1:n_rep, function(r) extract(fit_list[[r]], permuted=FALSE, pars='E')[len,1,])
    idx <- 1:n_rep
    for (r in idx4ex(n_rep, e)) {
      w <- exp((inv_T[r] - inv_T[r+1]) * (E[r] - E[r+1]))
      if (runif(1,0,1) < w) {
        idx[r]   <- r+1
        idx[r+1] <- r
      }
    }
    E_tbl[e,] <- E
    idx_tbl[e,] <- idx

    # update init_list
    init_list <- lapply(seq_len(n_rep), function(r) {
      ms <- extract(fit_list[[idx_tbl[e,r]]], permuted=FALSE, pars=names(par_list))[len,1,]
      init <- lapply(names(par_list), function(p) {
        pos <- grep(paste0('^', p, '(\\[.*\\])?$'), names(ms))
        if (identical(par_list[[p]], 1)) unname(ms[pos]) else array(ms[pos], dim=par_list[[p]])
      })
      names(init) <- names(par_list)
      init
    })
  }
  colnames(ms_T1) <- names(fit_list[[1]])
  return(list(ms_T1=ms_T1, idx_tbl=idx_tbl, E_tbl=E_tbl))
}


source('generate-data.R')
data <- list(N=N, X=X, Y=Y)
init <- list(b=24.17, s_y=0.4048)
stanmodel <- stan_model(file='model/model.stan')
# fit <- sampling(stanmodel, data=c(data, Inv_T=1), seed=10)

N_rep <- 10    # number of replicas
N_ex <- 100    # number of exchanges
Inv_T <- 0.5^seq(0, -log(0.02)/log(2), len=N_rep)

registerDoParallel(3)
res <- replica.exchange.mcmc(inv_T=Inv_T, n_ex=N_ex,
  stanmodel=stanmodel, data=data, par_list=list(b=1, s_y=1), init=init, iter=70, warmup=50)
source('print-result.R')
stopImplicitCluster()

generate-data.R

set.seed(123)

N <- 50
b <- 0.6
s_y <- 0.4
X <- seq(from=0.1, to=4*pi, length=N)
Y <- sin(b * X) + rnorm(N, 0, s_y)

model.stan

data {
  int<lower=1> N;
  vector[N] Y;
  vector[N] X;
  real<lower=0> Inv_T;
}

parameters {
  real<lower=0> b;
  real<lower=0> s_y;
}

transformed parameters {
  real E;
  {
    vector[N] mu;
    for (n in 1:N)
      mu[n] <- sin(b * X[n]);
    E <- 0;
    E <- E - normal_log(b, 0, 50);
    E <- E - student_t_log(s_y, 4, 0, 5);
    E <- E - normal_log(Y, mu, s_y);
  }
}

model {
  increment_log_prob(-Inv_T * E);
}

print-result.R

library(ggplot2)

ix <- expand.grid(1:N_rep, 1:N_ex)
d_ex <- as.data.frame(matrix(ncol=3, nrow=2*nrow(ix)))
colnames(d_ex) <- c('x', 'y', 'g')

for (i in 1:nrow(ix)) {
  r <- ix[i,1]
  e <- ix[i,2]
  d_ex[2*i-1, ] <- c(e, Inv_T[r], i)
  d_ex[2*i, ]   <- c(e+1, Inv_T[res$idx_tbl[e,r]], i)
}

p <- ggplot(data=d_ex, aes(x=x, y=y, group=g))
p <- p + theme(text=element_text(size=18)) + labs(x='Exchange Time', y='Inverse T')
p <- p + geom_line()
ggsave(p, file='output/exchange-invT.png', dpi=300, w=8, h=6)


d_E <- reshape2::melt(res$E_tbl)
colnames(d_E) <- c('Exchange', 'Replica', 'E')
d_E$Replica <- as.factor(d_E$Replica)

p <- ggplot(data=d_E, aes(x=Exchange, y=E, group=Replica, color=Replica))
p <- p + theme(text=element_text(size=18)) + labs(x='Exchange Time', y='E')
p <- p + geom_line()
ggsave(p, file='output/energy-whole.png', dpi=300, w=8, h=6)
p <- p + ylim(NA, 100)
ggsave(p, file='output/energy-zoom.png', dpi=300, w=8, h=6)


d_ms <- reshape2::melt(res$ms_T1)
colnames(d_ms) <- c('Step', 'parameter', 'value')
p <- ggplot(data=d_ms, aes(x=Step, y=value))
p <- p + theme(text=element_text(size=18))
p <- p + facet_wrap(~parameter, scales='free_y')
p <- p + geom_line()
ggsave(p, file='output/traceplot.png', dpi=300, w=8, h=6)