# 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-3.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( 8 ) # 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 = round( 1e+06 * runif( 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 ) ) )

}

# Multi-threaded run (assuming 8 chains)

iters <- 12500

set.seed( 1 )

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-3.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 )

#   user  system elapsed 
#   0.21    0.01   14.88

# 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-3.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 )

#   user  system elapsed 
#  80.08    0.48   80.66

80.66 / 14.88

# [1] 5.420699

# speed-up seem only to be by about a factor of 5.4,
# even with 8 threads in use

# this seems to be because, even though i didn't ask for it,
# each chain goes through an undocumented
# burn-in of 1000 iterations before monitoring:

str( jags.parsamples )

# List of 8
#  $ : mcmc [1:12500, 1:4] 403 405 404 404 404 ...
#   ..- attr(*, "dimnames")=List of 2
#   .. ..$ : NULL
#   .. ..$ : chr [1:4] "mu" "nu" "sigma" "y.new"
#   ..- attr(*, "mcpar")= num [1:3] 1001 13500 1
#
# and so on (8 of these )
#
#  $ : mcmc [1:12500, 1:4] 405 404 404 405 405 ...
#   ..- attr(*, "dimnames")=List of 2
#   .. ..$ : NULL
#   .. ..$ : chr [1:4] "mu" "nu" "sigma" "y.new"
#   ..- attr(*, "mcpar")= num [1:3] 1001 13500 1
#  - attr(*, "class")= chr "mcmc.list"

# so the parallel routine is creating 8 * 13500 = 108000
# samples in 14.88 seconds, which is ( 108000 / 14.88 ) =
# 7258 samples per second, and the single-thread routine
# is creating 101000 samples in 80.66 seconds, which is
# ( 101000 / 80.66 ) = 1252 samples per second; and
# ( 7258 / 1252 ) = 5.8 -- the real speed-up is about
# 5.8 with 8 threads

plot( jags.parsamples ) 

library( lattice )

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

summary( jags.parsamples )

# Iterations = 1001:13500
# Thinning interval = 1 
# Number of chains = 8 
# Sample size per chain = 12500 

# 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.4746 0.001501       0.001906
# nu      3.654 1.1945 0.003777       0.007818
# sigma   3.866 0.4403 0.001392       0.002458
# y.new 404.310 6.7623 0.021384       0.021425

# 2. Quantiles for each variable:

#          2.5%     25%     50%     75%   97.5%
# mu    403.367 403.974 404.296 404.613 405.231
# nu      2.142   2.806   3.399   4.207   6.712
# sigma   3.079   3.558   3.840   4.147   4.807
# y.new 392.642 401.377 404.303 407.213 415.975

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 )