using bkmrhat with multiple imputation

bkmrhat_MI.R

# using BKMRhat with multiple imputation

library(bkmrhat)
library(mice)
library(bkmr)
library(future)

# create simulated dataset for testing
dat <- bkmr::SimData(n = 50, M = 4)

zdf = as.data.frame(dat$Z)
set.seed(1232)
#set 10% of one of the exposures to missing
zdf$z1 = ifelse(runif(nrow(zdf))>.1, zdf$z1, NA)

# multiple imputation with some arbitrary method
num_imputations = 40
zdfmi = mice(zdf, m=num_imputations)



# this is basically source code from kmbayes_parallel, modified for using MI datasets
future::plan(strategy = future::multisession, workers=8)
ff <- list()
ss = round(runif(num_imputations) * .Machine$integer.max)
for (ii in 1:num_imputations) {
  ff[[ii]] <- future({
    # note, need to explicitly invoke package for future implementation
    bkmr::kmbayes(y=dat$y, X=dat$X, Z = mice::complete(zdfmi, ii), iter = 5000)
  }, seed = ss[ii])
}
parallelfits <- future::value(ff)
class(parallelfits) <- c("bkmrfit.list", class(res))

# parallel estimates across chains, 50% burnin
ors_parallel = bkmrhat::OverallRiskSummaries_parallel(parallelfits)


# pooling sample to get summary estimate, 50% burnin
pooledfit = kmbayes_combine(res)

ors_pooled = bkmr::OverallRiskSummaries(pooledfit)


# Rubin's rule estimator
Vb = tapply(ors_parallel$est, ors_parallel$quantile, var)   # "between" variance
Vw = tapply(ors_parallel$sd^2, ors_parallel$quantile, mean) # "within" variance
mn = tapply(ors_parallel$est, ors_parallel$quantile, mean)  # point estimate
Vr = Vw + Vb*(1+1/num_imputations)
quant = tapply(ors_parallel$quantile, ors_parallel$quantile, mean) # silly trick

# pooled estimate using Rubin's rules (this seems to over-estimate variance in short chains)
data.frame(quantile = quant, est = mn, sd = sqrt(Vr))

#pooled estimate using standard Bayes (seems slightly more efficient)
ors_pooled

This works on an intel-based Mac with R + packages mostly up-to-date in 12/23