# Written by Chris Severs, edited by David Draper

# Some helpers and example usage of rjags in parallel mode in R. 

# The output from the foreach functions below give an mcmc object that 
# can be manipulated with the coda package, which is automatically loaded 
# with rjags.

# The base code is from: http://andrewgelman.com/2011/07/parallel-jags-rngs/
# with some updates and a nice combine function. 

# This builds on the NB10 rjags example in the file called 
# "How to MCMC sample with R-JAGS" on the course web page:

# First, change directory to the location where you previously stored the
# file "nb10-model-1.txt"

# Second, if the following 3 packages have not yet been installed, 
# install them using the instructions in the file "How to MCMC sample with R-JAGS"

#  doParallel
#  rjags
#  random

library( doParallel )

library( rjags )

library( random )

# For OS X/Linux you can optionally use this instead of the version below

# registerDoParallel( ) 

# Change to registerDoParallel( n ) to use n processes

# For Windows, you need to use the following:

cl <- makeCluster( 4 ) # change to makeCluster( n ) to use n processes

registerDoParallel( cl )

# Function to generate the initial values and random seeds for each model. 

nb10.inits.1 <- function( ) {

  return( list( mu = 404.59, tau = 0.05, nu = 5.0, y.new = 400.0, 
    .RNG.name = "lecuyer::RngStream", 
    .RNG.seed = randomNumbers( n = 1, min = 1, max = 1e+06, col = 1 ) ) )
}

# Data to use

nb10.data.1 <- list( 
  y = c( 409., 400., 406., 399., 402., 406., 401., 403., 401., 403.,
         398., 403., 407., 402., 401., 399., 400., 401., 405., 402., 
         408., 399., 399., 402., 399., 397., 407., 401., 399., 401., 
         403., 400., 410., 401., 407., 423., 406., 406., 402., 405., 
         405., 409., 399., 402., 407., 406., 413., 409., 404., 402., 
         404., 406., 407., 405., 411., 410., 410., 410., 401., 402., 
         404., 405., 392., 407., 406., 404., 403., 408., 404., 407., 
         412., 406., 409., 400., 408., 404., 401., 404., 408., 406., 
         408., 406., 401., 412., 393., 437., 418., 415., 404., 401., 
         401., 407., 412., 375., 409., 406., 398., 406., 403., 404.), 
  n = 100 )

# Helper function to combine multiple mcmc lists into a single one. 

mcmc.combine <- function( ... ) {

  return( as.mcmc.list( sapply( list( ... ), mcmc ) ) )

}

# Usage example with some timing to see how much speedup we get with 
# parallel runs. On Chris Severs's machine the time difference is about 3x;
# on my machine it's about 3.6 

# Multi-threaded run (assuming 4 chains)

iters <- 25000

time.1 <- system.time(

  jags.parsamples <- foreach( i = 1:getDoParWorkers( ), .inorder = FALSE, 
    .packages = c( 'rjags', 'random' ), .combine = "mcmc.combine", 
    .multicombine = TRUE ) 

  %dopar% {

    load.module( "lecuyer" )

    model.jags <- jags.model( data = nb10.data.1, file = "nb10-model-1.txt", 
      inits = nb10.inits.1 )

    result <- coda.samples( model.jags, variable.names = c( "mu", "sigma", 
      "nu", "y.new" ), n.iter = iters )

    return( result )

  }

)

print( time.1 )
 
# Single-threaded run

iters <- 100000

time.2 <- system.time(

  jags.singlesample <- foreach( i = 1:1, .combine = "mcmc.combine", 
    .multicombine = TRUE) 

  %do% {

    load.module("lecuyer") 

    model.jags <- jags.model( data = nb10.data.1, file = "nb10-model-1.txt", 
      inits = nb10.inits.1 )

    result <- coda.samples( model.jags, variable.names = c( "mu", "sigma",
      "nu", "y.new" ), n.iter = iters )

    return( result )

  }

)

print( time.2 )

plot( jags.parsamples ) 

xyplot( jags.parsamples[ , c( "mu", "nu" ) ] )

summary( jags.parsamples )

# Iterations = 1001:26000
# Thinning interval = 1 
# Number of chains = 4 
# Sample size per chain = 25000 

# 1. Empirical mean and standard deviation for each variable,
#    plus standard error of the mean:

#          Mean     SD Naive SE Time-series SE
# mu    404.295 0.4719 0.001492       0.001920
# nu      3.637 1.1875 0.003755       0.007821
# sigma   3.863 0.4375 0.001384       0.002462
# y.new 404.295 6.8578 0.021686       0.021686

# 2. Quantiles for each variable:

#          2.5%     25%     50%     75%   97.5%
# mu    403.374 403.978 404.296 404.611 405.224
# nu      2.139   2.804   3.387   4.177   6.638
# sigma   3.078   3.558   3.840   4.141   4.793
# y.new 392.631 401.356 404.282 407.200 416.000

summary( jags.singlesample )

# Iterations = 1001:101000
# Thinning interval = 1 
# Number of chains = 1 
# Sample size per chain = 1e+05 

# 1. Empirical mean and standard deviation for each variable,
#    plus standard error of the mean:

#          Mean     SD Naive SE Time-series SE
# mu    404.299 0.4731 0.001496       0.001917
# nu      3.631 1.1653 0.003685       0.007388
# sigma   3.864 0.4348 0.001375       0.002436
# y.new 404.315 6.6972 0.021178       0.021178

# 2. Quantiles for each variable:

#          2.5%     25%     50%     75%   97.5%
# mu    403.379 403.980 404.295 404.615 405.236
# nu      2.142   2.807   3.387   4.165   6.567
# sigma   3.081   3.560   3.841   4.142   4.782
# y.new 392.684 401.398 404.307 407.196 416.070

autocorr.plot( jags.parsamples )