R/eba.R

In bivarhr: Bivariate Hurdle Regression with Bayesian Model Averaging

#' Extreme-Bounds Analysis (EBA) over Control-Variable Combinations
#'
#' Runs an Extreme-Bounds Analysis (EBA) over a predefined set of control
#' variable combinations, fitting (or re-fitting) the bivariate hurdle
#' model for each combination and extracting posterior mean coefficients
#' for all regression blocks (\code{mu_I}, \code{pi_I}, \code{mu_C},
#' \code{pi_C}).
#'
#' @param DT A \code{data.table} or \code{data.frame} with the data passed
#'   to \code{fit_one()}.
#' @param spec Character scalar; model specification (e.g.\ \code{"A"},
#'   \code{"B"}, \code{"C"}, \code{"D"}) passed to \code{fit_one()}.
#' @param k_bma_table Optional object (typically a named list or
#'   \code{list}-like structure) indexed by control-combination tags that
#'   indicates for which combinations a BMA selection table already exists.
#'   If \code{k_bma_table[[tag]]} is \code{NULL} or \code{bma_weights_*}
#'   CSV is missing, the function falls back to a default fit with
#'   \code{k = 2} and default horseshoe hyperparameters.
#' @param seed Integer; base random seed for the fits. For different
#'   control combinations, the seed is jittered to avoid identical
#'   pseudo-random sequences.
#' @param control_combos A named list whose names are control tags
#'   (e.g.\ \code{"None"}, \code{"X1+X2"}, \code{"X1+X3+X4"}) and whose
#'   elements are the corresponding character vectors of control-variable
#'   names, defining which control sets to explore.
#' @param dir_csv Character scalar or \code{NULL}; directory where BMA
#'   weight CSV files are read from and where the coefficient table
#'   (\code{"eba_coefficients.csv"}) is written. If \code{NULL} (default),
#'   no BMA files are read (each combination uses the default fit) and no
#'   CSV is written.
#'
#' @details
#' This function relies on \code{\link{fit_one}()}, which requires
#' \pkg{cmdstanr} and a working CmdStan installation.
#'
#' For each control-combination tag \code{tag}:
#' \itemize{
#'   \item If a BMA weights file
#'     \code{"bma_weights_spec<spec>_ctrl<tag>.csv"} exists in
#'     \code{dir_csv} and \code{k_bma_table[[tag]]} is not \code{NULL},
#'     the top-weighted row (highest \code{weight}) is used to select
#'     \code{k} and horseshoe hyperparameters (\code{hs_tau0},
#'     \code{hs_slab_scale}, \code{hs_slab_df}) for the fit.
#'   \item Otherwise, the model is fit with \code{k = 2} and default
#'     horseshoe hyperparameters.
#'   \item Posterior means of the regression coefficients with prefixes
#'     \code{"b_mu_I"}, \code{"b_pi_I"}, \code{"b_mu_C"},
#'     \code{"b_pi_C"} are extracted and mapped back to the corresponding
#'     column names of the design matrices.
#' }
#'
#' All coefficient summaries are stacked into a single table and, when
#' \code{dir_csv} is supplied, written to \code{"eba_coefficients.csv"} in
#' that directory.
#'
#' @return A \code{data.frame} with the columns:
#' \itemize{
#'   \item \code{name}: name of the covariate (design-matrix column).
#'   \item \code{mean}: posterior mean of the corresponding coefficient.
#'   \item \code{block}: block identifier (\code{"mu_I"}, \code{"pi_I"},
#'     \code{"mu_C"}, \code{"pi_C"}).
#'   \item \code{combo}: control-combination tag used for that fit.
#' }
#'
#' @examples
#' \donttest{
#' # This example fits Stan models and therefore runs only in an
#' # interactive session with 'cmdstanr' and a working CmdStan.
#' if (interactive() && requireNamespace("cmdstanr", quietly = TRUE)) {
#'   DT <- data.table::data.table(
#'     I = rpois(21, lambda = 4),
#'     C = rpois(21, lambda = 3),
#'     zI = rnorm(21),
#'     zC = rnorm(21),
#'     t_norm = seq(-1, 1, length.out = 21),
#'     t_poly2 = seq(-1, 1, length.out = 21)^2,
#'     Regime = factor(sample(c("A", "B"), 21, replace = TRUE)),
#'     trans_PS = sample(0:1, 21, replace = TRUE),
#'     trans_SF = sample(0:1, 21, replace = TRUE),
#'     trans_FC = sample(0:1, 21, replace = TRUE),
#'     log_exposure50 = rep(0, 21),
#'     X1 = rnorm(21), X2 = rnorm(21), X3 = rnorm(21)
#'   )
#'
#'   combos <- list(
#'     None       = character(0),
#'     "X1+X2"    = c("X1", "X2"),
#'     "X1+X2+X3" = c("X1", "X2", "X3")
#'   )
#'
#'   eba_tab <- run_eba(DT, spec = "C", control_combos = combos, seed = 123)
#'   print(head(eba_tab))
#' }
#' }
#'
#' @export

run_eba <- function(DT, spec = "C", control_combos, k_bma_table = NULL,
                    seed = 123, dir_csv = NULL) {
  combs <- names(control_combos)
  eba_all <- list()
  progressr::with_progress({
    p <- progressr::progressor(steps = length(combs))
    for (tag in combs) {
      controls <- if (tag=="None") character(0) else unlist(strsplit(tag, "\\+"))
      bma_path <- if (is.null(dir_csv)) NA_character_ else
        file.path(dir_csv, sprintf("bma_weights_spec%s_ctrl%s.csv", spec, if(tag=="None") "None" else tag))
      if (is.na(bma_path) || !file.exists(bma_path) || is.null(k_bma_table[[tag]])) {
        fit <- fit_one(DT, k=2, spec=spec, controls=controls, seed=seed)
      } else {
        tb <- readr::read_csv(bma_path, show_col_types = FALSE) %>% arrange(desc(weight)) %>% slice(1)
        fit <- fit_one(DT, k=tb$k, spec=spec, controls=controls, seed=seed+which(combs==tag),
                       hs_tau0 = tb$hs_tau0, hs_slab_scale = tb$hs_slab_scale, hs_slab_df = tb$hs_slab_df)
      }
      draws <- posterior::as_draws_df(fit$fit$draws())
      
      map_coefs <- function(prefix, X_cols, block){
        nm <- grep(paste0("^",prefix,"\\["), names(draws), value = TRUE)
        if (!length(nm)) return(NULL)
        idx <- as.integer(sub(".*\\[(\\d+)\\]","\\1", nm))
        vals <- colMeans(as.matrix(draws[, nm, drop = FALSE]))
        data.frame(name = X_cols[idx], mean = vals, block = block, stringsAsFactors = FALSE)
      }
      
      tab <- dplyr::bind_rows(
        map_coefs("b_mu_I", colnames(fit$des$X_mu_I), "mu_I"),
        map_coefs("b_pi_I", colnames(fit$des$X_pi_I), "pi_I"),
        map_coefs("b_mu_C", colnames(fit$des$X_mu_C), "mu_C"),
        map_coefs("b_pi_C", colnames(fit$des$X_pi_C), "pi_C")
      )
      
      tab$combo <- tag
      eba_all[[tag]] <- tab
      p(message = sprintf("EBA combo: %s", tag))
    }
  })
  eba <- dplyr::bind_rows(eba_all)
  if (!is.null(dir_csv)) {
    if (!dir.exists(dir_csv)) dir.create(dir_csv, recursive = TRUE)
    readr::write_csv(eba, file.path(dir_csv, "eba_coefficients.csv"))
  }
  eba
}