This tutorial demonstrates the workflow implemented in
swcrt_bart_analysis.R for calibrated Bayesian analysis of
stepped-wedge cluster-randomized trials (SW-CRTs) using mixed-effects
Bayesian additive regression trees (mxBART) as the working outcome
model.
The pipeline supports:
gform) and model-robust standardization
(mr);The pipeline is organized into nine functions chained by a top-level
wrapper run_swcrt_analysis(). The functions are summarized
in Section 4; full argument-by-argument documentation lives in the
comment block above each function in swcrt_bart_analysis.R.
Two end-to-end application examples follow in Sections 5 (subgroup
analysis) and 6 (overall analysis).
The pipeline depends on the mxBART package for the
working outcome model. The other R packages (stats,
utils) are part of base R.
# install.packages("mxBART") # if not already installed
library(mxBART)Source the analysis functions. Adjust the path to point to where
swcrt_bart_analysis.R is saved.
source("swcrt_bart_analysis.R")After sourcing, the following nine functions are loaded into the workspace:
prepare_data(), make_design_matrix(),
bootstrap_clusters_idx(),
compute_ATEs_gform(), compute_ATEs_mr(),
fit_mxbart_ate(), calibrated_variance(),
format_results(), run_swcrt_analysis().
To illustrate the pipeline, we simulate a small SW-CRT with 30 clusters observed over 4 periods. Each cluster crosses over from control to intervention at one of three staggered steps, with 10 clusters crossing per step (a balanced rollout).
The dataset uses generic variable names (cluster_id,
period, intervention, x1,
x2, y) to keep the example
application-agnostic. In real SW-CRTs these variables map to substantive
quantities. For example, in a healthcare quality-improvement trial:
cluster_id could index participating hospitals or
clinics;period could index calendar quarters or months over the
rollout window;intervention could indicate whether a clinical pathway
(e.g., a sepsis bundle, a discharge protocol, a palliative-care program)
has been rolled out at that cluster in that period;x1 could be a binary patient characteristic such as a
comorbidity indicator (presence of dementia, ESRD, COPD), or a binary
risk classification;x2 could be a standardized continuous baseline
covariate such as a frailty score, a severity index, or age centered at
the cohort mean;y could be a continuous patient-level outcome such as
log length of stay, a symptom-scale score, or a quality-of-life
measure.Subgroup analyses are typically motivated by clinically meaningful
binary characteristics, which is why we use a binary x1
here as the subgroup variable.
set.seed(2026)
I <- 30 # number of clusters
J <- 4 # number of periods
stopifnot(I %% (J - 1) == 0)
# Rollout schedule: 10 clusters cross over at each step (periods 2, 3, 4)
rollout <- rep(2:J, each = I / (J - 1))
# Period-varying coefficients (secular trends, covariate effects)
beta0j <- 0.25 + (0.02 / (J - 1)) * (0:(J - 1)) # intercepts
beta1j <- (3 * (1:J)) / 2 # x1 coefficient by period
beta2j <- (1:J) / J # x2 coefficient by period
# Variance components
sig2_b <- 0.05 # cluster random intercept variance
sig2_cp <- 0.05 # cluster x period random intercept variance
sig2_e <- 0.90 # residual variance
alpha_i <- rnorm(I, mean = 0, sd = sqrt(sig2_b)) # cluster random effects
# Sample sizes per cluster-period (grow slightly across periods)
N_mat <- matrix(NA_integer_, nrow = I, ncol = J)
for (j in 1:J) {
lo <- 5 + 5 * j
hi <- 45 + 5 * j
N_mat[, j] <- sample(lo:hi, size = I, replace = TRUE)
}
global_mean_N <- mean(N_mat)
EN_i1 <- (10 + 50) / 2
# Build long-format dataset
dat_list <- vector("list", I * J)
idx <- 0
for (i in 1:I) {
for (j in 1:J) {
idx <- idx + 1
z_ij <- as.numeric(j >= rollout[i])
delta_ij <- rnorm(1, mean = 0, sd = sqrt(sig2_cp))
n_ij <- N_mat[i, j]
x1 <- rbinom(n_ij, size = 1, prob = 0.5)
x2 <- rnorm(n_ij, mean = 0, sd = 0.1)
mu_fixed <- beta0j[j] +
beta1j[j] * x1 +
beta2j[j] * (x2^2)
# Heterogeneous treatment effect (varies by period, x1, x2,
# and the cluster-period sample size)
theta_ijk <- 1 +
j * sin(x1) -
(j^2) * exp(-x2) +
sqrt(n_ij) / global_mean_N +
3 * log(n_ij) * (n_ij^2) / (EN_i1^2)
eps <- rnorm(n_ij, mean = 0, sd = sqrt(sig2_e))
Y_0_true <- mu_fixed + alpha_i[i] + delta_ij + eps
Y_1_true <- Y_0_true + theta_ijk
y <- z_ij * Y_1_true + (1 - z_ij) * Y_0_true
dat_list[[idx]] <- data.frame(
cluster_id = i,
period = j,
intervention = z_ij,
x1 = x1,
x2 = x2,
y = y
)
}
}
sim_dat <- do.call(rbind, dat_list)cat("Dataset:", nrow(sim_dat), "rows x", ncol(sim_dat), "cols\n")#> Dataset: 4590 rows x 6 colscat("Subgroup (x1 = 1) prevalence:", round(mean(sim_dat$x1), 3), "\n")#> Subgroup (x1 = 1) prevalence: 0.496# Rollout schedule (1 = cluster is treated in that period)
with(sim_dat, tapply(intervention, list(cluster_id, period), max))#> 1 2 3 4
#> 1 0 1 1 1
#> 2 0 1 1 1
#> 3 0 1 1 1
#> 4 0 1 1 1
#> 5 0 1 1 1
#> 6 0 1 1 1
#> 7 0 1 1 1
#> 8 0 1 1 1
#> 9 0 1 1 1
#> 10 0 1 1 1
#> 11 0 0 1 1
#> 12 0 0 1 1
#> 13 0 0 1 1
#> 14 0 0 1 1
#> 15 0 0 1 1
#> 16 0 0 1 1
#> 17 0 0 1 1
#> 18 0 0 1 1
#> 19 0 0 1 1
#> 20 0 0 1 1
#> 21 0 0 0 1
#> 22 0 0 0 1
#> 23 0 0 0 1
#> 24 0 0 0 1
#> 25 0 0 0 1
#> 26 0 0 0 1
#> 27 0 0 0 1
#> 28 0 0 0 1
#> 29 0 0 0 1
#> 30 0 0 0 1The rollout table shows the staggered design: clusters 1–10 are treated from period 2 onward, clusters 11–20 from period 3, and clusters 21–30 from period 4.
The pipeline estimates two weighted average treatment effects:
The two estimands coincide when cluster-period sizes are non-informative, but can differ meaningfully when cluster-period sizes are informative. The pipeline returns posterior draws of both quantities from a single mxBART fit; the weighting scheme is applied at the aggregation step.
| Function | Purpose |
|---|---|
prepare_data() |
Rename design columns; build helper cols (period_f,
cp, N_ij). |
make_design_matrix() |
Build fixed-effect design matrix for mxBART from a covariate formula. |
bootstrap_clusters_idx() |
Resample whole clusters and return row-index mapping for
Y0/Y1. |
compute_ATEs_gform() |
Posterior draws of h-iATE and h-cATE under model-based g-computation (overall or subgroup). |
compute_ATEs_mr() |
Posterior draws of h-iATE and h-cATE under model-robust standardization (overall or subgroup). |
fit_mxbart_ate() |
Fit mxBART and return n x B counterfactual prediction matrices
Y0, Y1. |
calibrated_variance() |
Cluster-bootstrap-corrected variance of posterior mean ATEs. |
format_results() |
Two-row uncalibrated/calibrated summary table for h-iATE and h-cATE. |
run_swcrt_analysis() |
Top-level wrapper that chains the above. |
Full argument-by-argument documentation is in the comment block above
each function in swcrt_bart_analysis.R.
run_swcrt_analysis()The top-level wrapper takes the following arguments:
| Argument | Description |
|---|---|
data |
Cleaned data.frame, one row per individual, with the outcome on its final analysis scale. |
cluster_col |
String. Column name in data for the cluster ID. |
period_col |
String. Column name for the period (integer 1..J). |
treat_col |
String. Column name for the treatment indicator (0/1). |
outcome_col |
String. Column name for the outcome. |
covariate_formula |
One-sided formula (e.g., ~ age + factor(race)) listing
adjustment covariates. Pass ~ 1 or NULL for no
covariate adjustment. |
estimator |
"gform" (model-based g-computation) or
"mr" (model-robust standardization). Default
"gform". |
re_structure |
"simple_exch" (cluster RE only) or
"nested_exch" (cluster + cluster-period REs). Default
"simple_exch". |
subgroup_var |
Column name to restrict to for subgroup analysis, or
NULL for overall analysis. |
subgroup_val |
Value of subgroup_var defining the subgroup. Required
when subgroup_var is non-NULL. |
mctotal |
Total MCMC iterations. Default 4000. |
mcburnin |
Number of burn-in iterations to discard. Default 2000. |
ntree |
Number of trees in the BART ensemble. Default 50. |
thin_rate |
MCMC thinning rate (keepevery argument to
mxbart). Default 1 (no thinning). |
num_bootstrap |
Number of cluster-bootstrap resamples for calibrated variance. Default 100. |
seed |
Integer seed for reproducibility. Default 123. |
verbose |
If TRUE, print bootstrap progress every 10
iterations. |
The function returns a list with elements summary
(uncalibrated/calibrated table), posterior_ATEs (posterior
draws of h-iATE and h-cATE), Y0, Y1 (n x B
counterfactual prediction matrices), fit (raw mxBART fit),
and calibration (calibrated variance components).
Arguments for the lower-level functions follow the same vocabulary;
see the comment block above each function in
swcrt_bart_analysis.R for full details.
This example estimates the treatment effect within the subgroup
defined by x1 == 1, using the model-robust
standardization estimator with the
nested-exchangeable random-effects structure. The full
mxBART pipeline runs end-to-end via
run_swcrt_analysis().
res_subgroup <- run_swcrt_analysis(
data = sim_dat,
cluster_col = "cluster_id",
period_col = "period",
treat_col = "intervention",
outcome_col = "y",
covariate_formula = ~ x1 + x2,
estimator = "mr",
re_structure = "nested_exch",
subgroup_var = "x1",
subgroup_val = 1,
mctotal = 2000,
mcburnin = 1000,
num_bootstrap = 100,
ntree = 50,
seed = 123,
verbose = TRUE
)#> *****Calling gbmm: type=1
#> *****Mixed BART Settings
#> *****Number of levels: 3
#> **Level: 2
#> Units: 30
#> Random Effects (per unit): 1
#> prior,df 1 3
#> scale: 1
#> **Level: 3
#> Units: 120
#> Random Effects (per unit): 1
#> prior,df 1 3
#> scale: 1
#> *****Data:
#> data:n,p,np: 4590, 8, 9180
#> y1,yn: -11.622991, 23.636970
#> x1,x[n*p]: 1.000000, 61.000000
#> xp1,xp[np*p]: 1.000000, 61.000000
#> *****Number of Trees: 50
#> *****Number of Cut Points: 1 ... 46
#> *****burn,nd,thin: 1000,1000,1
#> *****Prior:beta,alpha,tau,nu,lambda,offset: 2,0.95,2.19911,3,460.167,13.5752
#> *****sigma: 14.860383
#> *****Dirichlet:sparse,theta,omega,a,b,rho,augment: 0,0,1,0.5,1,8,0
#> *****printevery: 100
#>
#> MCMC
#> done 0 (out of 2000)
#> done 100 (out of 2000)
#> done 200 (out of 2000)
#> done 300 (out of 2000)
#> done 400 (out of 2000)
#> done 500 (out of 2000)
#> done 600 (out of 2000)
#> done 700 (out of 2000)
#> done 800 (out of 2000)
#> done 900 (out of 2000)
#> done 1000 (out of 2000)
#> done 1100 (out of 2000)
#> done 1200 (out of 2000)
#> done 1300 (out of 2000)
#> done 1400 (out of 2000)
#> done 1500 (out of 2000)
#> done 1600 (out of 2000)
#> done 1700 (out of 2000)
#> done 1800 (out of 2000)
#> done 1900 (out of 2000)
#> time: 8s
#> trcnt,tecnt: 1000,1000
#> Bootstrap iteration 10 / 100
#> Bootstrap iteration 20 / 100
#> Bootstrap iteration 30 / 100
#> Bootstrap iteration 40 / 100
#> Bootstrap iteration 50 / 100
#> Bootstrap iteration 60 / 100
#> Bootstrap iteration 70 / 100
#> Bootstrap iteration 80 / 100
#> Bootstrap iteration 90 / 100
#> Bootstrap iteration 100 / 100The two-row table reports point estimates and 95% credible intervals
(both uncalibrated and calibrated) for h-cATE and h-iATE within the
subgroup x1 == 1. The point estimate is the mean of the
posterior draws of the corresponding estimand.
print(res_subgroup$summary, row.names = FALSE)#> Posterior Inference h-cATE Estimate h-cATE 95% C.I. h-iATE Estimate
#> Uncalibrated 20.567 (20.339, 20.853) 21.904
#> Calibrated 20.567 (17.196, 23.938) 21.904
#> h-iATE 95% C.I.
#> (21.693, 22.158)
#> (18.779, 25.030)The uncalibrated 95% CIs are the 2.5th and 97.5th percentiles of the posterior draws and reflect parameter uncertainty only. The calibrated 95% CIs are constructed as the point estimate ± 1.96 × calibrated SE, where the calibrated SE adds the cluster-bootstrap variance of the posterior mean to the within-data posterior variance. Calibration is especially important for the model-robust estimator, because its formulation depends jointly on the model parameters and on aggregated observed cluster-period outcomes. Posterior parameter uncertainty alone does not capture the sampling variability of the observed outcomes, so naive uncalibrated intervals can substantially under-cover the true estimand. The calibrated variance estimator uses cluster-level bootstrap that captures this additional variability and restores valid frequentist coverage.
The returned object also contains the raw posterior draws and the
calibrated variance components. In posterior_ATEs, the
column i_ATE stores the h-iATE draws and the column
c_ATE stores the h-cATE draws.
summary(res_subgroup$posterior_ATEs)#> i_ATE c_ATE
#> Min. :21.59 Min. :20.25
#> 1st Qu.:21.79 1st Qu.:20.44
#> Median :21.91 Median :20.57
#> Mean :21.90 Mean :20.57
#> 3rd Qu.:22.00 3rd Qu.:20.67
#> Max. :22.30 Max. :20.99cat("Calibrated SE for h-cATE:", round(res_subgroup$calibration$SE_c, 4), "\n")#> Calibrated SE for h-cATE: 1.7199cat("Calibrated SE for h-iATE:", round(res_subgroup$calibration$SE_i, 4), "\n")#> Calibrated SE for h-iATE: 1.5945The full mxBART fit object is available as
res_subgroup$fit if posterior diagnostics such as
traceplots or Geweke statistics are desired, and the n x B posterior
prediction matrices are available as res_subgroup$Y0 and
res_subgroup$Y1.
This example estimates the overall (population-level) treatment
effect using the model-based g-computation estimator
with the simple-exchangeable random-effects structure.
To run an overall analysis, simply omit the subgroup_var
and subgroup_val arguments from
run_swcrt_analysis(); the pipeline then aggregates over the
full sample rather than restricting to a subgroup.
res_overall <- run_swcrt_analysis(
data = sim_dat,
cluster_col = "cluster_id",
period_col = "period",
treat_col = "intervention",
outcome_col = "y",
covariate_formula = ~ x1 + x2,
estimator = "gform",
re_structure = "simple_exch",
# subgroup_var and subgroup_val omitted -> overall analysis
mctotal = 2000,
mcburnin = 1000,
num_bootstrap = 100,
ntree = 50,
seed = 123,
verbose = TRUE
)#> *****Calling gbmm: type=1
#> *****Mixed BART Settings
#> *****Number of levels: 2
#> **Level: 2
#> Units: 30
#> Random Effects (per unit): 1
#> prior,df 1 3
#> scale: 1
#> *****Data:
#> data:n,p,np: 4590, 8, 9180
#> y1,yn: -11.622991, 23.636970
#> x1,x[n*p]: 1.000000, 61.000000
#> xp1,xp[np*p]: 1.000000, 61.000000
#> *****Number of Trees: 50
#> *****Number of Cut Points: 1 ... 46
#> *****burn,nd,thin: 1000,1000,1
#> *****Prior:beta,alpha,tau,nu,lambda,offset: 2,0.95,2.19911,3,460.167,13.5752
#> *****sigma: 14.860383
#> *****Dirichlet:sparse,theta,omega,a,b,rho,augment: 0,0,1,0.5,1,8,0
#> *****printevery: 100
#>
#> MCMC
#> done 0 (out of 2000)
#> done 100 (out of 2000)
#> done 200 (out of 2000)
#> done 300 (out of 2000)
#> done 400 (out of 2000)
#> done 500 (out of 2000)
#> done 600 (out of 2000)
#> done 700 (out of 2000)
#> done 800 (out of 2000)
#> done 900 (out of 2000)
#> done 1000 (out of 2000)
#> done 1100 (out of 2000)
#> done 1200 (out of 2000)
#> done 1300 (out of 2000)
#> done 1400 (out of 2000)
#> done 1500 (out of 2000)
#> done 1600 (out of 2000)
#> done 1700 (out of 2000)
#> done 1800 (out of 2000)
#> done 1900 (out of 2000)
#> time: 8s
#> trcnt,tecnt: 1000,1000
#> Bootstrap iteration 10 / 100
#> Bootstrap iteration 20 / 100
#> Bootstrap iteration 30 / 100
#> Bootstrap iteration 40 / 100
#> Bootstrap iteration 50 / 100
#> Bootstrap iteration 60 / 100
#> Bootstrap iteration 70 / 100
#> Bootstrap iteration 80 / 100
#> Bootstrap iteration 90 / 100
#> Bootstrap iteration 100 / 100print(res_overall$summary, row.names = FALSE)#> Posterior Inference h-cATE Estimate h-cATE 95% C.I. h-iATE Estimate
#> Uncalibrated 19.032 (18.765, 19.247) 20.623
#> Calibrated 19.032 (16.259, 21.806) 20.623
#> h-iATE 95% C.I.
#> (20.367, 20.828)
#> (18.262, 22.983)The table reports the overall h-cATE and h-iATE under model-based g-computation. As in the subgroup example, the point estimate is the posterior mean and the two rows give the uncalibrated and calibrated 95% credible intervals.
summary(res_overall$posterior_ATEs)#> i_ATE c_ATE
#> Min. :20.23 Min. :18.62
#> 1st Qu.:20.55 1st Qu.:18.95
#> Median :20.63 Median :19.04
#> Mean :20.62 Mean :19.03
#> 3rd Qu.:20.71 3rd Qu.:19.12
#> Max. :20.95 Max. :19.36cat("Calibrated SE for h-cATE:", round(res_overall$calibration$SE_c, 4), "\n")#> Calibrated SE for h-cATE: 1.4152cat("Calibrated SE for h-iATE:", round(res_overall$calibration$SE_i, 4), "\n")#> Calibrated SE for h-iATE: 1.2042By switching the estimator and re_structure
arguments, the same pipeline returns estimates under the other three
combinations (mr + simple_exch,
gform + nested_exch, mr +
nested_exch). Adding subgroup_var and
subgroup_val turns any of these into a subgroup
analysis.