bkmr
and bkmrhat
bkmr
is a package to implement Bayesian kernel machine regression (BKMR) using Markov chain Monte Carlo (MCMC). Notably, bkmr
is missing some key features in Bayesian inference and MCMC diagnostics: 1) no facility for running multiple chains in parallel 2) no inference across multiple chains 3) limited posterior summary of parameters 4) limited diagnostics. The bkmrhat
package is a lightweight set of function that fills in each of those gaps by enabling post-processing of bkmr
output in other packages and building a small framework for parallel processing.
bkmrhat
packagekmbaryes
function from bkmr
, or use multiple parallel chains kmbayes_parallel
from bkmrhat
kmbayes_diagnose
function (uses functions from the rstan
package) OR convert the BKMR fit(s) to mcmc
(one chain) or mcmc.list
(multiple chains) objects from the coda
package using as.mcmc
or as.mcmc.list
from the bkmrhat
package. The coda
package has a whole host of inference and diagnostic procedures (but may lag behind some of the diagnostics functions from rstan
).coda
functions or combine chains from a kmbayes_parallel
fit using kmbayes_combine
. Final posterior inferences can be made on the combined object, which enables use of bkmr
package functions for visual summaries of independent and joint effects of exposures in the bkmr
model.First, simulate some data from the bkmr
function
library("bkmr") library("bkmrhat") library("coda") set.seed(111) dat <- bkmr::SimData(n = 50, M = 5, ind=1:3, Zgen="realistic") y <- dat$y Z <- dat$Z X <- cbind(dat$X, rnorm(50)) head(cbind(y,Z,X))
There is some overhead in parallel processing when using the future
package, so the payoff when using parallel processing may vary by the problem. Here it is about a 2-4x speedup, but you can see more benefit at higher iterations. Note that this may not yield as many usable iterations as a single large chain if a substantial burnin period is needed, but it will enable useful convergence diagnostics. Note that the future package can implement sequential processing, which effectively turns the kmbayes_parallel into a loop, but still has all other advantages of multiple chains.
# enable parallel processing (up to 4 simultaneous processes here) future::plan(strategy = future::multisession) # single run of 4000 observations from bkmr package set.seed(111) system.time(kmfit <- suppressMessages(kmbayes(y = y, Z = Z, X = X, iter = 4000, verbose = FALSE, varsel = FALSE))) # 4 runs of 1000 observations from bkmrhat package set.seed(111) system.time(kmfit5 <- suppressMessages(kmbayes_parallel(nchains=4, y = y, Z = Z, X = X, iter = 1000, verbose = FALSE, varsel = FALSE)))
The diagnostics from the rstan package come from the monitor
function (see the help files for that function in the rstan pacakge)
# Using rstan functions (set burnin/warmup to zero for comparability with coda numbers given later # posterior summaries should be performed after excluding warmup/burnin) singlediag = kmbayes_diagnose(kmfit, warmup=0, digits_summary=2) # Using rstan functions (multiple chains enable R-hat) multidiag = kmbayes_diagnose(kmfit5, warmup=0, digits_summary=2) # using coda functions, not using any burnin (for demonstration only) kmfitcoda = as.mcmc(kmfit, iterstart = 1) kmfit5coda = as.mcmc.list(kmfit5, iterstart = 1) # single chain trace plot traceplot(kmfitcoda)
The trace plots look typical, and fine, but trace plots don't give a full picture of convergence. Note that there is apparent quick convergence for a couple of parameters demonstrated by movement away from the starting value and concentration of the rest of the samples within a narrow band.
Seeing visual evidence that different chains are sampling from the same marginal distributions is reassuring about the stability of the results.
# multiple chain trace plot traceplot(kmfit5coda)
Now examine "cross correlation", which can help identify highly correlated parameters in the posterior, which can be problematic for MCMC sampling. Here there is a block {r3,r4,r5} which appear to be highly correlated. All other things equal, having highly correlated parameters in the posterior means that more samples are needed than would be needed with uncorrelated parameters.
# multiple cross-correlation plot (combines all samples) crosscorr(kmfit5coda) crosscorr.plot(kmfit5coda)
Now examine "autocorrelation" to identify parameters that have high correlation between subsequent iterations of the MCMC sampler, which can lead to inefficient MCMC sampling. All other things equal, having highly autocorrelated parameters in the posterior means that more samples are needed than would be needed with low-autocorrelation parameters.
# multiple chain trace plot #autocorr(kmfit5coda) # lots of output autocorr.plot(kmfit5coda)
Graphical tools can be limited, and are sometimes difficult to use effectively with scale parameters (of which bkmr
has many). Additionally, no single diagnostic is perfect, leading many authors to advocate the use of multiple, complementary diagnostics. Thus, more formal diagnostics are helpful.
Gelman's r-hat diagnostic gives an interpretable diagnostic: the expected reduction in the standard error of the posterior means if you could run the chains to an infinite size. These give some idea about when is a fine idea to stop sampling. There are rules of thumb about using r-hat to stop sampling that are available from several authors (for example you can consult the help files for rstan
and coda
).
Effective sample size is also useful - it estimates the amount of information in your chain, expressed in terms of the number of independent posterior samples it would take to match that information (e.g. if we could just sample from the posterior directly).
# Gelman's r-hat using coda estimator (will differ from rstan implementation) gelman.diag(kmfit5coda) # effective sample size effectiveSize(kmfitcoda) effectiveSize(kmfit5coda)
Posterior kernel marginal densities, 1 chain
# posterior kernel marginal densities using `mcmc` and `mcmc` objects densplot(kmfitcoda)
Posterior kernel marginal densities, multiple chains combined. Look for multiple modes that may indicate non-convergence of some chains
# posterior kernel marginal densities using `mcmc` and `mcmc` objects densplot(kmfit5coda)
Other diagnostics from the coda
package are available here.
Finally, the chains from the original kmbayes_parallel
fit can be combined into a single chain (see the help files for how to deal with burn-in, the default in bkmr
is to use the first half of the chain, which is respected here). The kmbayes_combine
function smartly first combines the burn-in iterations and then combines the iterations after burnin, such that the burn-in rules of subsequent functions within the bkmr
package are respected. Note that unlike the as.mcmc.list
function, this function combines all iterations into a single chain, so trace plots will not be good diagnotistics in this combined object, and it should be used once one is assured that all chains have converged and the burn-in is acceptable.
With this combined set of samples, you can follow any of the post-processing functions from the bkmr
functions, which are described here: https://jenfb.github.io/bkmr/overview.html. For example, see below the estimation of the posterior mean difference along a series of quantiles of all exposures in Z.
# posterior summaries using `mcmc` and `mcmc` objects summary(kmfitcoda) summary(kmfit5coda) # highest posterior density intervals using `mcmc` and `mcmc` objects HPDinterval(kmfitcoda) HPDinterval(kmfit5coda) # combine multiple chains into a single chain fitkmccomb = kmbayes_combine(kmfit5) # For example: summary(fitkmccomb) mean.difference <- suppressWarnings(OverallRiskSummaries(fit = fitkmccomb, y = y, Z = Z, X = X, qs = seq(0.25, 0.75, by = 0.05), q.fixed = 0.5, method = "exact")) mean.difference with(mean.difference, { plot(quantile, est, pch=19, ylim=c(min(est - 1.96*sd), max(est + 1.96*sd)), axes=FALSE, ylab= "Mean difference", xlab = "Joint quantile") segments(x0=quantile, x1=quantile, y0 = est - 1.96*sd, y1 = est + 1.96*sd) abline(h=0) axis(1) axis(2) box(bty='l') })
These results parallel previous session and are given here without comment, other than to note that no fixed effects (X variables) are included, and that it is useful to check the posterior inclusion probabilities to ensure they are similar across chains.
set.seed(111) system.time(kmfitbma.list <- suppressWarnings(kmbayes_parallel(nchains=4, y = y, Z = Z, X = X, iter = 1000, verbose = FALSE, varsel = TRUE))) bmadiag = kmbayes_diagnose(kmfitbma.list) # posterior exclusion probability of each chain lapply(kmfitbma.list, function(x) t(ExtractPIPs(x))) kmfitbma.comb = kmbayes_combine(kmfitbma.list) summary(kmfitbma.comb) ExtractPIPs(kmfitbma.comb) # posterior inclusion probabilities mean.difference2 <- suppressWarnings(OverallRiskSummaries(fit = kmfitbma.comb, y = y, Z = Z, X = X, qs = seq(0.25, 0.75, by = 0.05), q.fixed = 0.5, method = "exact")) mean.difference2 with(mean.difference2, { plot(quantile, est, pch=19, ylim=c(min(est - 1.96*sd), max(est + 1.96*sd)), axes=FALSE, ylab= "Mean difference", xlab = "Joint quantile") segments(x0=quantile, x1=quantile, y0 = est - 1.96*sd, y1 = est + 1.96*sd) abline(h=0) axis(1) axis(2) box(bty='l') })
bkmrhat
also has ported versions of the native posterior summarization functions to compare how these summaries vary across parallel chains. Note that these should serve as diagnostics, and final posterior inference should be done on the combined chain. The easiest of these functions to demonstrate is the OverallRiskSummaries_parallel
function, which simply runs OverallRiskSummaries
(from the bkmr
package) on each chain and combines the results. Notably, this function fixes the y-axis at zero for the median, so it under-represents overall predictive variation across chains, but captures variation in effect estimates across the chains. Ideally, that variation is negligible - e.g. if you see differences between chains that would result in different interpretations, you should re-fit the model with more iterations. In this example, the results are reasonably consistent across chains, but one might want to run more iterations if, say, the differences seen across the upper error bounds are of such a magnitude as to be practically meaningful.
set.seed(111) system.time(kmfitbma.list <- suppressWarnings(kmbayes_parallel(nchains=4, y = y, Z = Z, X = X, iter = 1000, verbose = FALSE, varsel = TRUE))) meandifference_par = OverallRiskSummaries_parallel(kmfitbma.list, y = y, Z = Z, X = X ,qs = seq(0.25, 0.75, by = 0.05), q.fixed = 0.5, method = "exact") head(meandifference_par) nchains = length(unique(meandifference_par$chain)) with(meandifference_par, { plot.new() plot.window(ylim=c(min(est - 1.96*sd), max(est + 1.96*sd)), xlim=c(min(quantile), max(quantile)), ylab= "Mean difference", xlab = "Joint quantile") for(cch in seq_len(nchains)){ width = diff(quantile)[1] jit = runif(1, -width/5, width/5) points(jit+quantile[chain==cch], est[chain==cch], pch=19, col=cch) segments(x0=jit+quantile[chain==cch], x1=jit+quantile[chain==cch], y0 = est[chain==cch] - 1.96*sd[chain==cch], y1 = est[chain==cch] + 1.96*sd[chain==cch], col=cch) } abline(h=0) axis(1) axis(2) box(bty='l') legend("bottom", col=1:nchains, pch=19, lty=1, legend=paste("chain", 1:nchains), bty="n") }) regfuns_par = PredictorResponseUnivar_parallel(kmfitbma.list, y = y, Z = Z, X = X ,qs = seq(0.25, 0.75, by = 0.05), q.fixed = 0.5, method = "exact") head(regfuns_par) nchains = length(unique(meandifference_par$chain)) # single variable with(regfuns_par[regfuns_par$variable=="z1",], { plot.new() plot.window(ylim=c(min(est - 1.96*se), max(est + 1.96*se)), xlim=c(min(z), max(z)), ylab= "Predicted Y", xlab = "Z") pc = c("#000000", "#E69F00", "#56B4E9", "#009E73", "#F0E442", "#0072B2", "#D55E00", "#CC79A7", "#999999") pc2 = c("#0000001A", "#E69F001A", "#56B4E91A", "#009E731A", "#F0E4421A", "#0072B21A", "#D55E001A", "#CC79A71A", "#9999991A") for(cch in seq_len(nchains)){ ribbonX = c(z[chain==cch], rev(z[chain==cch])) ribbonY = c(est[chain==cch] + 1.96*se[chain==cch], rev(est[chain==cch] - 1.96*se[chain==cch])) polygon(x=ribbonX, y = ribbonY, col=pc2[cch], border=NA) lines(z[chain==cch], est[chain==cch], pch=19, col=pc[cch]) } axis(1) axis(2) box(bty='l') legend("bottom", col=1:nchains, pch=19, lty=1, legend=paste("chain", 1:nchains), bty="n") })
Sometimes you just need to run more samples in an existing chain. For example, you run a bkmr fit for 3 days, only to find you don't have enough samples. A "continued" fit just means that you can start off at the last iteration you were at and just keep building on an existing set of results by lengthening the Markov chain. Unfortunately, due to how the kmbayes
function accepts starting values (for the official install version), you can't quite do this exactly in many cases (The function will relay a message and possible solutions, if any. bkmr
package authors are aware of this issue). The kmbayes_continue
function continues a bkmr
fit as well as the bkmr
package will allow. The r
parameters from the fit must all be initialized at the same value, so kmbayes_continue
starts a new MCMC fit at the final values of all parameters from the prior bkmr fit, but sets all of the r
parameters to the mean at the last iteration from the previous fit. Additionally, if h.hat
parameters are estimated, these are fixed to be above zero to meet similar constraints, either by fixing them at their posterior mean or setting to a small positive value. One should inspect trace plots to see whether this will cause issues (e.g. if the traceplots demonstrate different patterns in the samples before and after the continuation). Here's an example with a quick check of diagnostics of the first part of the chain, and the combined chain (which could be used for inference or extended again, if necessary). We caution users that this function creates 2 distinct, if very similar Markov chains, and to use appropriate caution if traceplots differ before and after each continuation. Nonetheless, in many cases one can act as though all samples are from the same Markov chain.
Note that if you install the developmental version of the bkmr
package you can continue fits from exactly where they left off, so you get a true, single Markov chain. You can install that via the commented code below
# install dev version of bkmr to allow true continued fits. #install.packages("devtools") #devtools::install_github("jenfb/bkmr") set.seed(111) # run 100 initial iterations for a model with only 2 exposures Z2 = Z[,1:2] kmfitbma.start <- suppressWarnings(kmbayes(y = y, Z = Z2, X = X, iter = 500, verbose = FALSE, varsel = FALSE)) kmbayes_diag(kmfitbma.start) # run 2000 additional iterations moreiterations = kmbayes_continue(kmfitbma.start, iter=2000) kmbayes_diag(moreiterations) TracePlot(moreiterations, par="beta") TracePlot(moreiterations, par="r")
Thanks to Haotian "Howie" Wu for invaluable feedback on early versions of the package.
Any scripts or data that you put into this service are public.
Add the following code to your website.
For more information on customizing the embed code, read Embedding Snippets.