# Migrate-n convergence and analysis script for use with R
# S Fourdrilis, adapted from the script of C R Young

# R is a free software environment for statistical computing and graphics, and can be downloaded
# from: http://cran.freestatistics.org/web/packages/coda/index.html

# You will need to make sure that the CODA package is installed in R before attempting to use
# this script.

# The outfile "bayesallfile" generated by Migrate-n should be placed in the working directory
# (root by default) of R. On Windows 7 for instance, this directory is located in:
# Local Disk (C:)/Program Files/R/R-3.0.2/bin/x64

# Ensure there are no extra columns in the bayesallfile filled with 0. If so, delete these
# columns.

# This script is meant to analyse migrate runs with 3 replicates. If you did more replicates, 
# you will need to add the following commands as many times as there are replicates, and to
# replace "X" by the number of the replicate:
# In the "Gelman and Rubin's diagnostic" section:
#         MCMCdata_chainX <- subset(MCMCdata, Replicate=="X")
#         size_chainX = dim(MCMCdata_chainX)[1]
#         num.parms = dim(MCMCdata_chainX)[2]
#         All.parms_chainX <- MCMCdata_chainX[, 4:num.parms]
#         MCMCdata.All_chainX = mcmc(data=All.parms_chainX, start = 1, end = size_chainX, thin = 1)
# In the "Heidelberger-Welch diagnostic" section:
#         heidel.diag(MCMCdata.All_chainX, eps=0.1, pvalue=0.05)
# Also, edit the following command to add as many replicates as you have:
# In the "Gelman and Rubin's diagnostic" section:
# MCMCdata.All_chains <- mcmc.list(MCMCdata.All_chain1, MCMCdata.All_chain2, MCMCdata.All_chain3)

# This complete file can be copied and pasted into the R console.

# Output of this script include:
# ------------------------------
# Effective sample size estimates
# Geweke's convergence diagnostic
# Raftery and Lewis's diagnostic
# Gelman-Rubin diagnostic (requires "MCMCdata_chain1.out" and "MCMCdata_chain2.out". Modifications to script are necessary if more chains are tested.
# Heidelberger-Welch diagnostic
# Parameter traces over the length of the chain

# After you ran R, save the outputs:
# 1) save the R console: File/Save to File.../
# 2) save the 2 pdf files containing the traces and density plots.

# --------------------------------------------------------------------------------------------
# Load the coda library
# --------------------------------------------------------------------------------------------

library(coda)

# --------------------------------------------------------------------------------------------
# Clear R's memory
# --------------------------------------------------------------------------------------------

rm(list = ls(all = TRUE))

# --------------------------------------------------------------------------------------------
# Read the contents of the datafile "bayesallfile" that should be in R's working directory
# --------------------------------------------------------------------------------------------

MCMCdata <- read.table("bayesallfile", header=TRUE)
# The bayesallfile is stored as the object MCMCdata.

# --------------------------------------------------------------------------------------------
# Turn MCMCdata into MCMC object
# --------------------------------------------------------------------------------------------

size = dim(MCMCdata)[1]
# to calculate the size of the 1st dimension = total number of rows in the bayesallfile
num.parms = dim(MCMCdata)[2]
# to calculate the size of the 2nd dimension = total number of columns in the bayesallfile
All.parms <- MCMCdata[, 4:num.parms]
# bayesallfile from column 4 to 24 = 3 first columns discarded, because we don't use the 3 first
# columns = steps, locus, replicate) to represent their traces.
MCMCdata.All = mcmc(data=All.parms, start = 1, end = size, thin = 1)
# conversion of the All.parms matrix into mcmc object

# --------------------------------------------------------------------------------------------
# Calculates effective sample sizes  (all parameters)
# --------------------------------------------------------------------------------------------

effectiveSize(MCMCdata.All)

# --------------------------------------------------------------------------------------------
# Convergence Diagnostic (Geweke's diagnostic)
# --------------------------------------------------------------------------------------------
# z-score for difference in means of first 10% of chain and last 50% (stationarity)
# -1.96 < z-score < 1.96 is reasonable
#
# Geweke (1992) proposed a convergence diagnostic for Markov chains based on a test for
# equality of the means of the first and last part of a Markov chain (by default the first
# 10% and the last 50%). If the samples are drawn from the stationary distribution of the
# chain, the two means are equal and Geweke's statistic has an asymptotically standard
# normal distribution.
#
# The test statistic is a standard Z-score: the difference between the two sample means
# divided by its estimated standard error. The standard error is estimated from the spectral
# density at zero and so takes into account any autocorrelation.
#
# The Z-score is calculated under the assumption that the two parts of the chain are
# asymptotically independent, which requires that the sum of frac1 and frac2 be strictly
# less than 1.

geweke.diag(MCMCdata.All, frac1=0.1, frac2=0.5)

# --------------------------------------------------------------------------------------------
# Convergence Diagnostic (Raftery and Lewis's diagnostic)
# --------------------------------------------------------------------------------------------
# Estimate for how long to run the chain for specified accuracy
#
# raftery.diag is a run length control diagnostic based on a criterion of accuracy of
# estimation of the quantile q. It is intended for use on a short pilot run of a Markov
# chain. The number of iterations required to estimate the quantile q to within an accuracy
# of +/- r with probability p is calculated. Separate calculations are performed for each
# variable within each chain.
#
# If the number of iterations in data is too small, an error message is printed indicating
# the minimum length of pilot run. The minimum length is the required sample size for a chain
# with no correlation between consecutive samples. Positive autocorrelation will increase the
# required sample size above this minimum value. An estimate I (the `dependence factor') of
# the extent to which autocorrelation inflates the required sample size is also provided.
# Values of I larger than 5 indicate strong autocorrelation which may be due to a poor choice
# of starting value, high posterior correlations or `stickiness' of the MCMC algorithm.
#
# The number of `burn in' iterations to be discarded at the beginning of the chain is also calculated.

raftery.diag(MCMCdata.All, q=0.05, r=0.01,  s=0.95, converge.eps=0.001)
raftery.diag(MCMCdata.All, q=0.50, r=0.025, s=0.95, converge.eps=0.001)
raftery.diag(MCMCdata.All, q=0.95, r=0.01,  s=0.95, converge.eps=0.001)

# --------------------------------------------------------------------------------------------
# Convergence Diagnostic (Gelman and Rubin's diagnostic)
# --------------------------------------------------------------------------------------------
# Gelman and Rubin (1992) propose a general approach to monitoring convergence of MCMC output in
# which two or more parallel chains are run with starting values that are overdispersed relative
# to the posterior distribution. Convergence is diagnosed when the chains have `forgotten' their
# initial values, and the output from all chains is indistinguishable. The gelman.diag diagnostic
# is applied to a single variable from the chain. It is based a comparison of within-chain and
# between-chain variances, and is similar to a classical analysis of variance.
#
# This script assumes that you have 1 chain per replicate named: "MCMCdata_chain1",
# "MCMCdata_chain2" and "MCMCdata_chain3" of the THE SAME SIZE (in number of iterations)!
# You have to subset replicates from MCMCdata, then convert them to mcmc object (not possible from
# MCMCdata.All because the 3 1st columns are missing).
#
# The test returns a list of psrf values labelled "Point est." and their upper confidence limits
# labelled "Upper C.I."
# Values substantially above 1 indicate lack of convergence.

MCMCdata_chain1 <- subset(MCMCdata, Replicate=="1")    # the bayesallfile MUST have a header!
MCMCdata_chain2 <- subset(MCMCdata, Replicate=="2")
MCMCdata_chain3 <- subset(MCMCdata, Replicate=="3")
size_chain1 = dim(MCMCdata_chain1)[1]                  # to calculate the size of the 1st dimension = total number of rows
size_chain2 = dim(MCMCdata_chain2)[1]
size_chain3 = dim(MCMCdata_chain3)[1]
num.parms = dim(MCMCdata_chain1)[2]                    # to calculate the size of the 2nd dimension = total number of columns
num.parms = dim(MCMCdata_chain2)[2]
num.parms = dim(MCMCdata_chain3)[2]
All.parms_chain1 <- MCMCdata_chain1[, 4:num.parms]     # to trim the 3 first columns
All.parms_chain2 <- MCMCdata_chain2[, 4:num.parms]
All.parms_chain3 <- MCMCdata_chain3[, 4:num.parms]
MCMCdata.All_chain1 = mcmc(data=All.parms_chain1, start = 1, end = size_chain1, thin = 1)   # conversion into mcmc object
MCMCdata.All_chain2 = mcmc(data=All.parms_chain2, start = 1, end = size_chain2, thin = 1)
MCMCdata.All_chain3 = mcmc(data=All.parms_chain3, start = 1, end = size_chain3, thin = 1)

MCMCdata.All_chains <- mcmc.list(MCMCdata.All_chain1, MCMCdata.All_chain2, MCMCdata.All_chain3)

gelman.diag(MCMCdata.All_chains, confidence = 0.95, transform=FALSE, autoburnin=FALSE, multivariate=FALSE)
# Error message: Error in chol.default(W): the leading minor of order 4 is not positive definite
# The error is solved when multivariate is set to FALSE instead of TRUE.

# --------------------------------------------------------------------------------------------
# Convergence Diagnostic (Heidelberger-Welch diagnostic)
# --------------------------------------------------------------------------------------------
# The convergence test uses the Cramer-von-Mises statistic to test the null hypothesis that the
# sampled values come from a stationary distribution. The test is successively applied, firstly
# to the whole chain, then after discarding the first 10%, 20%, ... of the chain until either the
# null hypothesis is accepted, or 50% of the chain has been discarded. The latter outcome constitutes
# `failure' of the stationarity test and indicates that a longer MCMC run is needed. If the stationarity
# test is passed, the number of iterations to keep and the number to discard are reported.
#
# The half-width test calculates a 95% confidence interval for the mean, using the portion of the chain
# which passed the stationarity test. Half the width of this interval is compared with the estimate of
# the mean. If the ratio between the half-width and the mean is lower than eps, the halfwidth test is
# passed. Otherwise the length of the sample is deemed not long enough to estimate the mean with sufficient
# accuracy.
#
# The test does not work when the bayesallfile contains replicates! Apply it to replicates!

heidel.diag(MCMCdata.All_chain1, eps=0.1, pvalue=0.05)
heidel.diag(MCMCdata.All_chain2, eps=0.1, pvalue=0.05)
heidel.diag(MCMCdata.All_chain3, eps=0.1, pvalue=0.05)

# --------------------------------------------------------------------------------------------
# Plot traces (all parameters) of the 3 chains together
# --------------------------------------------------------------------------------------------     
#
# The pdf("Traces.pdf") command, as well as the SavePlot() function, do not work because it can't
# overwrite pdf files while they are opened by another application, like a pdf viewer:
# Error in pdf(file = "Traces.pdf") : cannot open file 'Traces.pdf'
# Instead, to create plots, create temporary files:

pdf(file = tf <- tempfile(fileext = "Traces.pdf"))
plot(MCMCdata.All, trace = TRUE, density = FALSE, smooth = TRUE, auto.layout = FALSE, ask = FALSE)
dev.off()              # to close the file
shell.exec(tf)         # to open the pdf viewer

# --------------------------------------------------------------------------------------------
# Plot traces (all parameters) of the 3 single chains
# --------------------------------------------------------------------------------------------     
#
# The command line given in the sript of C.R. Young does not work:
# gelman.plot(MCMCdata.All_chains, bin.width = 10, max.bins = 50, confidence = 0.95, transform = FALSE, ask = TRUE)
# because it returns the error message:
# Error in chol.default(W): the leading minor of order 4 is not positive definite
# This is a matter of variables that are not enough variable (variance = 0) or highly correlated
# (correlation = 1). The error is solved by using the 4 commands below:

pdf(file = tf <- tempfile(fileext = "Gelman-Rubin_plot.pdf"))
plot(MCMCdata.All_chains, trace = TRUE, density = TRUE, ask = FALSE)    # to plot the 3 traces of the replicates on the same graph
dev.off()              # to close the file
shell.exec(tf)         # to open the pdf viewer