Skip to content

Instantly share code, notes, and snippets.

@ptitle
Created September 29, 2019 14:49
Show Gist options
  • Save ptitle/4c98cdcfb4ff507b2aec4ba8b9a914ef to your computer and use it in GitHub Desktop.
Save ptitle/4c98cdcfb4ff507b2aec4ba8b9a914ef to your computer and use it in GitHub Desktop.
BAMM convergence diagnosis plot
#BAMM run evaluation function
# libraries required: coda, data.table, BAMMtools
# if mcmcFile or chainFile are NULL, then it will look for the most likely files
# returns effective size and Geweke z-score for the log likelihood and the number of shifts
# as well as frequency of hot/cold chain swapping, if multi-chain analysis.
# ideal:
# effective size > 200
# Geweke z-score between -2 and 2
# chain swap frequency between 20 and 60%.
bammCheck <- function(expectedNumberOfShifts = 1, burnin=0.15, mcmcFile = NULL, chainFile = NULL, plotToPDF = FALSE, title = NULL) {
if (is.null(mcmcFile)) {
files <- list.files(pattern='mcmc_out')
if (length(files) == 1) {
mcmcFile <- files
} else {
stop('Please specify an mcmc_out file. Several were detected.')
}
}
if (is.null(chainFile)) {
files <- list.files(pattern='chain_swap')
if (length(files) == 1) {
chainFile <- files
} else if (length(files) == 0) {
chainFile <- 'chain_swap'
cat('\n\tNo chain swap file detected. Assuming single chain analysis.\n')
} else {
stop('Please specify an mcmc_out file. Several were detected.')
}
}
if (class(mcmcFile) == 'character') {
mcmc <- data.table::fread(mcmcFile, data.table = FALSE)
}
mcmc2 <- mcmc[floor(burnin * nrow(mcmc)):nrow(mcmc), ]
# if BAMM was running with a single chain, there will be no chain swap file.
if (chainFile %in% list.files()) {
chainswap <- data.table::fread(chainFile, data.table = FALSE)
chainswap <- chainswap[(burnin * nrow(chainswap)):nrow(chainswap),]
chainswap <- chainswap[which(chainswap$rank_1 == 1), ]
success <- length(which(chainswap$swapAccepted == 1)) / nrow(chainswap)
success <- round(success, digits = 2)
}
#get prior distribution of shifts
obsK <- seq(from=0, to = max(mcmc2[, "N_shifts"]), by = 1)
prior <- sapply(obsK, prob.k, poissonRatePrior=1 / expectedNumberOfShifts)
prior <- data.frame(N_shifts = obsK, prob = prior)
#get posterior distribution of shifts
posterior <- sapply(obsK, function(x) length(which(mcmc2[,'N_shifts'] == x))) / nrow(mcmc2)
names(posterior) <- obsK
posterior <- data.frame(N_shifts = names(posterior), prob=posterior)
if (!plotToPDF) {
dev.new(width = 10, height = 7)
}
par(mfrow = c(2, 2), mar = c(4, 4, 3, 3))
plot(mcmc[, 1], mcmc[, 4], type = 'l', lwd = 0.5, xlab = 'generations', ylab = 'loglik')
abline(v = max(mcmc[,1]) * burnin, lty = 2)
if (!is.null(title)) {
title(main = title)
}
#barplot(table(mcmc2$N_shifts), xlab='n shifts')
barplot(prior[,2], names.arg = prior[,1], ylim = c(0, max(c(prior[,2], posterior[,2]))), border = 'black', col = 'light blue', xlab = 'n shifts')
barplot(posterior[,2], add = TRUE, border = 'black', col = BAMMtools::transparentColor('red', 0.4), axes=FALSE)
legend('topright', legend = c('prior','posterior'), fill = c('light blue', BAMMtools::transparentColor('red', 0.4)), bty = 'n', cex=1.5)
plot(mcmc[,1], mcmc[,2], type='p', pch=20, cex=0.5, xlab='generations', ylab='n Events')
abline(v=max(mcmc[,1]) * burnin, lty=2)
plot(c(0, 1), c(0, 1), ann = F, bty = 'n', type = 'n', xaxt = 'n', yaxt = 'n', xpd=NA)
text(x=0.5, y=0.9, paste("effective size for log lik: ", round(coda::effectiveSize(mcmc2[,4]), 2), sep=''))
text(x=0.5, y=0.7, paste("effective size for n Events: ", round(coda::effectiveSize(mcmc2[,2]), 2), sep=''))
text(x=0.5, y=0.5, paste("Geweke z-score for log lik: ", round(coda::geweke.diag(coda::as.mcmc(mcmc2[,4]))$z, 2), sep=''))
text(x=0.5, y=0.3, paste("Geweke z-score for n Events: ", round(coda::geweke.diag(coda::as.mcmc(mcmc2[,2]))$z, 2), sep=''))
if (chainFile %in% list.files()) {
text(x=0.5, y=0.1, paste("successful swap frequency with cold chain: ", success, sep=''))
}
}
prob.k <- function(k, poissonRatePrior) {
Denom <- (poissonRatePrior + 1) ^ (k + 1)
Prob <- poissonRatePrior / Denom
return(Prob)
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment