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