Parallel chains with gam.mh()

sample_gam_posterior.R

sample_gam_posterior <- function(multinomial_model,
                                 chains = 4,
                                 cores = min(chains, parallel::detectCores()), ...) {

  samps <-
    purrr::transpose(parallel::mclapply(
      X = seq_len(chains),
      FUN = function(X) {
        post <- gam.mh(multinomial_model, ...)
        post
      },
      mc.cores = cores,
      mc.set.seed = TRUE
    ))

  gam_posterior <-
    structure(
      aperm(abind::abind(samps$bs, along = 3), c(1, 3, 2)),
      rw.accept = unlist(samps$rw.accept),
      accept = unlist(samps$accept)
    )
  dimnames(gam_posterior) <- list(
    Iteration = NULL, Chain = NULL, Parameter = dimnames(gam_posterior)[[3]]
  )

  gam_posterior
}

calc_posterior_stats <- function(gam_posterior) {
  posterior_stats <- tibble(
    parameter = dimnames(gam_posterior)[[3]],
    Rhat = apply(gam_posterior, 3, rstan::Rhat),
    ess_bulk = apply(gam_posterior, 3, rstan::ess_bulk),
    ess_tail = apply(gam_posterior, 3, rstan::ess_tail)
  )

  posterior_stats
}