# ============================================================================ # swcrt_bart_analysis.R # # Calibrated Bayesian analysis of stepped-wedge cluster-randomized trials # using mixed-effects BART (mxBART). Supports: # - g-computation ("gform") and model-robust ("mr") estimators # - Simple-exchangeable and nested-exchangeable random effects # - Overall and subgroup analyses # - User-specified covariate adjustment via formula # # Depends: mxBART # ============================================================================ # --------------------------------------------------------------------------- # 1. prepare_data() # --------------------------------------------------------------------------- # Goal: # Take a cleaned analysis data frame and add the helper columns needed by # the downstream pipeline: standardized variable names for the design # columns (cluster, period, treat, y), a factor version of period, # a cluster-period interaction column (cp), and cluster-period sizes (N_ij). # # Arguments: # data : data.frame. One row per individual. Must contain the four # design columns named below, plus all covariates referenced # later via covariate_formula. # cluster_col : string. Column name for cluster ID. # period_col : string. Column name for period / step (integer 1..J). # treat_col : string. Column name for treatment indicator (0/1). # outcome_col : string. Column name for the (already transformed) outcome. # # Output: # A data.frame with the original columns plus: # cluster - factor version of cluster ID # period - integer period (renamed from period_col) # period_f - factor version of period # treat - numeric 0/1 (renamed from treat_col) # y - numeric outcome (renamed from outcome_col) # cp - factor, interaction(cluster, period_f) # N_ij - integer, cluster-period sample size repeated on each row # row_id - integer 1..n # --------------------------------------------------------------------------- prepare_data <- function(data, cluster_col, period_col, treat_col, outcome_col) { dat <- data # Rename user-supplied column names to the standardized vocabulary # the rest of the pipeline expects. names(dat)[names(dat) == cluster_col] <- "cluster" names(dat)[names(dat) == period_col] <- "period" names(dat)[names(dat) == treat_col] <- "treat" names(dat)[names(dat) == outcome_col] <- "y" dat$treat <- as.numeric(dat$treat) dat$cluster <- factor(dat$cluster) dat$period_f <- factor(dat$period) dat$cp <- interaction(dat$cluster, dat$period_f, drop = TRUE) dat$row_id <- seq_len(nrow(dat)) # Cluster-period size, repeated on every row of that cluster-period dat$N_ij <- ave(dat$y, dat$cluster, dat$period, FUN = function(z) sum(!is.na(z))) dat } # --------------------------------------------------------------------------- # 2. make_design_matrix() # --------------------------------------------------------------------------- # Goal: # Build the fixed-effect design matrix for mxBART. This always includes # period dummies, the treatment indicator, and the cluster-period size # N_ij (all required by the SW-CRT analysis), plus any user-specified # covariates passed through covariate_formula. # # Arguments: # data : data.frame produced by prepare_data(). # covariate_formula : one-sided formula (e.g. ~ age + gender + factor(race)) # listing the additional covariates to adjust for. # Pass ~ 1 or NULL to adjust for nothing beyond the # defaults (period dummies, treat, N_ij). # # Output: # Numeric matrix with n rows. Column order: # period dummies (one per period, no intercept) | treat | covariates | N_ij # --------------------------------------------------------------------------- make_design_matrix <- function(data, covariate_formula = NULL) { # Period dummies (no intercept => one dummy per period level) P <- model.matrix(~ -1 + period_f, data = data) # User-specified covariates (drop the intercept column from model.matrix) if (is.null(covariate_formula) || identical(deparse(covariate_formula), "~1")) { Xcov <- NULL } else { Xcov <- model.matrix(covariate_formula, data = data) if ("(Intercept)" %in% colnames(Xcov)) { Xcov <- Xcov[, colnames(Xcov) != "(Intercept)", drop = FALSE] } } out <- cbind(P, treat = as.numeric(data$treat)) if (!is.null(Xcov)) out <- cbind(out, Xcov) out <- cbind(out, N_ij = as.numeric(data$N_ij)) out } # --------------------------------------------------------------------------- # 3. bootstrap_clusters_idx() # --------------------------------------------------------------------------- # Goal: # Resample whole clusters with replacement (all periods/rows of each # sampled cluster), relabel sampled clusters to 1..I, and return both # the resampled data frame and the integer index mapping back to the # original rows (in bootstrap order). The index is needed so that the # posterior-prediction matrices Y0, Y1 can be reordered to match the # bootstrap data, since the BART model is only fit once and predictions # are re-used across bootstrap resamples. # # Arguments: # data : data.frame produced by prepare_data(). # cluster_var : string, name of cluster column (default "cluster"). # period_var : string, name of integer period column (default "period"). # period_f_var : string, name of factor period column (default "period_f"). # # Output: # List with two elements: # dat - resampled data.frame, clusters relabeled 1..I, cp rebuilt # index - integer vector of original-row indices, length nrow(dat). # Use this to reorder Y0, Y1: e.g., Y0_boot <- Y0[index, ]. # --------------------------------------------------------------------------- bootstrap_clusters_idx <- function(data, cluster_var = "cluster", period_var = "period", period_f_var = "period_f") { cl <- as.integer(factor(data[[cluster_var]])) ucl <- sort(unique(cl)) samp <- sample(ucl, length(ucl), replace = TRUE) idx_vec <- integer(0) out <- vector("list", length(samp)) for (new_id in seq_along(samp)) { rows <- which(cl == samp[new_id]) d <- data[rows, , drop = FALSE] d[[cluster_var]] <- new_id out[[new_id]] <- d idx_vec <- c(idx_vec, rows) } res <- do.call(rbind, out) res[[period_f_var]] <- factor(res[[period_var]], levels = levels(data[[period_f_var]])) res$cp <- interaction(res[[cluster_var]], res[[period_f_var]], drop = TRUE) list(dat = res, index = idx_vec) } # --------------------------------------------------------------------------- # 4. compute_ATEs_gform() # --------------------------------------------------------------------------- # Goal: # Compute posterior draws of i-ATE and c-ATE under the g-computation # estimator, for the full sample or for a user-specified subgroup. # # Arguments: # data : data.frame produced by prepare_data(). # Y0, Y1 : n x B posterior prediction matrices under treat = 0 and # treat = 1 respectively (B = number of posterior draws). # subgroup_var : NULL (default) for overall analysis, or string naming a # column in `data` to restrict to. # subgroup_val : value of subgroup_var defining the subgroup of interest. # Required when subgroup_var is non-NULL. # cluster_var, # period_f_var : column names (defaults match prepare_data() output). # # Output: # data.frame with B rows and two columns (i_ATE, c_ATE). Each row is one # posterior draw. # # Note on subgroup handling: # The cluster-period structure is kept intact and the non-subgroup # individuals are masked out of the aggregation via S_r * y. N_ij is # replaced by the subgroup count within each cluster-period (N_cp); # periods/clusters with zero subgroup members are dropped from the # average. When subgroup_var is NULL, S_r is all 1's and the function # reduces to the overall g-computation analysis. # --------------------------------------------------------------------------- compute_ATEs_gform <- function(data, Y0, Y1, subgroup_var = NULL, subgroup_val = NULL, cluster_var = "cluster", period_f_var = "period_f") { stopifnot(nrow(Y0) == nrow(data), nrow(Y1) == nrow(data), ncol(Y0) == ncol(Y1)) if (!is.null(subgroup_var) && is.null(subgroup_val)) { stop("subgroup_val must be specified when subgroup_var is given.") } # Subgroup indicator: 1 for subgroup members, 0 otherwise. # When subgroup_var is NULL, S_r is all 1s -> reduces to overall analysis. S_r <- if (is.null(subgroup_var)) { rep(1L, nrow(data)) } else { as.integer(data[[subgroup_var]] == subgroup_val) } dat <- data i_vec <- as.integer(dat[[cluster_var]]) j_vec <- as.integer(dat[[period_f_var]]) Jmax <- max(j_vec) cp_df <- unique(data.frame(i = i_vec, j = j_vec)) cp_df <- cp_df[order(cp_df$i, cp_df$j), ] row_key <- paste(i_vec, j_vec, sep = "_") cp_key <- paste(cp_df$i, cp_df$j, sep = "_") cp_index <- factor(row_key, levels = cp_key) cp_i <- cp_df$i cp_j <- cp_df$j keep_mid <- (cp_j >= 2 & cp_j <= (Jmax - 1)) ## subgroup cp sizes: N^{(r)}_{ij} N_cp <- as.numeric(rowsum(S_r, cp_index)) one_draw <- function(y0, y1) { cp_y0 <- as.numeric(rowsum(y0 * S_r, cp_index)) cp_y1 <- as.numeric(rowsum(y1 * S_r, cp_index)) # already have cp_y0, cp_y1 masked by S, and N_cp for subgroup sizes mean_y0_by_j <- tapply(cp_y0, cp_j, mean) mean_y1_by_j <- tapply(cp_y1, cp_j, mean) mean_Ncp_by_j <- tapply(N_cp, cp_j, mean) idx <- 2:(Jmax - 1) # drop periods with zero subgroup everywhere (avoid 0/0) valid_idx <- idx[ mean_Ncp_by_j[idx] > 0 ] i_ATE <- if (length(valid_idx) == 0L) NA_real_ else (sum(mean_y1_by_j[valid_idx]) - sum(mean_y0_by_j[valid_idx])) / sum(mean_Ncp_by_j[valid_idx]) # subgroup c-ATE with mid-period normalization by subgroup sizes sum_y1_by_i <- tapply(cp_y1[keep_mid], cp_i[keep_mid], sum) sum_y0_by_i <- tapply(cp_y0[keep_mid], cp_i[keep_mid], sum) sum_Nr_by_i <- tapply(N_cp [keep_mid], cp_i[keep_mid], sum) # subgroup counts nonzero <- sum_Nr_by_i > 0 c_ATE <- mean( (sum_y1_by_i[nonzero] / sum_Nr_by_i[nonzero]) - (sum_y0_by_i[nonzero] / sum_Nr_by_i[nonzero]) ) c(i_ATE = i_ATE, c_ATE = c_ATE) } B <- ncol(Y0) out <- matrix(NA_real_, B, 2, dimnames = list(NULL, c("i_ATE","c_ATE"))) for (b in seq_len(B)) out[b, ] <- one_draw(Y0[, b], Y1[, b]) as.data.frame(out) } # --------------------------------------------------------------------------- # 5. compute_ATEs_mr() # --------------------------------------------------------------------------- # Goal: # Compute posterior draws of i-ATE and c-ATE under the model-robust # estimator, for the full sample or for a user-specified subgroup. # # Arguments: # data : data.frame produced by prepare_data(). # Y0, Y1 : n x B posterior prediction matrices under treat = 0 and # treat = 1 respectively (B = number of posterior draws). # subgroup_var : NULL (default) for overall analysis, or string naming a # column in `data` to restrict to. # subgroup_val : value of subgroup_var defining the subgroup of interest. # Required when subgroup_var is non-NULL. # cluster_var, # period_f_var : column names (defaults match prepare_data() output). # # Output: # data.frame with B rows and two columns (i_ATE, c_ATE). Each row is one # posterior draw. # # Note on subgroup handling: # Unlike the g-computation estimator (which uses masking), the model- # robust estimator filters data, Y0, Y1 down to the subgroup before # computing the residualized and plug-in pieces. The residualization # step needs observed ybar_ij and predicted muhat_ij on the same units, # so filtering (rather than masking) is the natural choice here. # --------------------------------------------------------------------------- compute_ATEs_mr <- function(data, Y0, Y1, subgroup_var = NULL, subgroup_val = NULL, cluster_var = "cluster", period_f_var = "period_f") { dat <- data stopifnot(nrow(Y0) == nrow(dat), nrow(Y1) == nrow(dat), ncol(Y0) == ncol(Y1)) if (!is.null(subgroup_var) && is.null(subgroup_val)) { stop("subgroup_val must be specified when subgroup_var is given.") } # --- subgroup filter (skip if overall) --- if (!is.null(subgroup_var)) { idx <- dat[[subgroup_var]] == subgroup_val dat <- dat[idx, , drop = FALSE] Y0 <- Y0[idx, , drop = FALSE] Y1 <- Y1[idx, , drop = FALSE] } ## ----- cluster-period indexing ----- i_vec <- as.integer(dat[[cluster_var]]) j_vec <- as.integer(dat[[period_f_var]]) Jmax <- max(j_vec) cp_df <- unique(data.frame(i = i_vec, j = j_vec)) cp_df <- cp_df[order(cp_df$i, cp_df$j), ] row_key <- paste(i_vec, j_vec, sep = "_") cp_key <- paste(cp_df$i, cp_df$j, sep = "_") cp_index <- factor(row_key, levels = cp_key) cp_i <- cp_df$i cp_j <- cp_df$j Nij_cp <- as.numeric(table(cp_index)) # N_{ij} keep_mid <- (cp_j >= 2 & cp_j <= (Jmax - 1)) # j = 2..J-1 ## ----- observed cluster-period means: Ybar_ij ----- ybar_cp <- as.numeric(rowsum(dat$y, cp_index)) / Nij_cp ## ----- cluster-period treatment indicators Z_{ij} (should be constant within cp) ----- z_cp <- as.integer(tapply(dat$treat, cp_index, function(v) { u <- unique(v); if (length(u) != 1L) stop("Mixed Z within a cluster-period."); u })) ## ----- weights ----- w_ij_i <- Nij_cp w_j_i <- as.numeric(tapply(w_ij_i, cp_j, sum)) ## c-ATE Ni_total_mid_by_i <- tapply(Nij_cp[keep_mid], cp_i[keep_mid], sum) # Σ_{s=2..J-1} N_is Ni_total_mid <- Ni_total_mid_by_i[as.character(cp_i)] # align to cp order # compute per-cluster-period weights w_ij_c <- Nij_cp / Ni_total_mid w_j_c <- as.numeric(tapply(w_ij_c, cp_j, sum)) ## helper: compute for one posterior draw and one weighting scheme mu_mr_onez <- function(z, muhat_cp, w_ij, w_j) { idx_mid <- which(keep_mid) # indices of cp_df that are middle periods Js <- sort(unique(cp_j[keep_mid])) # the middle periods as integers Js <- Js[w_j[Js] > 0] if (length(Js) == 0L) return(NA_real_) denom_period <- sum(w_j[Js]) acc <- 0 for (jj in Js) { in_j <- (cp_j == jj) in_j_z <- in_j & (z_cp == z) ## residualized term denom_D <- sum(w_ij[in_j_z]) A_j <- if (denom_D > 0) { sum(w_ij[in_j_z] * (ybar_cp[in_j_z] - muhat_cp[in_j_z])) / denom_D } else { 0 } ## plug-in term B_j <- sum(w_ij[in_j] * muhat_cp[in_j]) / w_j[jj] acc <- acc + w_j[jj] * (A_j + B_j) } acc / denom_period } ## ----- loop over posterior draws B ----- B <- ncol(Y0) out <- matrix(NA_real_, B, 2, dimnames = list(NULL, c("i_ATE","c_ATE"))) for (b in seq_len(B)) { ## predicted cluster-period means under z=0 and z=1 for draw b muhat0_cp <- as.numeric(rowsum(Y0[, b], cp_index)) / Nij_cp muhat1_cp <- as.numeric(rowsum(Y1[, b], cp_index)) / Nij_cp ## i-ATE mu_i_0 <- mu_mr_onez(0, muhat0_cp, w_ij_i, w_j_i) mu_i_1 <- mu_mr_onez(1, muhat1_cp, w_ij_i, w_j_i) i_mr <- mu_i_1 - mu_i_0 ## c-ATE mu_c_0 <- mu_mr_onez(0, muhat0_cp, w_ij_c, w_j_c) mu_c_1 <- mu_mr_onez(1, muhat1_cp, w_ij_c, w_j_c) c_mr <- mu_c_1 - mu_c_0 out[b, ] <- c(i_mr, c_mr) } as.data.frame(out) } # --------------------------------------------------------------------------- # 6. fit_mxbart_ate() # --------------------------------------------------------------------------- # Goal: # Fit a mixed-effects BART (mxBART) model to the SW-CRT data and return # posterior prediction matrices Y0 (under treat = 0) and Y1 (under # treat = 1) with random effects set to zero at test time (achieved by # omitting id.test, which is the mxBART default behavior). Switches # between simple-exchangeable (cluster RE only) and nested-exchangeable # (cluster + cluster-period REs) via the re_structure argument. # # Arguments: # data : data.frame produced by prepare_data(). # covariate_formula : one-sided formula passed to make_design_matrix(). # re_structure : "simple_exch" (cluster RE only) or # "nested_exch" (cluster + cluster-period REs). # mctotal : total MCMC iterations. # mcburnin : number of burn-in iterations to discard. # thin_rate : keepevery argument to mxbart (1 = no thinning). # ntree : number of trees in the BART ensemble. # seed : integer seed passed to set.seed() before fitting. # # Output: # List with: # Y0 - n x B matrix of posterior predictions under treat = 0 # Y1 - n x B matrix of posterior predictions under treat = 1 # fit - the raw mxBART fit object (for diagnostics) # --------------------------------------------------------------------------- fit_mxbart_ate <- function(data, covariate_formula, re_structure = "simple_exch", mctotal = 4000, mcburnin = 2000, thin_rate = 1, ntree = 50L, seed = 123) { if (!re_structure %in% c("simple_exch", "nested_exch")) { stop('re_structure must be "simple_exch" or "nested_exch".') } set.seed(seed) ## -------- design matrix (train + counterfactual test) -------- x_train <- make_design_matrix(data, covariate_formula) y_train <- data$y nd0 <- data; nd0$treat <- 0 nd1 <- data; nd1$treat <- 1 X0 <- make_design_matrix(nd0, covariate_formula) X1 <- make_design_matrix(nd1, covariate_formula) x_test <- rbind(X0, X1) ## -------- random-effect design (switches by re_structure) -------- cluster_id <- as.integer(data$cluster) cp_id <- as.integer(data$cp) n <- nrow(data) if (re_structure == "simple_exch") { z_train <- list(matrix(1, n, 1)) id_train <- cbind(cluster_id) mxps <- list(list(prior = 1, df = 3, scale = 1)) } else { # nested_exch z_train <- list(matrix(1, n, 1), matrix(1, n, 1)) id_train <- cbind(cluster_id, cp_id) mxps <- list(list(prior = 1, df = 3, scale = 1), list(prior = 1, df = 3, scale = 1)) } ## -------- mxBART fit -------- ndpost_val <- mctotal - mcburnin sigest_est <- stats::sd(y_train) fit <- mxBART::mxbart( y.train = y_train, x.train = x_train, id.train = id_train, z.train = z_train, x.test = x_test, sigest = sigest_est, mxps = mxps, nskip = mcburnin, ndpost = ndpost_val, ntree = ntree, printevery = 100L, keepevery = thin_rate ) ## -------- extract Y0, Y1 from fhat.test -------- posterior <- fit$fhat.test Y0 <- t(posterior[, 1:n, drop = FALSE]) # n x ndpost Y1 <- t(posterior[, (n + 1):(2*n), drop = FALSE]) list(Y0 = Y0, Y1 = Y1, fit = fit) } # --------------------------------------------------------------------------- # 7. calibrated_variance() # --------------------------------------------------------------------------- # Goal: # Compute the calibrated (cluster-bootstrap-corrected) variance of the # posterior mean estimates of i-ATE and c-ATE. The calibrated variance is # Var(row means of V_matrix) + Var(uncalibrated posterior draws), # which corrects the under-coverage of naive posterior quantile intervals # that ignore the cluster-correlated sampling variability. # # Arguments: # data : data.frame from prepare_data(). # Y0, Y1 : n x B posterior prediction matrices from # fit_mxbart_ate(). # posterior_ATEs : data.frame from compute_ATEs_gform() or # compute_ATEs_mr() (B rows, columns i_ATE, c_ATE). # estimator : "gform" or "mr". Must match what produced # posterior_ATEs. # subgroup_var, # subgroup_val : as in compute_ATEs_*(); NULL for overall. # num_bootstrap : number of cluster-bootstrap resamples (default 100). # verbose : if TRUE, print progress every 10 iterations. # # Output: # List with: # SE_c, SE_i - calibrated standard errors for c-ATE and i-ATE # Var_c, Var_i - calibrated variances # --------------------------------------------------------------------------- calibrated_variance <- function(data, Y0, Y1, posterior_ATEs, estimator = "gform", subgroup_var = NULL, subgroup_val = NULL, num_bootstrap = 100, verbose = TRUE) { if (!estimator %in% c("gform", "mr")) { stop('estimator must be "gform" or "mr".') } B <- ncol(Y0) V_c <- matrix(NA_real_, nrow = num_bootstrap, ncol = B) V_i <- matrix(NA_real_, nrow = num_bootstrap, ncol = B) for (m in seq_len(num_bootstrap)) { if (verbose && m %% 10 == 0) { cat(sprintf("Bootstrap iteration %d / %d\n", m, num_bootstrap)) flush.console() } boot <- bootstrap_clusters_idx(data) idx_m <- boot$index Y0_boot <- Y0[idx_m, , drop = FALSE] Y1_boot <- Y1[idx_m, , drop = FALSE] deltas_m <- if (estimator == "gform") { compute_ATEs_gform(boot$dat, Y0_boot, Y1_boot, subgroup_var = subgroup_var, subgroup_val = subgroup_val) } else { compute_ATEs_mr(boot$dat, Y0_boot, Y1_boot, subgroup_var = subgroup_var, subgroup_val = subgroup_val) } V_c[m, ] <- deltas_m$c_ATE V_i[m, ] <- deltas_m$i_ATE } row_means_c <- apply(V_c, 1, mean, na.rm = TRUE) row_means_i <- apply(V_i, 1, mean, na.rm = TRUE) Var_c <- var(row_means_c, na.rm = TRUE) + var(posterior_ATEs$c_ATE, na.rm = TRUE) Var_i <- var(row_means_i, na.rm = TRUE) + var(posterior_ATEs$i_ATE, na.rm = TRUE) list(SE_c = sqrt(Var_c), SE_i = sqrt(Var_i), Var_c = Var_c, Var_i = Var_i) } # --------------------------------------------------------------------------- # 8. format_results() # --------------------------------------------------------------------------- # Goal: # Take posterior draws and the calibrated SEs and assemble a two-row # summary table showing Uncalibrated (posterior quantile) and Calibrated # (cluster-bootstrap-corrected) 95% credible intervals. # # Arguments: # posterior_ATEs : data.frame from compute_ATEs_gform() or # compute_ATEs_mr() (B rows, i_ATE and c_ATE columns). # cal : list from calibrated_variance() (must have SE_c, SE_i). # digits : number of digits to round point estimates and CIs to. # # Output: # data.frame with columns: # Posterior Inference, h-cATE Estimate, h-cATE 95% C.I., # h-iATE Estimate, h-iATE 95% C.I. # --------------------------------------------------------------------------- format_results <- function(posterior_ATEs, cal, digits = 3) { est_c <- mean(posterior_ATEs$c_ATE, na.rm = TRUE) est_i <- mean(posterior_ATEs$i_ATE, na.rm = TRUE) ci_c_uncal <- quantile(posterior_ATEs$c_ATE, c(.025, .975), na.rm = TRUE) ci_i_uncal <- quantile(posterior_ATEs$i_ATE, c(.025, .975), na.rm = TRUE) ci_c_cal <- est_c + c(-1.96, 1.96) * cal$SE_c ci_i_cal <- est_i + c(-1.96, 1.96) * cal$SE_i fmt_ci <- function(lo, hi) sprintf(paste0("(%.", digits, "f, %.", digits, "f)"), lo, hi) data.frame( `Posterior Inference` = c("Uncalibrated", "Calibrated"), `h-cATE Estimate` = c(round(est_c, digits), round(est_c, digits)), `h-cATE 95% C.I.` = c(fmt_ci(ci_c_uncal[1], ci_c_uncal[2]), fmt_ci(ci_c_cal[1], ci_c_cal[2])), `h-iATE Estimate` = c(round(est_i, digits), round(est_i, digits)), `h-iATE 95% C.I.` = c(fmt_ci(ci_i_uncal[1], ci_i_uncal[2]), fmt_ci(ci_i_cal[1], ci_i_cal[2])), check.names = FALSE ) } # --------------------------------------------------------------------------- # 9. run_swcrt_analysis() -- top-level wrapper # --------------------------------------------------------------------------- # Goal: # End-to-end driver: prepare data, fit mxBART, compute posterior ATEs, # run cluster bootstrap for calibrated variance, and return both the # formatted summary table and the underlying posterior objects (for # diagnostics or further analysis). # # Arguments: # data : raw analysis data.frame (already cleaned, with the # outcome on its final scale). # cluster_col, # period_col, # treat_col, # outcome_col : column names of the four design variables. # covariate_formula : one-sided formula for adjustment covariates. # estimator : "gform" or "mr". # re_structure : "simple_exch" or "nested_exch". # subgroup_var : NULL for overall, or column name to restrict to. # subgroup_val : level of subgroup_var defining the subgroup. # mctotal, mcburnin, ntree, thin_rate : mxBART MCMC settings. # num_bootstrap : number of cluster-bootstrap resamples. # seed : integer seed. # verbose : print bootstrap progress. # # Output: # List with: # summary - data.frame, two-row Uncalibrated/Calibrated table # posterior_ATEs - data.frame with B rows of (i_ATE, c_ATE) draws # Y0, Y1 - n x B posterior prediction matrices # fit - raw mxBART fit (for diagnostics) # calibration - list returned by calibrated_variance() # --------------------------------------------------------------------------- run_swcrt_analysis <- function(data, cluster_col, period_col, treat_col, outcome_col, covariate_formula, estimator = "gform", re_structure = "simple_exch", subgroup_var = NULL, subgroup_val = NULL, mctotal = 4000, mcburnin = 2000, ntree = 50L, thin_rate = 1, num_bootstrap = 100, seed = 123, verbose = TRUE) { if (!estimator %in% c("gform", "mr")) { stop('estimator must be "gform" or "mr".') } if (!re_structure %in% c("simple_exch", "nested_exch")) { stop('re_structure must be "simple_exch" or "nested_exch".') } # 1. Prepare data dat <- prepare_data(data, cluster_col = cluster_col, period_col = period_col, treat_col = treat_col, outcome_col = outcome_col) # 2. Fit mxBART fit_obj <- fit_mxbart_ate(dat, covariate_formula = covariate_formula, re_structure = re_structure, mctotal = mctotal, mcburnin = mcburnin, ntree = ntree, thin_rate = thin_rate, seed = seed) # 3. Posterior ATE draws (overall or subgroup, gform or mr) posterior_ATEs <- if (estimator == "gform") { compute_ATEs_gform(dat, fit_obj$Y0, fit_obj$Y1, subgroup_var = subgroup_var, subgroup_val = subgroup_val) } else { compute_ATEs_mr(dat, fit_obj$Y0, fit_obj$Y1, subgroup_var = subgroup_var, subgroup_val = subgroup_val) } # 4. Calibrated variance via cluster bootstrap cal <- calibrated_variance(dat, fit_obj$Y0, fit_obj$Y1, posterior_ATEs, estimator = estimator, subgroup_var = subgroup_var, subgroup_val = subgroup_val, num_bootstrap = num_bootstrap, verbose = verbose) # 5. Format results table tab <- format_results(posterior_ATEs, cal) list(summary = tab, posterior_ATEs = posterior_ATEs, Y0 = fit_obj$Y0, Y1 = fit_obj$Y1, fit = fit_obj$fit, calibration = cal) }