# IM convergence and analysis script for use with R
# C R Young

# R is a free software environment for statistical computing and graphics, and can be downloaded from:
# http://www.r-project.org/

# You will need to make sure that the CODA and 
# KernSmooth packages are installed in R before 
# attempting to use this script!!!!!!

# These functions assume that the file "MCMCdata.out"
# is present in the working directory (root by default) 
# of R. This file is the output file generated by 
# MCMCthin. 

# Instructions for installing packages in R can be
# found in the html help documentation of R. 

# This file contains a set of R functions that I find useful in diagnosing convergence of IM runs
# and in understanding the fitted model parameters, both by visualizing them and by examining 
# numerical output. This complete file can be copied and pasted into the R console. Several plots
# (*.pdf) are saved to the root directory of R. Most of the plotting functions are fairly automatic,
# and the parameters of the plotting functions should be reasonable for most data. However, you may 
# need to adjust the xlim and ylim in some plots. This script is meant to provide a starting point
# for analysis of IM runs, but you will probably need to tailor it to your particular data.

# Descriptions of some of the methods used in this script are included in notes from a short MCMC
# workshop given by C R Young at MBARI in 2005. The R documentation is also a good resource. 

# Output of this script include:
# ------------------------------
# Parameter traces over the length of the chain
# Cumulative quantile plots (0.025%, 50%, and 97.5%)
# Marginal kernel density estimates (KDEs)
# Joint KDEs
# 2D contour plots of joint KDEs
# Summary statistics of marginal distributions (including cross correlations and autocorrelations)
# 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

# Passing MCMC diagnostics does NOT guarantee that a chain is stationary!!!!!!

# I suggest using MCMCthin to generate and output file
# on the order of 5000 lines. 

# Further information on these functions can be obtained from:
# http://www.maths.lth.se/help/R/.R/library/coda/html/00Index.html
# or the html documentation in R


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

library(coda)


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

rm(list = ls(all = TRUE))


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

MCMCdata <- read.table("MCMCdata.out", header=TRUE)


# --------------------------------------------------------------------------------------------
# Turn MCMCdata into MCMC object
# --------------------------------------------------------------------------------------------
size = dim(MCMCdata)[1]
MCMCdata = mcmc(data= MCMCdata, start = 1, end = size, thin = 1)


# --------------------------------------------------------------------------------------------
# Plots Cumulative Estimates of 0.025%, 50%, and 97.5% quantiles (all parameters)
# --------------------------------------------------------------------------------------------
postscript("ZZZ1.ps")
cumuplot(MCMCdata, probs=c(0.025,0.5,0.975), ylab="",lty=c(2,1), lwd=c(1,2), type="l", ask=FALSE, auto.layout=TRUE, col=1)
dev.off()


# --------------------------------------------------------------------------------------------
# Plot traces (all parameters)
# --------------------------------------------------------------------------------------------
postscript("ZZZ2.ps")
plot(MCMCdata, trace = TRUE, density = FALSE, smooth = TRUE, auto.layout = TRUE, ask = FALSE)
dev.off()

# --------------------------------------------------------------------------------------------
# Plot Kernel density estimates (all parameters)
# --------------------------------------------------------------------------------------------
postscript("ZZZ3.ps")
plot(MCMCdata, trace = FALSE, density = TRUE, smooth = TRUE, auto.layout = TRUE, ask = FALSE)
dev.off()

# --------------------------------------------------------------------------------------------
# Autocorrelation function (all parameters)
# --------------------------------------------------------------------------------------------
postscript("ZZZ4.ps")
autocorr.plot(MCMCdata, 50, auto.layout = TRUE, ask = FALSE)
dev.off()

# --------------------------------------------------------------------------------------------
# Partial Autocorrelation function (all parameters)
# --------------------------------------------------------------------------------------------
postscript("ZZZ5.ps")
pacf(MCMCdata, lag.max = 20, plot=TRUE, ask = FALSE)
dev.off()

# --------------------------------------------------------------------------------------------
# Cross correlation Plot
# --------------------------------------------------------------------------------------------
postscript("ZZZ6.ps")
crosscorr.plot (MCMCdata, col = topo.colors(10))
dev.off()

# --------------------------------------------------------------------------------------------
# Basic Summary Statistics  (all parameters)
# --------------------------------------------------------------------------------------------

summary(MCMCdata, quantiles = c(0.025, 0.25, 0.5, 0.75, 0.975))


# --------------------------------------------------------------------------------------------
# Numerical output for cross correlations
# --------------------------------------------------------------------------------------------

crosscorr(MCMCdata)

# --------------------------------------------------------------------------------------------
# Numerical output for autocorrelations  (all parameters)
# --------------------------------------------------------------------------------------------

autocorr.diag(MCMCdata, lags = c(0, 1, 5, 10, 50, 100, 500, 1000, 5000), relative=TRUE)


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

effectiveSize(MCMCdata)


# --------------------------------------------------------------------------------------------
# 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, frac1=0.1, frac2=0.5) 




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

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




# --------------------------------------------------------------------------------------------
# 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, q=0.05, r=0.01,  s=0.95, converge.eps=0.001)
raftery.diag(MCMCdata, q=0.50, r=0.025, s=0.95, converge.eps=0.001)
raftery.diag(MCMCdata, q=0.95, r=0.01,  s=0.95, converge.eps=0.001)


# --------------------------------------------------------------------------------------------
# --------------------------------------------------------------------------------------------
# 2D Kernel density estimates: Package "KernSmooth"
# --------------------------------------------------------------------------------------------
# --------------------------------------------------------------------------------------------
# dpik()
# Use direct plug-in methodology to select the bandwidth of a kernel density estimate.
# This method for selecting the bandwidth of a kernel density estimate was proposed by Sheather 
# and Jones (1991) and is described in Section 3.6 of Wand and Jones (1995).
# bkde2D()
# This is the binned approximation to the 2D kernel density estimate. Linear binning is used to 
# obtain the bin counts and the Fast Fourier Transform is used to perform the discrete convolutions. 
# For each x1,x2 pair the bivariate Gaussian kernel is centered on that location and the heights 
# of the kernel, scaled by the bandwidths, at each datapoint are summed. This sum, after a 
# normalization, is the corresponding fhat value in the output.

rm(list = ls(all = TRUE))
library(KernSmooth)
MCMCdata <- read.table("MCMCdata.out", header=TRUE)
bwq1 <- dpik(MCMCdata$q1)
bwq2 <- dpik(MCMCdata$q2)
bwm1 <- dpik(MCMCdata$m1)
bwm2 <- dpik(MCMCdata$m2)
bwqA <- dpik(MCMCdata$qA)
bwt  <- dpik(MCMCdata$t)

# --------------------------------------------------------------------------------------------
# 2D Kernel density estimates: q1 X q2
# --------------------------------------------------------------------------------------------
# horizontal view

x <- cbind(MCMCdata$q1, MCMCdata$q2)
est <- bkde2D(x, bandwidth=c(bwq1,bwq2))
par(ask = FALSE)
postscript("ZZZ8.ps")
persp(est$x1, est$x2, est$fhat, theta = 145, phi = 25, col = "lightblue1", xlab = "q1", ylab = "q2", zlab = "Pr(q1, q2| X)", main = "Joint Density of q1 and q2", axes = TRUE, nticks = 5, ticktype = "detailed", ltheta = -135, lphi = 145, shade = TRUE)
dev.off()

# Top view (including contour plot)
postscript("ZZZ9.ps")
x <- cbind(MCMCdata$q1, MCMCdata$q2)
est <- bkde2D(x, bandwidth=c(bwq1,bwq2))
contour(est$x1, est$x2, est$fhat, , xlim = c(0,50), ylim = c(0,50))
lines(c(0,100), c(0,100), type = "l", lty = "dashed")			# add line y=x to contour plot
dev.off()

postscript("ZZZ10.ps")
persp(est$x1, est$x2, est$fhat, theta = 325, phi = 70, d = 3, col = "lightblue1", xlab = "q1", ylab = "q2", zlab = "Pr(q1, q2| X)", main = "Joint Density of q1 and q2", axes = TRUE, nticks = 5, ticktype = "detailed", ltheta = -180, lphi = 145, shade = TRUE, box = TRUE)
dev.off()


# --------------------------------------------------------------------------------------------
# 2D Kernel density estimates: m1 X m2
# --------------------------------------------------------------------------------------------
# horizontal view

postscript("ZZZ11.ps")
x <- cbind(MCMCdata$m1, MCMCdata$m2)
est <- bkde2D(x, bandwidth=c(bwm1,bwm2))
persp(est$x1, est$x2, est$fhat, theta = 145, phi = 25, col = "lightblue1", xlab = "m1", ylab = "m2", zlab = "Pr(m1, m2| X)", main = "Joint Density of m1 and m2", axes = TRUE, nticks = 5, ticktype = "detailed", ltheta = -135, lphi = 145, shade = TRUE)
dev.off()

# Top view (including contour plot)

postscript("ZZZ12.ps")
x <- cbind(MCMCdata$m1, MCMCdata$m2)
est <- bkde2D(x, bandwidth=c(bwm1,bwm2))
contour(est$x1, est$x2, est$fhat, xlim = c(0,20), ylim = c(0,20))
lines(c(0,100), c(0,100), type = "l", lty = "dashed")  			# add line y=x to contour plot
persp(est$x1, est$x2, est$fhat, theta = 325, phi = 70, d = 3, col = "lightblue1", xlab = "m1", ylab = "m2", zlab = "Pr(m1, m2| X)", main = "Joint Density of m1 and m2", axes = TRUE, nticks = 5, ticktype = "detailed", ltheta = -180, lphi = 145, shade = TRUE, box = TRUE)
dev.off()



# --------------------------------------------------------------------------------------------
# 2D Kernel density estimates: qA X t
# --------------------------------------------------------------------------------------------
# horizontal view
postscript("ZZZ13.ps")
x <- cbind(MCMCdata$qA, MCMCdata$t)
est <- bkde2D(x, bandwidth=c(bwqA,bwt))
persp(est$x1, est$x2, est$fhat, theta = 145, phi = 25, col = "lightblue1", xlab = "qA", ylab = "t", zlab = "Pr(qA, t| X)", main = "Joint Density of qA and t", axes = TRUE, nticks = 5, ticktype = "detailed", ltheta = -135, lphi = 145, shade = TRUE)
dev.off()

# Top view (including contour plot)
postscript("ZZZ14.ps")
x <- cbind(MCMCdata$qA, MCMCdata$t)
est <- bkde2D(x, bandwidth=c(bwqA,bwt))
contour(est$x1, est$x2, est$fhat, xlim = c(0,50), ylim = c(0,50))
lines(c(0,100), c(0,100), type = "l", lty = "dashed")  			# add line y=x to contour plot
persp(est$x1, est$x2, est$fhat, theta = 325, phi = 70, d = 3, col = "lightblue1", xlab = "qA", ylab = "t", zlab = "Pr(qA, t| X)", main = "Joint Density of qA and t",  axes = TRUE, nticks = 5, ticktype = "detailed", ltheta = -180, lphi = 145, shade = TRUE, box = TRUE)
dev.off()