Heterogeneity of Treatment Effect analysis pipeline

This document runs heterogeneity of treatment effect (HTE) analyses for a randomized trial. Parts I to III are model-based: they estimate how the conditional average treatment effect (CATE) varies with observed baseline covariates. Part IV is randomization-based: it asks how many patients benefit and by how much, characterizing the distribution of individual treatment effects (ITE) without modeling covariates.

Part	Question	Method	Key output
I. Fit-the-fit	Which covariate-defined subgroups have different average treatment effects?	Stage 1: BART estimates the CATE for every patient. Stage 2: a CART tree is fitted to those estimated CATEs for subgroup discovery.	Per-patient CATE distributions, MCMC diagnostics, CART subgroup tree, variable importance and interaction heatmap.
II. Sensitivity and method comparison	Is the Part I CATE estimation robust to the BART prior specification and to the choice of estimation method?	BART prior-sensitivity grid evaluation; cross-method comparisons against BCF, causal forest, and XGBoost T- and S-learners	Prior-grid CATE summaries, cross-method CATE distribution overlays.
III. HTE by prognosis score	Does the treatment effect vary with baseline prognosis?	Prognostic score estimation (three approaches), quartile stratification with observed effects, CATE modeled as a function of the prognostic score	Effect estimates by prognostic stratum, CATE-by-prognosis curves, and CART partitioning on prognostic scores.
IV. ITE quantiles	What fraction of patients benefit, and by how much?	Randomization-based inference for quantiles of the individual treatment effects	Simultaneous confidence bands for effect quantiles, confidence bounds on the proportion of patients benefiting.

---

Setup and Configuration

Package installation

## ---- CRAN packages ----
install.packages(c("BART", "dbarts", "bartMan", "caret", "rpart", "rpart.plot",
                   "flextable", "officer", "ggplot2", "dplyr", "tidyr", "purrr",
                   "ggh4x", "readxl", "gridExtra", "scales", "bcf", "xgboost",
                   "grf", "extraDistr", "devtools"))

## ---- RIQITE: Randomization-inference engine for Part IV ----
devtools::install_github("li-xinran/RIQITE")

Source files

File	Role	Acquisition
clusterfunctions.R	Parallel processing wrappers for Bayesian Additive Regression Trees (BART).	Local directory.
helper_functions.R	Helper functions supporting the randomization-based methodology in Part IV.	Hosted on https://github.com/Zhe-Chen-1999/Enhanced_inference_for_ITE_quantiles. You can manually clone the repository or execute the code chunk below to download it automatically.

download.file(
  paste0("https://raw.githubusercontent.com/Zhe-Chen-1999/",
         "Enhanced_inference_for_ITE_quantiles/main/helper_functions.R"),
  destfile = "helper_functions.R")

## ---- Parts I, II, III: Model-based HTE ----
library(BART)
library(dbarts)
library(bartMan)   # VIVI-VSUP heatmaps
library(caret)
library(rpart)
library(rpart.plot)
library(flextable)
library(officer)
library(ggplot2)
library(dplyr)
library(tidyr)
library(purrr)
library(ggh4x)
library(readxl)
library(gridExtra)
library(scales)
library(bcf)
library(xgboost)
library(grf)

source("clusterfunctions.R")    # lbart.cluster(), wbart.cluster()

## ---- Part IV: randomization-based ITE-quantile inference ----
library(RIQITE)      # called from within helper_functions.R
library(extraDistr)  # required by helper_functions.R

source("helper_functions.R")

Making re-knits fast

A full run is dominated by a handful of expensive chunks, including BART hyperparameter tuning, prior-sensitivity grids, and Bayesian Causal Forest (BCF) estimation. Re-running all of them just to fix a typo in the prose would be painful, so every expensive result is written to an .rds file under cache/ the first time it is computed and simply reloaded afterwards.

The mechanism is one helper, cached(), defined in the configuration chunk:

step1.res <- cached("step1_bart_binary", {   # <- name of the .rds file
  bart_fit_binary_outcome(...)               # <- only runs if the file is absent
})

What this means in practice.

You want to…	Do this
Knit for the first time	Nothing. Every expensive chunk computes and saves. Budget an hour or more.
Edit prose, a caption, or a plot, then re-knit	Nothing. Cached results reload in seconds.
Force one analysis to recompute	Delete its file, e.g. cache/step1_bart_binary.rds, or pass refresh = TRUE to that cached() call.
Force everything to recompute	Set REFRESH_ALL <- TRUE in the configuration chunk, or delete the whole cache/ directory.
Change the data or a seed	Set REFRESH_ALL <- TRUE for one knit, then set it back. Otherwise you will silently keep results computed under the old settings.

---

Global Configuration

## ============================ EDIT THIS BLOCK =============================
## ---- Outcome type ----
## Governs Parts I-III: Determines the appropriate BART likelihood (logit vs. Gaussian), XGBoost 
## objective function, prognostic score family, and output summaries. 
## Risk ratio outputs are restricted to binary outcomes.
OUTCOME_TYPE <- "binary"       # "binary" or "continuous"

## ---- Execution & Caching Parameters ----
GLOBAL_SEED <- 123
CACHE_DIR   <- "cache"   # expensive results are cached here

## Set to TRUE to recompute every cached result on this knit. 
REFRESH_ALL <- FALSE

## ---- Which Parts to run ----
RUN_PRIOR_GRID    <- TRUE   # Part II: BART prior-sensitivity grid
RUN_CROSS_METHOD  <- TRUE   # Part II: Cross-method comparison (BCF, Causal Forest, XGBoost)
RUN_BASELINE_RISK <- TRUE   # Part III: Baseline-risk heterogeneity analysis
## =========================================================================

OUTCOME_TYPE <- match.arg(OUTCOME_TYPE, c("binary", "continuous"))
BINARY <- OUTCOME_TYPE == "binary"

## Label for the additive effect scale in plot titles and summary tables.
EFFECT_LABEL <- if (BINARY) "Risk Difference" else "Mean Difference"

set.seed(GLOBAL_SEED)
if (!dir.exists(CACHE_DIR)) dir.create(CACHE_DIR, showWarnings = FALSE)

#' Compute an expensive result once, then reload it on later knits.
#'
#' Writes evaluation of 'expr' to CACHE_DIR/<name>.rds.
#' Delete that file, pass refresh = TRUE, or set REFRESH_ALL <- TRUE to recompute.
.cache_log <- new.env(parent = emptyenv())

cached <- function(name, expr, refresh = FALSE) {
  # # Build the cache filename
  path <- file.path(CACHE_DIR, paste0(name, ".rds")) 
  
  # If the result already exists and you did not request or globally force recomputation, reload the result.
  if (!refresh && !REFRESH_ALL && file.exists(path)) {
    assign(name, list(status = "reloaded", secs = 0), envir = .cache_log) # Record that it was reloaded
    return(readRDS(path))
  }
  
  # Otherwise compute
  t0 <- Sys.time()
  value <- force(expr) # evaluates the expression
  saveRDS(value, path)
  assign(name, list(status = "computed",
                    secs = as.numeric(difftime(Sys.time(), t0, units = "secs"))), # record runtime
         envir = .cache_log)
  value # return the result
}

#' Report what this knit computed versus reloaded (printed in the appendix).
cache_report <- function() {
  nms <- ls(.cache_log)
  if (!length(nms)) return(data.frame())
  data.frame(
    Result   = nms,
    Status   = vapply(nms, function(n) get(n, .cache_log)$status, character(1)), # "computed" or "reloaded"
    Seconds  = round(vapply(nms, function(n) get(n, .cache_log)$secs, numeric(1)), 1), # computation time
    row.names = NULL)
}

---

Part I. A “fit-the-fit” approach

1.1 BART estimation of conditional average treatment effects

By default, Bayesian Additive Regression Trees (BART) utilize an ensemble of 200 trees, anchored by a base prior of 0.95 and a power prior of 2.0. These priors regularize the model by penalizing deep tree structures, thereby limiting complex interactions and mitigating overfitting. These defaults work well in general but are not necessarily optimal for a given dataset, so the hyperparameters were evaluated by 10-fold cross-validation across a grid of 27 configurations: power \(\in\) {1, 2, 3}, base \(\in\) {0.25, 0.5, 0.95}, and number of trees \(\in\) {50, 200, 400}. The configuration minimizing cross-validated mean squared error (MSE) is selected for the final estimation.

Posterior inference is achieved via Markov Chain Monte Carlo (MCMC) sampling with nchains independent chains, nskip burn-in iterations, and ndpost retained draws.

Binary outcome Model (logit BART)

bart_fit_binary_outcome <- function(X, y, z, seed = 123, nfolds = 10,
                          Power = 1:3, Base = c(0.25, 0.5, 0.95), Ntrees = c(50, 200, 400),
                          k = 2.0,
                          nchains = 4, keepevery = 1, ndpost=1000L, nskip = 500){

  # Predictors: combine treatment indicator z with covariates X
  x <- cbind(z, X)

  # Create datasets under counterfactual treatment scenarios (z = TRUE and z = FALSE)
  x1 <- x0 <- x
  x1$z <- TRUE
  x0$z <- FALSE

  best_hp <- NULL
  cvoutput <- NULL

  ## Cross-validation step for hyperparameter tuning (Power, Base, Ntrees)
  if (length(Power) * length(Base) * length(Ntrees) > 1) {

    # Create dataset folds for cross-validation
    set.seed(seed) # Ensures the CV folds are generated the same way every time
    folds <- createFolds(y, k = nfolds, list = TRUE, returnTrain = FALSE)

    # Initialize output matrices for prediction error from each model
    cvoutput <- expand.grid(Power, Base, Ntrees, NA)
    colnames(cvoutput) <- c("Power", "Base", "Ntrees", "CVMSE")
    mse <- array(NA, dim = c(nrow(cvoutput), nfolds))

    # Perform cross-validation for hyperparameter tuning  (may take >2 hours)
    for (hp in 1:nrow(cvoutput)) {
      for (i in 1:nfolds) {

        # Fit BART model with specific hyperparameters
        bartmod <- lbart.cluster(x.train = x[-folds[[i]],],
                             y.train = y[-folds[[i]]],
                             x.test = x[folds[[i]],],
                             power = cvoutput$Power[hp],
                             base = cvoutput$Base[hp],
                             ntree = cvoutput$Ntrees[hp],
                             k = k,
                             nchains = nchains,
                             nskip=nskip, ndpost = ndpost,
                             keepevery = keepevery,
                             # Dynamically adjusting the seed prevents identical folds from repeating seeds
                             seed = seed + hp * nfolds + i) # Ensure unique reproducible seeds 

        # Predictions and MSE calculation
        bartmod$yhat.test.collapse <- apply(bartmod$yhat.test, 2, rbind)
        pred <- colMeans(plogis(bartmod$yhat.test.collapse))
        mse[hp, i] <- mean((y[folds[[i]]] - pred)^2)
      }
    }

    # Calculate cross-validation error (CVMSE) for each hyperparameter combination
    cvoutput$CVMSE <- rowMeans(mse)

    # Fit final model with best hyperparameters based on minimum CVMSE
    best_hp <- which.min(cvoutput$CVMSE)
    bartmod1 <- lbart.cluster(x.train = x, y.train = y,
                              x.test = x1, nchains = nchains,
                              power = cvoutput$Power[best_hp],
                              base = cvoutput$Base[best_hp],
                              ntree = cvoutput$Ntrees[best_hp],
                              k = k,
                              nskip=nskip, ndpost = ndpost,
                              keepevery = keepevery,
                              seed = seed) # Explicit seed

    bartmod0 <- lbart.cluster(x.train = x, y.train = y,
                              x.test = x0, nchains = nchains,
                              power = cvoutput$Power[best_hp],
                              base = cvoutput$Base[best_hp],
                              ntree = cvoutput$Ntrees[best_hp],
                              k = k,
                              nskip=nskip, ndpost = ndpost,
                              keepevery = keepevery,
                              seed = seed) # Same seed ensures identical posterior trees

    # Print the optimal hyperparameters for future reference
    cat("Optimal hyperparameters index:", best_hp,"\n")
    cat("Power", "Base", "Ntrees:", "\n")
    cat(cvoutput[best_hp, 1], cvoutput[best_hp, 2], cvoutput[best_hp, 3],"\n")

  }else{
    # Fit BART model directly with pre-specified hyperparameters if no cross-validation needed
    bartmod1 <- lbart.cluster(x.train = x, y.train = y,
                              x.test = x1, nchains = nchains,
                              power = Power,
                              base = Base,
                              ntree = Ntrees,
                              k = k,
                              nskip=nskip, ndpost = ndpost,
                              keepevery = keepevery,
                              seed = seed) # Explicit seed

    bartmod0 <- lbart.cluster(x.train = x, y.train = y,
                              x.test = x0, nchains = nchains,
                              power = Power,
                              base = Base,
                              ntree = Ntrees,
                              k = k,
                              nskip=nskip, ndpost = ndpost,
                              keepevery = keepevery,
                              seed = seed) # Same seed ensures identical posterior trees
  }

  # Collapse predictions across chains for certain calculations
  bartmod1$yhat.test.collapse <- apply(bartmod1$yhat.test, 2, rbind)
  bartmod0$yhat.test.collapse <- apply(bartmod0$yhat.test, 2, rbind)
  bartmod1$yhat.train.collapse <- apply(bartmod1$yhat.train, 2, rbind)
  bartmod0$yhat.train.collapse <- apply(bartmod0$yhat.train, 2, rbind)

  # Transform the fitted outcomes to probabilities using inverse logit transformation
  ## P(Y=1 | Z=1, X)
  prob_treated <- exp(bartmod1$yhat.test.collapse) / (1 +   exp(bartmod1$yhat.test.collapse)) # plogis(bartmod1$yhat.test.collapse)
  ## P(Y=1 | Z=0, X)
  prob_control <- exp(bartmod0$yhat.test.collapse) / (1 + exp(bartmod0$yhat.test.collapse))

  # Calculate the CATE (risk difference) for each individual for each MCMC sample
  cate_posterior <- prob_treated - prob_control

  # Calculate the Risk Ratio (RR) for each individual for each MCMC sample
  epsilon <- 1e-10 # Add small constant to avoid division by zero
  rr_posterior <- prob_treated / (prob_control + epsilon)

  # Calculate equal-tailed 95% credible intervals for CATE
  res <- data.frame(
    # Risk Difference metrics
    cate = apply(cate_posterior, 2, mean), # mean CATE across MCMC samples
    cate_lower = apply(cate_posterior, 2, quantile, probs = 0.025),
    cate_upper = apply(cate_posterior, 2, quantile, probs = 0.975),

    # Risk Ratio metrics
    rr = apply(rr_posterior, 2, mean), # mean RR across all MCMC samples
    rr_lower = apply(rr_posterior, 2, quantile, probs = 0.025),
    rr_upper = apply(rr_posterior, 2, quantile, probs = 0.975),

    # Expected probabilities under each treatment
    prob_treated_mean = apply(prob_treated, 2, mean),
    prob_control_mean = apply(prob_control, 2, mean))

  return(list(cate_results = res,
              bartmod0 = bartmod0,
              bartmod1 = bartmod1,
              # posterior distributions
              cate_posterior = cate_posterior,
              rr_posterior = rr_posterior,
              # counterfactual probability draws (used by fit_bart_fixed() in Part II)
              prob_treated = prob_treated,
              prob_control = prob_control,
              optimal_priors = if (!is.null(best_hp)) cvoutput[best_hp, ] else
                     data.frame(Power = Power, Base = Base, Ntrees = Ntrees, CVMSE = NA)))
}

Continuous outcome Model (standard Gaussian BART)

bart_fit_continuous_outcome <- function(X, y, z, seed = 123, nfolds = 10,
                          Power = 1:3, Base = c(0.25, 0.5, 0.95),
                          Ntrees = c(50, 200, 400), k = 2.0,
                          ndpost=1000L, nskip = 500,
                          nchains = 4, keepevery = 1L) {

  x <- cbind(z, X)
  x1 <- x0 <- x
  x1$z <- TRUE
  x0$z <- FALSE
  best_hp <- NULL
  cvoutput <- NULL

  if (length(Power) * length(Base) * length(Ntrees) > 1) {

    set.seed(seed) 
    folds <- createFolds(y, k = nfolds, list = TRUE, returnTrain = FALSE)
    cvoutput <- expand.grid(Power, Base, Ntrees, NA)
    colnames(cvoutput) <- c("Power", "Base", "Ntrees", "CVMSE")
    mse <- array(NA, dim = c(nrow(cvoutput), nfolds))

    for (hp in 1:nrow(cvoutput)) {
      for (i in 1:nfolds) {
        bartmod <- wbart.cluster(x.train = x[-folds[[i]], ],
                                 y.train = y[-folds[[i]]],
                                 x.test = x[folds[[i]], ],
                                 power = cvoutput$Power[hp],
                                 base = cvoutput$Base[hp],
                                 ntree = cvoutput$Ntrees[hp],
                                 k = k,
                                 nchains = nchains,
                                 ndpost = ndpost,
                                 nskip = nskip, keepevery = keepevery,
                                 seed = seed + hp * nfolds + i) 
        bartmod$yhat.test.collapse <- apply(bartmod$yhat.test, 2, rbind)
        pred <- colMeans(bartmod$yhat.test.collapse)
        mse[hp, i] <- mean((y[folds[[i]]] - pred)^2)
      }
    }

    cvoutput$CVMSE <- rowMeans(mse)
    best_hp <- which.min(cvoutput$CVMSE)

    bartmod1 <- wbart.cluster(x.train = x, y.train = y,
                              x.test = x1, nchains = nchains,
                              power = cvoutput$Power[best_hp],
                              base = cvoutput$Base[best_hp],
                              ntree = cvoutput$Ntrees[best_hp],
                              k = k,
                              nskip = nskip, ndpost = ndpost,
                              keepevery = keepevery,
                              seed = seed) # Explicit seed
    bartmod0 <- wbart.cluster(x.train = x, y.train = y,
                              x.test = x0, nchains = nchains,
                              power = cvoutput$Power[best_hp],
                              base = cvoutput$Base[best_hp],
                              ntree = cvoutput$Ntrees[best_hp],
                              k = k,
                              nskip = nskip, ndpost = ndpost,
                              keepevery = keepevery,
                              seed = seed) # Same seed ensures identical posterior trees

    cat("Optimal hyperparameters index:", best_hp, "\n")
    cat("Power", "Base", "Ntrees:", "\n")
    cat(cvoutput[best_hp, 1], cvoutput[best_hp, 2], cvoutput[best_hp, 3], "\n")
  } else {
    bartmod1 <- wbart.cluster(x.train = x, y.train = y,
                              x.test = x1, nchains = nchains,
                              power = Power, base = Base,
                              ntree = Ntrees, k = k,
                              nskip = nskip, ndpost = ndpost,
                              keepevery = keepevery,
                              seed = seed) # Explicit seed
    bartmod0 <- wbart.cluster(x.train = x, y.train = y,
                              x.test = x0, nchains = nchains,
                              power = Power, base = Base,
                              ntree = Ntrees, k = k,
                              nskip = nskip, ndpost = ndpost,
                              keepevery = keepevery,
                              seed = seed) # Same seed ensures identical posterior trees
    cvoutput <- NULL
    best_hp <- NULL
  }

  bartmod1$yhat.test.collapse <- apply(bartmod1$yhat.test, 2, rbind)
  bartmod0$yhat.test.collapse <- apply(bartmod0$yhat.test, 2, rbind)
  bartmod1$yhat.train.collapse <- apply(bartmod1$yhat.train, 2, rbind)
  bartmod0$yhat.train.collapse <- apply(bartmod0$yhat.train, 2, rbind)

  # CATE posterior: draw-by-draw difference (draws x patients)
  cate_posterior <- bartmod1$yhat.test.collapse - bartmod0$yhat.test.collapse

  res <- data.frame(
    cate = colMeans(cate_posterior),
    cate_lower = apply(cate_posterior, 2, quantile, probs = 0.025),
    cate_upper = apply(cate_posterior, 2, quantile, probs = 0.975)
  )

  return(list(
    cate_results = res,
    bartmod0 = bartmod0,
    bartmod1 = bartmod1,
    cate_posterior = cate_posterior,
    optimal_priors = if (!is.null(best_hp)) cvoutput[best_hp, ] else
                     data.frame(Power = Power, Base = Base, Ntrees = Ntrees, CVMSE = NA)
  ))
}

Per-patient CATE plot

Plots the estimated CATE (solid line) with 95% credible intervals (shaded area) for each individual.

plot_treatment_effects <- function(cate.res, 
                                   binary_outcome = TRUE,
                                   scales = NULL) {
  ## cate.res : data.frame with cate/cate_lower/cate_upper, optionally rr/rr_*
  ## binary_outcome : drives the default label if scales is NULL
  ## scales : character vector specifying panels to draw (e.g., "risk difference", "risk ratio" for a binary outcome, and "mean difference" for a continuous outcome; defaults to what the data supports
 
 # Default scales based on outcome type and available columns
  if (is.null(scales)) {
    scales <- if (binary_outcome) {
      c("risk difference", if ("rr" %in% names(cate.res)) "risk ratio")
    } else {
      "mean difference"
    }
  }
  
  plots <- list()
  
  # Create a custom theme with a larger base font size for better readability
  custom_theme <- theme_bw(base_size = 14) +
    theme(axis.title = element_text(face = "bold", margin = margin(t = 10, r = 10)))
  
  # Absolute Effect Plot (Risk Difference or Mean Difference)
  if (any(c("risk difference", "mean difference") %in% scales)) {
    # Sort by CATE
    cate.res.sorted <- cate.res[order(cate.res$cate, decreasing = TRUE), ]
    cate.res.sorted$index <- 1:nrow(cate.res.sorted)

    active_label <- intersect(scales, c("risk difference", "mean difference"))[1]
    
    plots$rd <- ggplot(cate.res.sorted, aes(x = index, y = cate)) +
      geom_ribbon(aes(ymin = cate_lower, ymax = cate_upper),
                  fill = "grey", alpha = 0.4) +
      geom_line(linewidth = 1) +
      geom_hline(yintercept = 0, lty = 3, color = "red") +
      xlab(paste0("Participant (sorted by ", active_label, ")")) +
      ylab(active_label) +
      custom_theme
  }

  # Risk Ratio Plot (for binary outcomes)
  if ("risk ratio" %in% scales && binary_outcome) {
    # Sort by RR
    cate.res.sorted2 <- cate.res[order(cate.res$rr, decreasing = TRUE), ]
    cate.res.sorted2$index <- 1:nrow(cate.res.sorted2)

    plots$rr <- ggplot(cate.res.sorted2, aes(x = index, y = rr)) +
      geom_ribbon(aes(ymin = rr_lower, ymax = rr_upper),
                  fill = "lightblue", alpha = 0.4) +
      geom_line(linewidth = 1) +
      geom_hline(yintercept = 1, lty = 3, color = "red") +
      xlab("Participant (sorted by risk ratio)") +
      ylab("Risk Ratio") +
      custom_theme +
      # log scale for risk ratios to ensure visual symmetry around the null (1.0).
      # Because ratios are multiplicative, log makes a halving of risk (0.5) and
      # a doubling of risk (2.0) exactly equidistant from 1. 
      scale_y_continuous(trans = "log2", 
                         breaks = c(0.25, 0.5, 1, 2, 4),
                         labels = c("0.25", "0.5", "1", "2", "4"))
  }
  
  # Combine plots
  if (length(plots) == 1) plots[[1]] else do.call(gridExtra::grid.arrange, c(plots, ncol = 1))
}

Example Analysis: The SHAMROC Trial

The illustrative dataset throughout this document is derived from the SHAMROC trial, consisting of 431 patients with 6-month follow-up data (223 in the restrictive fluid group, 208 in the liberal fluid group).

Data Preprocessing

Covariates used: age65, adl_bin (baseline disability), liveathome, charl_congheart (heart failure), kidneypresdial (ESRD), d_sofa_gcs.

Imputation: the binary outcome adlmorelessf has one missing value imputed by the mode; the continuous 6-month scores have one missing value each imputed by the median.

clovers_merge6mo <- read.csv('shamroc_cart.csv')
clovers_merge6mo$age65 <- ifelse(clovers_merge6mo$scr_age >= 65, '>= 65', '<65')
clovers_merge6mo$adl_bin <- ifelse(clovers_merge6mo$adl_cat == '0',
                                   'No Baseline Disability', 'Baseline Disability')
clovers_merge6mo$liveathome <- ifelse(
  clovers_merge6mo$medhx_prehosp %in% c('Home independently',
                                        'Home with family help',
                                        'Home with professional help'),
  'Living at Home', 'Not Living At Home')
clovers_merge6mo$adlmoreless <- ifelse(clovers_merge6mo$adldiff6mo > 0,
                                       'MoreADLs', 'FewerADLs')
clovers_merge6mo$adlmorelessf <- as.factor(clovers_merge6mo$adlmoreless)

df <- clovers_merge6mo %>%
  dplyr::select(rand_trt, age65, adl_bin, liveathome, charl_congheart,
                kidneypresdial, d_sofa_gcs, adlmorelessf,
                d_score_mo_total.Month_6, d_score_mo_tscore.Month_6)
cat("Rows x columns:", dim(df), "\n")

## Rows x columns: 431 10

cat("Complete cases:", nrow(df[complete.cases(df), ]), "\n")

## Complete cases: 430

# Impute NA values in a character column with the mode
impute_mode <- function(column) {
  mode_value <- names(which.max(table(column, useNA = "no")))
  column[is.na(column)] <- mode_value
  return(column)
}
df$adlmorelessf <- impute_mode(df$adlmorelessf)


# Replace NA values with the median of each numerical column
impute_with_median <- function(x) {
  apply(x, 2, function(col) {
    col[is.na(col)] <- median(col, na.rm = TRUE)
    return(col)
  })
}
df[, c("d_score_mo_total.Month_6", "d_score_mo_tscore.Month_6")] <-
  as.data.frame(impute_with_median(
    df[, c("d_score_mo_total.Month_6", "d_score_mo_tscore.Month_6")]))

# Standardize continuous covariates
d_sofa_gcs_scale <- scale(df$d_sofa_gcs)

# Convert categorical covariates into factors
X_categorical <- df %>%
  select(age65, adl_bin, liveathome, charl_congheart, kidneypresdial)
X_categorical[] <- lapply(X_categorical, as.factor)

## ---- Analysis variables used by the HTE pipeline ----
z <- ifelse(df$rand_trt == "Restrictive Fluid Group", TRUE, FALSE)  # binary treatment
y <- ifelse(df$adlmorelessf == "MoreADLs", 1, 0)   # binary outcome, 1 = worse
y_cont <- df$d_score_mo_total.Month_6              # continuous outcome 
X_unscaled <- cbind(df$d_sofa_gcs, X_categorical)  # covariates, original scale
X <- cbind(d_sofa_gcs_scale, X_categorical)        # standardized covariates

head(X)

##   d_sofa_gcs_scale age65                adl_bin         liveathome
## 1       -0.7965854 >= 65    Baseline Disability     Living at Home
## 2        0.5188487   <65 No Baseline Disability     Living at Home
## 3        0.5188487 >= 65    Baseline Disability Not Living At Home
## 4       -0.7965854 >= 65 No Baseline Disability     Living at Home
## 5        1.3958048   <65 No Baseline Disability     Living at Home
## 6       -0.7965854   <65 No Baseline Disability     Living at Home
##   charl_congheart kidneypresdial
## 1              No             No
## 2              No             No
## 3              No             No
## 4              No             No
## 5              No             No
## 6              No             No

Direction of benefit. y = 1 means more ADL limitations (a negative clinical outcome), so a negative risk difference corresponds to a treatment benefit. y_cont is a function score where higher is better, so a positive mean difference corresponds to a treatment benefit.

BART Execution: estimate CATE for each participant

Cross-validation on this dataset optimized to a power prior of 3, a base prior of 0.25, and 50 trees.

## BART hyperparameters selected by 10-fold cross-validation 
## To re-run the CV search itself, call bart_fit_binary_outcome() with vector-valued Power/Base/Ntrees (takes > 2 hours).
CV_POWER  <- 3
CV_BASE   <- 0.25
CV_NTREE  <- 50

BART_NDPOST    <- 1000L   # retained posterior draws per chain
BART_NSKIP     <- 500     # burn-in
BART_NCHAINS   <- 4
BART_KEEPEVERY <- 1       # thinning; raise if the Geweke plot looks bad

step1.res <- cached(paste0("step1_bart_", OUTCOME_TYPE), {
  fit_fun <- if (BINARY) bart_fit_binary_outcome else bart_fit_continuous_outcome
  fit_fun(X, y, z,
          Power     = CV_POWER,
          Base      = CV_BASE,
          Ntrees    = CV_NTREE,
          seed      = GLOBAL_SEED,
          nchains   = BART_NCHAINS,
          ndpost    = BART_NDPOST,
          nskip     = BART_NSKIP,
          keepevery = BART_KEEPEVERY)
})

cate.res <- step1.res$cate_results
bartmod0 <- step1.res$bartmod0
bartmod1 <- step1.res$bartmod1

# Individual-level CATE plot with 95% credible intervals (shaded area)
plot_treatment_effects(cate.res, binary_outcome = TRUE)

Summary of the CATE estimates

ate_draws_rd <- rowMeans(step1.res$cate_posterior)

ate_df <- data.frame(
  Metric = c("Posterior Mean", "95% CrI lower", "95% CrI upper"),
  Effect = c(mean(ate_draws_rd),
             quantile(ate_draws_rd, 0.025),
             quantile(ate_draws_rd, 0.975)))

ind_df <- data.frame(
  Metric = c("Mean", "SD", "IQR lower (25th pctl)", "IQR upper (75th pctl)",
             "Min", "Max", "Prop. effect < 0"),
  Effect = c(mean(cate.res$cate), sd(cate.res$cate),
             quantile(cate.res$cate, 0.25), quantile(cate.res$cate, 0.75),
             min(cate.res$cate), max(cate.res$cate),
             mean(cate.res$cate < 0)))

## Risk ratio is defined only for a binary outcome
if (BINARY) {
  ate_draws_rr <- rowMeans(step1.res$rr_posterior)
  ate_df$Risk_Ratio <- c(mean(ate_draws_rr),
                         quantile(ate_draws_rr, 0.025),
                         quantile(ate_draws_rr, 0.975))
  ind_df$Risk_Ratio <- c(mean(cate.res$rr), sd(cate.res$rr),
                         quantile(cate.res$rr, 0.25), quantile(cate.res$rr, 0.75),
                         min(cate.res$rr), max(cate.res$rr),
                         mean(cate.res$rr < 1))
  ind_df$Metric[7] <- "Prop. posterior-mean RD<0 / RR<1"
}

knitr::kable(ate_df, digits = 4,
             col.names = c("", EFFECT_LABEL, if (BINARY) "Risk Ratio"),
             caption = "Average CATE among all patients")

Average CATE among all patients

		Risk Difference	Risk Ratio
	Posterior Mean	-0.0155	0.9763
2.5%	95% CrI lower	-0.1145	0.8227
97.5%	95% CrI upper	0.0635	1.1172

knitr::kable(ind_df, digits = 4,
             col.names = c("", EFFECT_LABEL, if (BINARY) "Risk Ratio"),
             caption = "Individual CATE (posterior means across patients) summary")

Individual CATE (posterior means across patients) summary

	Risk Difference	Risk Ratio
Mean	-0.0155	0.9763
SD	0.0024	0.0050
IQR lower (25th pctl)	-0.0168	0.9736
IQR upper (75th pctl)	-0.0155	0.9772
Min	-0.0176	0.9639
Max	-0.0050	0.9945
Prop. posterior-mean RD<0 / RR<1	1.0000	1.0000

MCMC diagnostics

Diagnostics follow Sparapani et al. (2021). First, the autocorrelation of the estimated response surface from BART for 10 randomly selected subjects, one panel per chain. It examines how independent sequential MCMC draws are from one another and evaluates mixing efficiency (how rapidly the MCMC chain explores the posterior space). Spikes should rapidly decay within the dashed white-noise band. Failure to decay indicates high autocorrelation, necessitating an increase in the thinning parameter (BART_KEEPEVERY).

#' Autocorrelation of the fitted response surface, one panel per chain.
#'
#' Examines mixing (efficiency): how fast the sampler forgets its own past. Not a
#' convergence test -- a chain stuck in the wrong place can look perfect here.
#' 
#' GOOD: spikes fall inside the dashed band within ~5-10 lags and stay there,
#' and all panels look alike. The band is +/- 1.96/sqrt(ndpost), inside which
#' an autocorrelation is indistinguishable from zero.
#' BAD: decay still visible at the right edge -> raise BART_KEEPEVERY.
plot_bart_acf <- function(bartmod, n_subjects = 10, seed = NULL) {
  if (!is.null(seed)) set.seed(seed)
  nchains <- dim(bartmod$yhat.train)[3]
  n_obs   <- dim(bartmod$yhat.train)[2]
  ndraw <- dim(bartmod$yhat.train)[1]
  
  # Panel grid follows nchains 
  nr <- floor(sqrt(nchains)); nc <- ceiling(nchains / nr) 
  op <- par(mfrow = c(nr, nc)); on.exit(par(op))
  
  ## Sample subjects ONCE, outside the loop, so a difference between panels is
  ## a difference between chains and not a different set of patients.
  idx <- sample(seq_len(n_obs), n_subjects)
  
  band <- 1.96 / sqrt(ndraw)   # white noise band
  cols <- grDevices::hcl.colors(n_subjects, "Dark 3")   # base palette recycled at 8
  j <- seq(-0.5, 0.4, length.out = n_subjects)
  
  for (ch in seq_len(nchains)) {
    auto.corr <- acf(bartmod$yhat.train[, idx, ch], plot = FALSE)
    max.lag <- max(auto.corr$lag[, 1, 1])

    plot(NA, xlim = c(0, max.lag + 1), ylim = c(-1, 1),
         ylab = "acf", xlab = "lag", main = paste("chain", ch))
    abline(h = 0, col = "grey70")
    abline(h = c(-band, band), lty = 2, col = "grey40")
    for (h in seq_len(n_subjects))
      lines(seq_len(max.lag) + j[h], auto.corr$acf[1 + seq_len(max.lag), h, h],
            type = "h", col = cols[h])
  }
  invisible(idx)
}

## bartmod0 and bartmod1 are ONE model queried at two test sets: they share a
## seed, so their yhat.train is identical. One call covers both.
plot_bart_acf(bartmod0, seed = GLOBAL_SEED)  

Next, we compute the Geweke Z-statistic for each individual, which tests for stationarity (convergence) within a single chain by comparing the mean of the first 10% of iterations against the last 50%. Test statistics follow a standard normal distribution under the null hypothesis of stationarity. If many points lie beyond the dark blue line or further, consider increasing the burn-in period (BART_NSKIP) or the thinning interval (BART_KEEPEVERY) during model fitting.

#' Geweke Z statistics per individual, one panel per chain.
#'
#' Examines within-chain stationarity: whether the first 10% of a chain has the
#' same mean as its last 50%. If the chain has settled into a stationary distribution, 
#' the early and late draws should look like they come from the exact same distribution
#' (Z-scores mostly between -2 and 2).
#' Says nothing about whether the chains agree with each other.
#' 
#' GOOD: a flat, structureless band centered on 0, no trend across i.
#' One test per patient, so exceedances are expected. Several beyond the blue
#' line (3.291), or any beyond green (3.891), is a real signal -> raise
#' BART_NSKIP and/or BART_KEEPEVERY.
plot_bart_geweke <- function(bartmod, main = NULL) {
  arr <- bartmod$yhat.train
  if (is.null(arr)) stop("bartmod has no yhat.train array.")
  n <- dim(arr)[2]; nchains <- dim(arr)[3]

  levels_z <- c(1.96, 2.576, 3.291, 3.891, 4.417)
  labs     <- c("0.95", "0.99", "0.999", "0.9999", "0.99999")
  cols     <- c(6, 5, 4, 3, 2)

  nr <- floor(sqrt(nchains)); nc <- ceiling(nchains / nr)
  op <- par(mfrow = c(nr, nc)); on.exit(par(op))

  out <- vector("list", nchains)
  for (ch in seq_len(nchains)) {
    z <- BART::gewekediag(arr[, , ch])$z     # per chain, NOT the stacked matrix
    out[[ch]] <- vapply(levels_z, function(L) sum(abs(z) > L), numeric(1))
    plot(z, pch = ".", cex = 2, ylab = "z", xlab = "i",
         xlim = c(-n / 10, n), ylim = c(-5, 5),
         main = if (is.null(main)) paste("chain", ch) else
                paste0(main, " - chain ", ch))
    for (i in seq_along(levels_z)) {
      abline(h = c(-levels_z[i], levels_z[i]), col = cols[i])
      text(c(1, 1), c(-levels_z[i], levels_z[i]), pos = 2, cex = 0.6, labels = labs[i])
    }
  }

  tab <- data.frame(
    line      = labs,
    threshold = levels_z,
    expected  = round(n * 2 * (1 - pnorm(levels_z)), 2),
    observed  = apply(do.call(cbind, out), 1, function(r) paste(r, collapse = "/")))
  names(tab)[4] <- paste0("observed (chain 1/../", nchains, ")")
  
  cat("\n====================================================================\n")
  cat(" Geweke Diagnostic Summary: Expected vs. Observed Z-Score Outliers\n")
  cat("====================================================================\n")
  cat("Interpretation: This table compares the expected number of Z-scores\n")
  cat("falling outside standard normal thresholds against the actual observed\n")
  cat("counts per chain. If observed counts drastically exceed expected counts,\n")
  cat("it suggests the MCMC chain has not reached a stationary distribution.\n\n")
  print(tab, row.names = FALSE)
  invisible(tab)
}

plot_bart_geweke(bartmod0)

## 
## ====================================================================
##  Geweke Diagnostic Summary: Expected vs. Observed Z-Score Outliers
## ====================================================================
## Interpretation: This table compares the expected number of Z-scores
## falling outside standard normal thresholds against the actual observed
## counts per chain. If observed counts drastically exceed expected counts,
## it suggests the MCMC chain has not reached a stationary distribution.
## 
##     line threshold expected observed (chain 1/../4)
##     0.95     1.960    21.55                3/2/19/8
##     0.99     2.576     4.31                 0/0/0/0
##    0.999     3.291     0.43                 0/0/0/0
##   0.9999     3.891     0.04                 0/0/0/0
##  0.99999     4.417     0.00                 0/0/0/0

---

1.2 Subgroup Discovery via Classification and Regression Trees (CART)

To find subgroups exhibiting heterogeneity of treatment effect, we fit a CART model recursively partitioning the Stage 1 posterior-mean CATE estimates utilizing the baseline covariates.

The resulting tree maps continuous variables on their original scale. The top value in each box is the estimated mean CATE in that subgroup; below it is the proportion of the trial sample in the subgroup; the interval in each leaf is a 95% credible region for the subgroup mean.

cart_fit <- function(point_est, posterior, X,
                     metric = "Risk Difference",
                     maxdepth = 3, cp = 0.01, digits = 3,
                     label = NULL,
                     varimp = TRUE) {
  ## point_est  : length-n vector of per-patient posterior-mean CATEs
  ## posterior  : (n_draws x n_patients) matrix of posterior draws
  ## X          : data.frame of covariates for CART splits
  ## metric     : label for titles and printed output
  ## varimp     : draw the variable-importance barplot

  stopifnot(ncol(posterior) == length(point_est))
  stopifnot(nrow(X) == length(point_est))

  n_total <- length(point_est)

  data_cart <- data.frame(X, outcome = point_est)
  cartmod <- rpart(
    outcome ~ ., data = data_cart, method = "anova",
    control = rpart.control(cp = cp, maxdepth = maxdepth),
    model = TRUE
  )

  ## Precompute posterior CrI for terminal nodes only
  frame <- cartmod$frame
  node_info <- vector("list", nrow(frame))

  for (i in seq_len(nrow(frame))) {
    if (frame$var[i] != "<leaf>") next
    idx <- which(cartmod$where == i)
    if (length(idx) == 0) next
    node_avg_draws <- rowMeans(posterior[, idx, drop = FALSE])
    node_info[[i]] <- list(
      mean = mean(node_avg_draws),
      lo   = quantile(node_avg_draws, 0.025),
      hi   = quantile(node_avg_draws, 0.975),
      n    = length(idx),
      pct  = round(100 * length(idx) / n_total)
    )
  }

  ## Custom node label: CrI only for leaves, default for internal nodes
  node_label <- function(x, labs, digits, varlen) {
    fr <- x$frame
    for (i in seq_along(labs)) {
      if (fr$var[i] == "<leaf>" && !is.null(node_info[[i]])) {
        info <- node_info[[i]]
        labs[i] <- sprintf("%s\n%d%%\n95%% CrI:\n(%s, %s)",
                           formatC(info$mean, format = "f", digits = 3),
                           info$pct,
                           formatC(info$lo,   format = "f", digits = 3),
                           formatC(info$hi,   format = "f", digits = 3))
      }
    }
    labs
  }

  ## Variable importance
  if (varimp && length(cartmod$variable.importance)) {
    barplot(sort(cartmod$variable.importance), horiz = TRUE, las = 1,
            cex.names = 0.6, cex.axis = 0.6,
            main = paste0("CART Variable Importance\n(", metric, ")"))
  }
  
  ## Tree with CrI labels
  rpart.plot(cartmod,
             main = ifelse(is.null(label), paste("CART for", metric), label),
             yesno = 2,
             node.fun = node_label)

  ## Printed summaries
  cat("\n====", toupper(metric), "SUBGROUPS ====\n")
  for (i in unique(cartmod$where)) {
    idx <- which(cartmod$where == i)
    node_avg_draws <- rowMeans(posterior[, idx, drop = FALSE])
    node_mean <- mean(node_avg_draws)
    node_ci <- quantile(node_avg_draws, c(0.025, 0.975))
    cat("\nNode:", i, "  n =", length(idx), "\n")
    cat("  Mean", metric, ":", round(node_mean, digits), "\n")
    cat("  95% Posterior CrI for Subgroup Mean: [",
        round(node_ci[1], digits), ",", round(node_ci[2], digits), "]\n")
  }

  invisible(cartmod)
}

# CART on the additive effect scale
cart_rd <- cart_fit(
  point_est = cate.res$cate,
  posterior = step1.res$cate_posterior,
  X = X_unscaled,
  metric = EFFECT_LABEL,
  maxdepth = 3, cp = 0.01, digits = 3
)

## 
## ==== RISK DIFFERENCE SUBGROUPS ====
## 
## Node: 4   n = 135 
##   Mean Risk Difference : -0.017 
##   95% Posterior CrI for Subgroup Mean: [ -0.123 , 0.066 ]
## 
## Node: 7   n = 233 
##   Mean Risk Difference : -0.016 
##   95% Posterior CrI for Subgroup Mean: [ -0.12 , 0.067 ]
## 
## Node: 9   n = 33 
##   Mean Risk Difference : -0.008 
##   95% Posterior CrI for Subgroup Mean: [ -0.068 , 0.036 ]
## 
## Node: 5   n = 13 
##   Mean Risk Difference : -0.015 
##   95% Posterior CrI for Subgroup Mean: [ -0.114 , 0.061 ]
## 
## Node: 8   n = 17 
##   Mean Risk Difference : -0.014 
##   95% Posterior CrI for Subgroup Mean: [ -0.106 , 0.058 ]

# CART on the risk ratio scale (binary outcome only)
cart_rr <- cart_fit(
  point_est = cate.res$rr,
  posterior = step1.res$rr_posterior,
  X = X_unscaled,
  metric = "Risk Ratio",
  maxdepth = 3, cp = 0.01, digits = 3
)

## 
## ==== RISK RATIO SUBGROUPS ====
## 
## Node: 4   n = 130 
##   Mean Risk Ratio : 0.972 
##   95% Posterior CrI for Subgroup Mean: [ 0.792 , 1.146 ]
## 
## Node: 7   n = 238 
##   Mean Risk Ratio : 0.976 
##   95% Posterior CrI for Subgroup Mean: [ 0.82 , 1.115 ]
## 
## Node: 9   n = 33 
##   Mean Risk Ratio : 0.991 
##   95% Posterior CrI for Subgroup Mean: [ 0.916 , 1.047 ]
## 
## Node: 5   n = 16 
##   Mean Risk Ratio : 0.977 
##   95% Posterior CrI for Subgroup Mean: [ 0.821 , 1.112 ]
## 
## Node: 8   n = 14 
##   Mean Risk Ratio : 0.982 
##   95% Posterior CrI for Subgroup Mean: [ 0.852 , 1.086 ]

1.3 Variable Importance and Interaction (VIVI) Heatmaps

Inglis et al. (2024) use Value Suppressing Uncertainty Palettes (VSUP) to build heatmaps that jointly display variable importance (VImp) and variable interactions (VInt), with the color scale encoding posterior uncertainty.

Below, the estimated CATEs from Part I are regressed on the covariates alone by fitting a new BART model, to identify which covariates are predictive of treatment effect heterogeneity. The resulting VIVI–VSUP heatmap summarizes variable importance and interactions from this model: diagonal elements correspond to marginal variable importance, while off-diagonal elements represent pairwise interaction strength. Color intensity reflects magnitude, and desaturation/lightness reflects posterior uncertainty (parameterized by the coefficient of variation, following the bartMan author’s preference).

zip <- bartMan:::zip

# scaled covariates matrix: expanding factors to a set of dummy variables
X_mat <- model.matrix(~ ., data = X)[, -1]

# Fit a BART model to estimate the CATE
# keeptrees = TRUE retains the posterior tree structures for downstream analysis
vivi_rd <- cached("vivi_vsup_rd", {
  set.seed(GLOBAL_SEED)
  m <- bart(x.train = X_mat, y.train = cate.res$cate,
            ntree = 50, ndpost = 500, nskip = 200, keeptrees = TRUE)
  
  # Extract posterior tree structures from the fitted BART model
  fData <- as.data.frame(cbind(X_mat, cate.est = cate.res$cate))
  td <- extractTreeData(model = m, data = fData)
  
  # Compute variable importance and interaction measures from the tree structures
  list(vsup = viviBartMatrix(td, type = 'vsup', metric = 'propMean',
                             metricError = "CV"),
       structure_head = head(td$structure, 5))
})

options(tibble.width = Inf) # Display all columns when printing tibbles  
vivi_rd$structure_head

VIVI_rd <- viviBartPlot(vivi_rd$vsup,
             max_desat = 1,
             pow_desat = 0.6,
             max_light = 0.6,
             pow_light = 1,
             label = 'CV') +
  labs(
    title = paste0("Variable Importance and Interaction Heatmap \n (", EFFECT_LABEL, " Scale)"),
  ) +
  theme(
    axis.text.x = element_text(angle = 90, hjust = 1, vjust = 0.5, size = 11),
    axis.text.y = element_text(size = 11),
    legend.key.width = unit(1.0, "cm"),
    legend.key.height = unit(0.7, "cm"),
    legend.text = element_text(size = 9),
    legend.title = element_text(size = 11),
    plot.title   = element_text(size = 13, face = "bold", hjust = 0),
    plot.margin  = margin(10, 25, 10, 10)    
  ) +
  coord_cartesian(clip = "off")

vivi_rr <- cached("vivi_vsup_rr", {
  set.seed(GLOBAL_SEED)
  m <- bart(x.train = X_mat, y.train = cate.res$rr,
            ntree = 50, ndpost = 500, nskip = 200, keeptrees = TRUE)
  fData <- as.data.frame(cbind(X_mat, cate.est = cate.res$rr))
  td <- extractTreeData(model = m, data = fData)
  list(vsup = viviBartMatrix(td, type = 'vsup', metric = 'propMean',
                             metricError = "CV"),
       structure_head = head(td$structure, 5))
})

options(tibble.width = Inf) # used to display full tibble in output
vivi_rr$structure_head

VIVI_rr <- viviBartPlot(vivi_rr$vsup,
             max_desat = 1,
             pow_desat = 0.6,
             max_light = 0.6,
             pow_light = 1,
             label = 'CV') +
  labs(
    title = "Variable Importance and Interaction Heatmap \n (Risk Ratio Scale)",
  ) +
  theme(
    axis.text.x = element_text(angle = 90, hjust = 1, vjust = 0.5, size = 11),
    axis.text.y = element_text(size = 11),
    legend.key.width = unit(1.0, "cm"),
    legend.key.height = unit(0.7, "cm"),
    legend.text = element_text(size = 9),
    legend.title = element_text(size = 11),
    plot.title   = element_text(size = 13, face = "bold", hjust = 0),
    plot.margin  = margin(10, 25, 10, 10)    
  ) +
  coord_cartesian(clip = "off")

print(VIVI_rd)

print(VIVI_rr)

Part II. Sensitivity analyses and cross-method comparison

2.1 Sensitivity analysis of the BART prior specification

To assess whether the Part I finding is robust to prior specification, BART is refit under a panel of prior configurations spanning under-regularized, default, CV-selected, and over-regularized regimes. The four hyperparameters varied are:

• k (leaf-value shrinkage): primarily controls effect size magnitude. Larger k shrinks fitted values more tightly toward the overall mean.
• power (depth penalty): primarily controls interaction capacity. Higher power penalizes deeper trees, making treatment-by-covariate interactions harder to learn.
• base (split probability): also controls interaction capacity. Lower base reduces the probability of splitting, discouraging the deep trees needed for interaction detection.
• ntree (ensemble size): more trees reduce each tree’s individual contribution.

prior_grid <- tibble::tribble(
  ~label,             ~power, ~base, ~k, ~ntree,
  "default",               2, 0.95, 2, 200,     # BART default
  "original_cv_selected", CV_POWER, CV_BASE, 2, CV_NTREE, # CV-selected in primary analysis
  "under_regularized",     1, 0.95, 1, 50,      # everything relaxed: deep trees, few, weak k
  "deep_trees_default_k",  1, 0.95, 2, 200,     # isolates tree-depth effect (power=1 vs 2)
  "low_k_default_struct",  2, 0.95, 1, 200,     # isolates leaf-shrinkage effect (k=1 vs 2)
  "high_k_default_struct", 2, 0.95, 5, 200,     # strong leaf shrinkage
  "mildly_over_reg",       2, 0.75, 3, 200,
  "over_regularized",      3, 0.50, 3, 400
)

knitr::kable(prior_grid, caption = "Prior configurations to be compared")

Prior configurations to be compared

label	power	base	k	ntree
default	2	0.95	2	200
original_cv_selected	3	0.25	2	50
under_regularized	1	0.95	1	50
deep_trees_default_k	1	0.95	2	200
low_k_default_struct	2	0.95	1	200
high_k_default_struct	2	0.95	5	200
mildly_over_reg	2	0.75	3	200
over_regularized	3	0.50	3	400

Helper: fit BART once with fixed priors

#' `fit_bart_fixed()` is a thin adapter over `bart_fit_binary_outcome()` from Part I,
#' called with a single hyperparameter combination so the cross-validation branch is
#' skipped. It exists only to reshape the output into the names this Part uses.
fit_bart_fixed <- function(X, y, z, power, base, k, ntree,
                           seed = GLOBAL_SEED, nchains = BART_NCHAINS,
                           ndpost = BART_NDPOST, nskip = BART_NSKIP,
                           keepevery = 1) {

  ## Dispatch on OUTCOME_TYPE 
  fit_fun <- if (BINARY) bart_fit_binary_outcome else bart_fit_continuous_outcome

  fit <- fit_fun(
    X = X, y = y, z = z, seed = seed,
    Power = power, Base = base, Ntrees = ntree, k = k,
    nchains = nchains, ndpost = ndpost, nskip = nskip, keepevery = keepevery
  )

  cate_post <- fit$cate_posterior   # (draws x patients), additive effect scale

  list(
    cate_mean = colMeans(cate_post),
    cate_sd   = apply(cate_post, 2, sd),
    cate_lo   = apply(cate_post, 2, quantile, 0.025),
    cate_hi   = apply(cate_post, 2, quantile, 0.975),
    cate_post = cate_post,
    config    = list(power = power, base = base, k = k, ntree = ntree)
  )
}

Run the prior-sensitivity grid

prior_results <- if (!RUN_PRIOR_GRID) NULL else cached(paste0("prior_results_", OUTCOME_TYPE), {
  res <- vector("list", nrow(prior_grid))
  names(res) <- prior_grid$label
  for (i in seq_len(nrow(prior_grid))) {
    cfg <- prior_grid[i, ]
    cat("\n--- Fitting", cfg$label, "---\n")
    cat("  power =", cfg$power, "  base =", cfg$base,
        "  k =", cfg$k, "  ntree =", cfg$ntree, "\n")
    res[[i]] <- fit_bart_fixed(X = X, y = y, z = z,
                               power = cfg$power, base = cfg$base,
                               k = cfg$k, ntree = cfg$ntree)
  }
  res
})

CATE distributions across prior configurations

Summary table

prior_summary <- purrr::map_dfr(seq_along(prior_results), function(i) {
  r <- prior_results[[i]]
  tibble::tibble(
    label = names(prior_results)[i],
    power = r$config$power,
    base = r$config$base,
    k = r$config$k,
    ntree = r$config$ntree,
    mean_CATE = mean(r$cate_mean),
    median_CATE = median(r$cate_mean),
    sd_across_pts = sd(r$cate_mean),
    avg_within_post_sd = mean(r$cate_sd),
    shrinkage_ratio = sd(r$cate_mean) / mean(r$cate_sd),
  )
})

knitr::kable(prior_summary, digits = 3,
             caption = "CATE summaries across hyperparameter configurations")

CATE summaries across hyperparameter configurations

label	power	base	k	ntree	mean_CATE	median_CATE	sd_across_pts	avg_within_post_sd	shrinkage_ratio
default	2	0.95	2	200	-0.028	-0.019	0.020	0.066	0.307
original_cv_selected	3	0.25	2	50	-0.015	-0.016	0.002	0.042	0.056
under_regularized	1	0.95	1	50	-0.027	-0.011	0.052	0.097	0.530
deep_trees_default_k	1	0.95	2	200	-0.027	-0.014	0.034	0.077	0.436
low_k_default_struct	2	0.95	1	200	-0.029	-0.013	0.047	0.086	0.542
high_k_default_struct	2	0.95	5	200	-0.019	-0.017	0.004	0.050	0.087
mildly_over_reg	2	0.75	3	200	-0.024	-0.022	0.006	0.055	0.117
over_regularized	3	0.50	3	400	-0.020	-0.020	0.003	0.049	0.053

The shrinkage ratio is the standard deviation of posterior-mean CATEs across patients divided by the average within-patient posterior standard deviation. Values near 0 indicate the posterior cannot separate patients (heavy shrinkage); values above 1 indicate clear between-patient separation relative to posterior uncertainty.

Density plot

plot_df <- purrr::map_dfr(seq_along(prior_results), function(i) {
  tibble::tibble(
    label = names(prior_results)[i],
    cate  = prior_results[[i]]$cate_mean
  )
})

ggplot(plot_df, aes(x = cate, color = label)) +
  geom_density(linewidth = 0.7) +
  geom_vline(xintercept = 0, linetype = "dotted") +
  labs(
    title = "CATE distribution under different BART priors",
    subtitle = paste0(EFFECT_LABEL, " scale."),
    x = "Per-patient posterior-mean CATE",
    y = "Density",
    color = "Prior configuration"
  ) +
  theme_bw()

Caterpillar plots with individual credible intervals

cat_df <- purrr::map_dfr(prior_grid$label, function(lab) {
  r <- prior_results[[lab]]
  tibble::tibble(
    label = lab,
    cate  = r$cate_mean,
    lo    = r$cate_lo,
    hi    = r$cate_hi
  ) %>%
    arrange(cate) %>%
    mutate(rank = row_number())
})

ggplot(cat_df, aes(x = rank, y = cate)) +
  geom_ribbon(aes(ymin = lo, ymax = hi), fill = "grey80", alpha = 0.6) +
  geom_line(linewidth = 0.6) +
  geom_hline(yintercept = 0, linetype = "dotted", color = "red") +
  facet_wrap(~ label, ncol = 1, scales = "free_y") +
  labs(
    title = "Per-patient CATE posterior mean (line) and 95% credible interval (ribbon)",
    subtitle = paste0(EFFECT_LABEL, " scale."),
    x = "Patient index (sorted by CATE within each panel)",
    y = "CATE with 95% credible interval"
  ) +
  theme_bw()

---

2.2 Cross-method comparison

Bayesian Causal Forest (BCF)

Bayesian Causal Forest (Hahn, Murray, and Carvalho 2020) separates the prognostic surface from the treatment effect surface using distinct tree ensembles.

# BCF requires a numeric covariate matrix: expand factors
X_mat_num <- model.matrix(~ . - 1, data = X)

# Ensure treatment is coded as a numeric 0/1 indicator
if (is.factor(z)) {
  if (!all(levels(z) %in% c("0", "1"))) {
    stop("Treatment factor z must have levels '0' and '1'.")
  }
  z_num <- as.numeric(as.character(z))
} else {
  z_num <- as.numeric(z)
}

stopifnot(all(z_num %in% c(0, 1)))

# BCF requires a propensity score estimate (even in an RCT, we estimate it to
# capture minor imbalances)
pihat_glm <- fitted(glm(z_num ~ ., data = as.data.frame(X_mat_num),
                        family = binomial))

## Cache the derived quantities
bcf_res <- cached(paste0("bcf_fit_", OUTCOME_TYPE), {
  set.seed(GLOBAL_SEED)
  sink(tempfile())
  fit <- bcf::bcf(
    y = y, z = z_num,
    x_control = X_mat_num, x_moderate = X_mat_num, pihat = pihat_glm,
    nburn = 2000, nsim = 2000, n_chains = 4,
    random_seed = GLOBAL_SEED, no_output = TRUE
  )
  sink()

  ## bcf::bcf fits a GAUSSIAN outcome model, y = mu(x) + tau(x) * z + eps.
  ## For a 0/1 outcome this is a linear probability model, which makes `tau` the risk difference directly.
  post <- fit$tau

  list(cate_post = post,
       cate_mean = colMeans(post),
       cate_sd   = apply(post, 2, sd),
       cate_lo   = apply(post, 2, quantile, 0.025),
       cate_hi   = apply(post, 2, quantile, 0.975))
})

cate_bcf_post <- bcf_res$cate_post
cate_bcf_mean <- bcf_res$cate_mean
cate_bcf_sd   <- bcf_res$cate_sd
cate_bcf_lo   <- bcf_res$cate_lo
cate_bcf_hi   <- bcf_res$cate_hi

cat("BCF (", EFFECT_LABEL, "scale) summary:\n")

## BCF ( Risk Difference scale) summary:

cat("mean of posterior-mean CATEs:", round(mean(cate_bcf_mean), 4), "\n")

## mean of posterior-mean CATEs: -0.0115

cat("SD across patients:", round(sd(cate_bcf_mean), 5), "\n")

## SD across patients: 0.01603

cat("mean within-patient post SD:", round(mean(cate_bcf_sd), 5), "\n")

## mean within-patient post SD: 0.05532

cat("shrinkage ratio:", round(sd(cate_bcf_mean) / mean(cate_bcf_sd), 3), "\n")

## shrinkage ratio: 0.29

XGBoost T-learner and S-learner

The T-learner (Two-model) fits separate outcome models on the treated and control arms and takes the per-patient difference as the CATE. Hyperparameters are selected via 5-fold cross-validation optimized for log-loss or RMSE with early stopping.

## Objective and evaluation metric follow OUTCOME_TYPE.
XGB_OBJECTIVE <- if (BINARY) "binary:logistic"  else "reg:squarederror"
XGB_METRIC    <- if (BINARY) "logloss"          else "rmse"
XGB_CV_COL    <- paste0("test_", XGB_METRIC, "_mean")

## Search grids. The T- and S-learner grids differ only in the depth range; the
## tuner itself is shared, so the two cannot drift apart again.
XGB_GRID_T <- expand.grid(
  eta = c(0.01, 0.05),
  max_depth = c(1, 2, 3),
  subsample = c(0.6, 0.8),        # fraction of observations sampled per round
  min_child_weight = c(10, 20))   # prevents leaves with < 10-20 patients

XGB_GRID_S <- expand.grid(
  eta = c(0.01, 0.05),
  max_depth = c(1, 2, 3),
  subsample = c(0.6, 0.8),
  min_child_weight = c(10, 20))

#' Grid-search XGBoost hyperparameters by 5-fold CV with early stopping.
#'
#' Returns list(loss, params, nrounds, tuned). `tuned` is FALSE when the search
#' found nothing and the fallback was used
tune_xgb <- function(X_train, y_train, param_grid = XGB_GRID_T,
                     nrounds_max = 300, nfold = 5, seed = GLOBAL_SEED) {
  dtrain <- xgb.DMatrix(data = as.matrix(X_train), label = y_train)
  best <- list(loss = Inf, params = NULL, nrounds = NULL)
  n_failed <- 0L

  for (i in seq_len(nrow(param_grid))) {
    cv_seed <- seed + i
    set.seed(cv_seed)  # ensure reproducible CV folds 
    
    cv <- tryCatch(
      xgb.cv(
        data = dtrain,
        params = list(
          objective = XGB_OBJECTIVE,
          eval_metric = XGB_METRIC,
          eta = param_grid$eta[i],
          max_depth = param_grid$max_depth[i],
          min_child_weight = param_grid$min_child_weight[i],
          subsample = param_grid$subsample[i],
          nthread = 1,
          seed = cv_seed
        ),
        nrounds = nrounds_max,
        nfold = nfold,
        early_stopping_rounds = 10,
        verbose = 0
      ),
      
      error = function(e) {
        message("  xgb.cv config ", i, " failed: ", conditionMessage(e))
        NULL
      }
    )
    if (is.null(cv)) { n_failed <- n_failed + 1L; next }

    best_iter <- if (!is.null(cv$early_stop)) {
      cv$early_stop$best_iteration
    } else {
      cv$best_iteration
    }

    if (is.null(best_iter) || length(best_iter) == 0 || best_iter == 0) {
      n_failed <- n_failed + 1L; next
    }

    if (!XGB_CV_COL %in% names(cv$evaluation_log)) {
      warning("evaluation_log has no column '", XGB_CV_COL, "'; available: ",
              paste(names(cv$evaluation_log), collapse = ", "))
      n_failed <- n_failed + 1L; next
    }

    best_loss <- cv$evaluation_log[[XGB_CV_COL]][best_iter]
    if (length(best_loss) == 0 || is.na(best_loss)) { n_failed <- n_failed + 1L; next }

    if (best_loss < best$loss) {
      best <- list(loss    = best_loss,
                   params  = as.list(param_grid[i, ]),
                   nrounds = best_iter)
    }
  }

  if (is.null(best$params)) {
    warning("No valid XGBoost configuration found: all ", nrow(param_grid),
            " candidates failed. Using fallback defaults - results are NOT CV-tuned.")
    best <- list(loss = NA, params = list(eta = 0.05, max_depth = 2,
                 min_child_weight = 10, subsample = 1.0), nrounds = 50)
    best$tuned <- FALSE
  } else {
    if (n_failed > 0L)
      message("  ", n_failed, " of ", nrow(param_grid), " configurations skipped")
    best$tuned <- TRUE
  }
  best
}

#' Fit one XGBoost model with a given parameter set.
fit_xgb <- function(X_train, y_train, params, nrounds, seed = GLOBAL_SEED) {
  xgb.train(
    data    = xgb.DMatrix(data = as.matrix(X_train), label = y_train),
    params  = c(list(objective   = XGB_OBJECTIVE,
                     eval_metric = XGB_METRIC,
                     nthread     = 1,
                     seed = seed), params),
    nrounds = nrounds,
    verbose = 0
  )
}

#' Format a tuning result for printing, flagging the fallback.
xgb_tune_label <- function(tune) {
  paste0(toString(sprintf("%s=%s", names(tune$params), tune$params)),
         " nrounds = ", tune$nrounds,
         if (isTRUE(tune$tuned)) "" else "   [FALLBACK - not CV-tuned]")
}

# XGBoost requires a numeric design matrix: expand factors
X_mat_xgb <- model.matrix(~ . - 1, data = X)

treated <- z_num == 1
control <- z_num == 0

xgb_T <- cached(paste0("xgb_tlearner_", OUTCOME_TYPE), {
  # Train on treated arm
  tune_t <- tune_xgb(X_mat_xgb[treated, ], y[treated], param_grid = XGB_GRID_T,
      seed = GLOBAL_SEED)
  
  mod_t <- fit_xgb(X_mat_xgb[treated, ], y[treated], tune_t$params, tune_t$nrounds,
                   seed = GLOBAL_SEED + 100)

  # Train on control arm
  tune_c <- tune_xgb(X_mat_xgb[control, ], y[control], param_grid = XGB_GRID_T,
                     seed = GLOBAL_SEED + 1)
   
  mod_c <- fit_xgb(X_mat_xgb[control, ], y[control], tune_c$params, tune_c$nrounds,
                   seed = GLOBAL_SEED + 101)

  # Predict counterfactuals for the whole cohort
  list(tune_t = tune_t, tune_c = tune_c,
       cate = predict(mod_t, as.matrix(X_mat_xgb)) -
              predict(mod_c, as.matrix(X_mat_xgb)))
})

tune_t <- xgb_T$tune_t; tune_c <- xgb_T$tune_c
cate_xgb_T <- xgb_T$cate   # RD scale, negative = benefit

cat("XGBoost T-learner:\n")

## XGBoost T-learner:

cat("Treated-arm tuning:", xgb_tune_label(tune_t), "\n")

## Treated-arm tuning: eta=0.01, max_depth=3, subsample=0.8, min_child_weight=10 nrounds = 1

cat("Control-arm tuning:", xgb_tune_label(tune_c), "\n")

## Control-arm tuning: eta=0.05, max_depth=1, subsample=0.6, min_child_weight=10 nrounds = 33

cat("Mean CATE:", round(mean(cate_xgb_T), 4),
    "SD across patients:", round(sd(cate_xgb_T), 4),
    "CATE range: [", round(min(cate_xgb_T), 4), ",", round(max(cate_xgb_T), 4), "]\n")

## Mean CATE: -0.0253 SD across patients: 0.0491 CATE range: [ -0.0912 , 0.0506 ]

The S-learner (Single-model) fits a single outcome model on the full sample with treatment assignment included as a covariate, and estimates per-patient CATEs by toggling the treatment indicator while holding baseline covariates fixed.

X_all <- cbind(z = z_num, X_mat_xgb)

xgb_S <- cached(paste0("xgb_slearner_", OUTCOME_TYPE), {
   
  tune_s <- tune_xgb(X_all, y, param_grid = XGB_GRID_S,
      seed = GLOBAL_SEED + 2)

  mod_s <- fit_xgb(X_all, y, tune_s$params, tune_s$nrounds,
      seed = GLOBAL_SEED + 102)

  # Counterfactual predictions: toggle Z for everyone
  X_treat <- X_all; X_treat[, "z"] <- 1
  X_ctrl  <- X_all; X_ctrl[, "z"]  <- 0
  list(tune_s = tune_s,
       cate = predict(mod_s, as.matrix(X_treat)) -
              predict(mod_s, as.matrix(X_ctrl)))
})

tune_s <- xgb_S$tune_s
cate_xgb_S <- xgb_S$cate   # RD scale, negative = benefit

cat("XGBoost S-learner:\n")

## XGBoost S-learner:

cat("Tuning:", xgb_tune_label(tune_s), "\n")

## Tuning: eta=0.05, max_depth=2, subsample=0.8, min_child_weight=20 nrounds = 5

cat("Mean CATE:", round(mean(cate_xgb_S), 4),
    "SD:", round(sd(cate_xgb_S), 4),
    "range: [", round(min(cate_xgb_S), 4), ",",
    round(max(cate_xgb_S), 4), "]\n")

## Mean CATE: 0.0011 SD: 9e-04 range: [ 0 , 0.0018 ]

Causal Forest

Causal forests (Athey, Tibshirani, and Wager 2019) use honest sample splitting to prevent overfitting and provide asymptotically valid inference for individualized treatment effects.

X_mat_grf <- model.matrix(~ . - 1, data = X)

cf <- cached(paste0("causal_forest_", OUTCOME_TYPE), {
  set.seed(GLOBAL_SEED)
  causal_forest(
    X = X_mat_grf, Y = y, W = z_num,
    num.trees = 2000, honesty = TRUE,
    tune.parameters = "all", seed = GLOBAL_SEED
  )
})

cate_grf <- predict(cf)$predictions
cate_grf_var <- predict(cf, estimate.variance = TRUE)$variance.estimates

ate <- average_treatment_effect(cf)
cat("Causal Forest results:\n")

## Causal Forest results:

cat("Mean CATE:", round(mean(cate_grf), 4),
    "SD:", round(sd(cate_grf), 4),
    "range: [", round(min(cate_grf), 4), ",", round(max(cate_grf), 4), "]\n")

## Mean CATE: -0.0212 SD: 0.0033 range: [ -0.0295 , -0.014 ]

cat("ATE estimate:", round(ate["estimate"], 4),
    "SE:", round(ate["std.err"], 4), "\n")

## ATE estimate: -0.0165 SE: 0.0476

Cross-method comparison

overlay_df <- bind_rows(
  if (RUN_PRIOR_GRID) tibble::tibble(method = "BART (default priors)",
                 cate   = prior_results[["default"]]$cate_mean),
  tibble::tibble(method = "BART (CV-selected, Part I)", cate = cate.res$cate),
  tibble::tibble(method = "Bayesian Causal Forest", cate = cate_bcf_mean),
  tibble::tibble(method = "Causal Forest",          cate = cate_grf),
  tibble::tibble(method = "XGBoost T-learner",      cate = cate_xgb_T),
  tibble::tibble(method = "XGBoost S-learner",      cate = cate_xgb_S)
)

# Violin plots
p_full <- ggplot(overlay_df, aes(x = method, y = cate, fill = method)) +
  geom_violin(alpha = 0.4, trim = FALSE) +
  geom_boxplot(width = 0.15, fill = "white", outlier.size = 0.8) +
  geom_hline(yintercept = 0, linetype = "dotted", color = "red") +
  labs(title = "CATE distribution across methods",
       subtitle = paste0(EFFECT_LABEL, " scale."),
       x = "", y = "Per-patient CATE") +
  theme_bw() +
  theme(legend.position = "none",
        axis.text.x = element_text(angle = 25, hjust = 1))
print(p_full)

# Zoomed companion: visualize the regularized methods
# print(p_full +
#  coord_cartesian(ylim = c(-0.3, 0.3)) +
#  labs(title = "CATE distribution across methods (zoomed to [-0.3, 0.3])",
#       subtitle = "Distributions extending beyond the axis limits are clipped."))

method_summary <- overlay_df %>%
  group_by(method) %>%
  summarize(
    mean   = mean(cate),
    sd     = sd(cate),
    Q25    = quantile(cate, 0.25),
    median = median(cate),
    Q75    = quantile(cate, 0.75),
    range = paste0("[", round(min(cate), 4), ", ", round(max(cate), 4), "]"),
    .groups = "drop"
  )

knitr::kable(method_summary, digits = 4,
             caption = paste0("Per-method CATE summary (", tolower(EFFECT_LABEL), " scale)"))

Per-method CATE summary (risk difference scale)

method	mean	sd	Q25	median	Q75	range
BART (CV-selected, Part I)	-0.0155	0.0024	-0.0168	-0.0158	-0.0155	[-0.0176, -0.005]
BART (default priors)	-0.0276	0.0204	-0.0528	-0.0195	-0.0130	[-0.0656, 0.0018]
Bayesian Causal Forest	-0.0115	0.0160	-0.0293	-0.0052	-0.0010	[-0.0378, 0.0361]
Causal Forest	-0.0212	0.0033	-0.0239	-0.0211	-0.0188	[-0.0295, -0.014]
XGBoost S-learner	0.0011	0.0009	0.0000	0.0018	0.0018	[0, 0.0018]
XGBoost T-learner	-0.0253	0.0491	-0.0812	-0.0137	0.0163	[-0.0912, 0.0506]

---

Part III. Treatment effect heterogeneity by prognostic score

The prognostic score denotes the expected untreated outcome conditional on baseline covariates, \(S_0(X) = E[Y(0)\mid X]\). For a binary outcome, this represents the baseline predicted risk, \(S_0(X) = P(Y(0)=1\mid X)\).

Two complementary analyses are considered: treatment effects within prognostic score strata (quartiles), and Bayesian causal machine learning using the prognostic score as the sole covariate.

3.1 Prognostic score estimation

When available, a high-quality, externally-developed, compatible prognostic model should be used to stratify trial results. When an external model is unavailable, three internal approaches to prognostic score estimation may be considered.

1. Control-arm only model. A prognostic model is fit using control-arm participants only and then applied to predict the untreated outcome for all participants. This targets baseline expected outcome, \(E[Y(0)\mid X]\), which is the natural prognostic quantity for assessing whether treatment effects vary by baseline prognosis. However, because control outcomes are reused in both prognostic score estimation and downstream HTE analysis, the score is estimated in-sample for control patients but out-of-sample for treated patients. This asymmetry may induce a spurious CATE-by-risk gradient that mimics HTE even when none exists, particularly when the number of covariates is large relative to the control-arm sample size.

2. Full-sample model without treatment indicator. Fit a regression of \(Y\) on \(X\) using all participants, excluding treatment assignment \(Z\). This may maintain appropriate type I error control (though with no guarantee) but could have lower power for detecting HTE. The resulting score also lacks a natural interpretation: it estimates risk averaged over treated and control conditions in a trial-dependent proportion, \(E(Y \mid X)\), rather than baseline risk under no treatment.

3. Sample splitting. Reserve a subset of control-arm patients to fit the prognostic model and exclude them from downstream HTE analyses. The remaining controls and all treated participants form the analysis sample. This prevents outcome reuse and ensures independence between prognostic score estimation and HTE analysis, but reduces sample size for both stages, leading to potentially less stable prognostic score estimation and lower power. Best suited to larger trials.

## Prognostic-model family follows OUTCOME_TYPE: logistic for a binary outcome,
## Gaussian (i.e. linear regression) for a continuous one. The predicted value is
## a baseline risk in the binary case and a baseline expected outcome otherwise;
## the quartile stratification below is identical either way.
PROG_FAMILY <- if (BINARY) binomial() else gaussian()

## ---------- Approach 1: control-arm only ----------
prog_fit_ctrl <- glm(y ~ ., data = data.frame(y = y[!z], X[!z, ]),
                     family = PROG_FAMILY)
risk_ctrl <- predict(prog_fit_ctrl, newdata = X, type = "response")

## ---------- Approach 2: full-sample, no treatment indicator ----------
prog_fit_full <- glm(y ~ ., data = data.frame(y = y, X), family = PROG_FAMILY)
risk_full <- predict(prog_fit_full, type = "response")

## ---------- Approach 3: sample-split control-arm ----------
split_frac <- 0.3  # fraction of controls reserved for score training
set.seed(GLOBAL_SEED)
n_ctrl <- sum(!z)
ctrl_idx <- which(!z)
split_train_idx <- sample(ctrl_idx, size = floor(split_frac * n_ctrl))
split_hte_ctrl  <- setdiff(ctrl_idx, split_train_idx)
split_hte_idx   <- sort(c(split_hte_ctrl, which(z))) # remaining controls + all treated

# Fit prognostic model on reserved controls only
prog_fit_split <- glm(y ~ ., data = data.frame(y = y[split_train_idx],
                                               X[split_train_idx, ]),
                      family = PROG_FAMILY)

# Score all HTE-analysis patients (remaining controls + treated)
risk_split <- predict(prog_fit_split, newdata = X[split_hte_idx, ],
                      type = "response")

# Subset outcome and treatment vectors for HTE analysis
y_split <- y[split_hte_idx]
z_split <- z[split_hte_idx]
X_split <- X[split_hte_idx, ]

cat(sprintf("Sample split approach:\n"))

## Sample split approach:

cat(sprintf("Controls reserved for scoring: %d\n", length(split_train_idx)))

## Controls reserved for scoring: 62

cat(sprintf("Patients in HTE analysis: %d (controls: %d, treated: %d)\n",
            length(split_hte_idx), length(split_hte_ctrl), sum(z)))

## Patients in HTE analysis: 369 (controls: 146, treated: 223)

cat("\nCorrelation between prognostic scores on the full sample:\n")

## 
## Correlation between prognostic scores on the full sample:

cat(sprintf("  Control-arm vs Full-sample:  %.4f\n", cor(risk_ctrl, risk_full)))

##   Control-arm vs Full-sample:  0.9082

cat("Correlation between prognostic scores on the HTE-analysis subset:\n")

## Correlation between prognostic scores on the HTE-analysis subset:

cat(sprintf("  Control-arm vs Sample-split: %.4f\n",
            cor(risk_ctrl[split_hte_idx], risk_split)))

##   Control-arm vs Sample-split: 0.7140

cat(sprintf("  Full-sample vs Sample-split: %.4f\n",
            cor(risk_full[split_hte_idx], risk_split)))

##   Full-sample vs Sample-split: 0.7175

3.2 Observed treatment effects across prognostic-score strata

Patients are stratified into quartiles of prognostic score. Within each stratum, we report the average treatment effect with 95% bootstrap confidence intervals.

compute_risk_strat <- function(risk_pred, y, z, B = 1000, label = "") {
  risk_q <- cut(risk_pred,
                breaks = quantile(risk_pred, probs = seq(0, 1, 0.25)),
                include.lowest = TRUE,
                labels = c("Q1 (Lowest Risk)", "Q2", "Q3", "Q4 (Highest Risk)"))

  strat_rows <- list()
  for (q in levels(risk_q)) {
    idx <- which(risk_q == q)
    yq  <- y[idx]; zq <- z[idx]
    rd_q <- mean(yq[zq]) - mean(yq[!zq])

    rd_boot <- replicate(B, {
      bidx <- sample(idx, replace = TRUE)
      mean(y[bidx][z[bidx]]) - mean(y[bidx][!z[bidx]])
    })

    strat_rows[[q]] <- tibble::tibble(
      approach = label,
      stratum  = q,
      n        = length(idx),
      n_treat  = sum(zq),
      risk_diff = rd_q,
      CI_lower  = quantile(rd_boot, 0.025, na.rm = TRUE),
      CI_upper  = quantile(rd_boot, 0.975, na.rm = TRUE)
    )
  }
  bind_rows(strat_rows)
}

set.seed(GLOBAL_SEED)
strat_all <- bind_rows(
  compute_risk_strat(risk_ctrl, y, z, label = "Control-arm only"),
  compute_risk_strat(risk_full, y, z, label = "Full-sample (no treatment indicator)"),
  compute_risk_strat(risk_split, y_split, z_split, label = "Sample-split control-arm")
)
knitr::kable(strat_all, digits = 4,
             caption = paste("Observed", tolower(EFFECT_LABEL), "by baseline risk quartile"))

Observed risk difference by baseline risk quartile

approach	stratum	n	n_treat	risk_diff	CI_lower	CI_upper
Control-arm only	Q1 (Lowest Risk)	129	73	0.0521	-0.1164	0.2450
Control-arm only	Q2	92	51	0.0660	-0.1490	0.2789
Control-arm only	Q3	105	47	-0.0686	-0.2445	0.1199
Control-arm only	Q4 (Highest Risk)	105	52	-0.0816	-0.2540	0.0873
Full-sample (no treatment indicator)	Q1 (Lowest Risk)	118	62	0.0161	-0.1718	0.2020
Full-sample (no treatment indicator)	Q2	115	64	-0.1158	-0.2922	0.0446
Full-sample (no treatment indicator)	Q3	93	48	0.1153	-0.0789	0.3123
Full-sample (no treatment indicator)	Q4 (Highest Risk)	105	49	-0.0510	-0.2097	0.1111
Sample-split control-arm	Q1 (Lowest Risk)	96	62	-0.0531	-0.2275	0.1563
Sample-split control-arm	Q2	95	52	-0.1400	-0.3363	0.0391
Sample-split control-arm	Q3	93	61	0.1183	-0.1044	0.3381
Sample-split control-arm	Q4 (Highest Risk)	85	48	-0.0546	-0.2389	0.1195

ggplot(strat_all, aes(x = stratum, y = risk_diff)) +
  geom_point(size = 2.5) +
  geom_errorbar(aes(ymin = CI_lower, ymax = CI_upper), width = 0.15) +
  geom_hline(yintercept = 0, linetype = "dotted", color = "red") +
  facet_wrap(~ approach, nrow = 1) +
  labs(title = "Treatment effect by baseline risk quartile",
       subtitle = "95% bootstrap CIs. Three prognostic score approaches compared.",
       x = "Baseline Risk Quartile (Q1: Lowest, Q4: Highest)",
       y = paste0(EFFECT_LABEL, " (Treatment - Control)")) +
  theme_bw() +
  theme(axis.text.x = element_text(angle = 30, hjust = 1, size = 7))

#' Per-stratum outcome summary by treatment arm.
#' Binary outcome  -> 2x2 count tables, event rates, and risk difference.
#' Continuous      -> n, mean, SD per arm, and mean difference.
print_strata_summary <- function(risk_pred, y, z, label = "") {
  risk_q <- cut(risk_pred,
                breaks = quantile(risk_pred, probs = seq(0, 1, 0.25)),
                include.lowest = TRUE,
                labels = c("Q1 (Lowest Risk)", "Q2", "Q3", "Q4 (Highest Risk)"))

  cat("\n\n##########", toupper(label), "##########\n")
  for (q in levels(risk_q)) {
    idx <- which(risk_q == q)
    yq  <- y[idx]; zq <- z[idx]

    if (BINARY) {
       cat("\n========== Risk quartile:", q, "(n =", length(idx), ") ==========\n")
      tab <- table(
        Arm = factor(ifelse(zq, "Treatment", "Control"),
                     levels = c("Treatment", "Control")),
        Outcome = factor(ifelse(yq == 1, "Worse", "Better"),
                         levels = c("Better", "Worse")))
      print(addmargins(tab))
      ev_treat   <- sum(yq[zq])  / sum(zq)
      ev_control <- sum(yq[!zq]) / sum(!zq)
      cat(sprintf("  Treatment event rate: %d/%d = %.3f\n",
                  sum(yq[zq]),  sum(zq),  ev_treat))
      cat(sprintf("  Control event rate:   %d/%d = %.3f\n",
                  sum(yq[!zq]), sum(!zq), ev_control))
      cat(sprintf("  Unadjusted RD:        %.3f\n", ev_treat - ev_control))
    } else {
       cat("\n========== Prognostic-score quartile:", q, "(n =", length(idx), ") ==========\n")
      print(data.frame(
        Arm  = c("Treatment", "Control"),
        n    = c(sum(zq), sum(!zq)),
        mean = c(mean(yq[zq]), mean(yq[!zq])),
        sd   = c(sd(yq[zq]),   sd(yq[!zq])),
        row.names = NULL), digits = 4)
      cat(sprintf("  Unadjusted mean difference: %.4f\n",
                  mean(yq[zq]) - mean(yq[!zq])))
    }
  }
}

print_strata_summary(risk_ctrl, y, z, "Control-arm only")

## 
## 
## ########## CONTROL-ARM ONLY ##########
## 
## ========== Risk quartile: Q1 (Lowest Risk) (n = 129 ) ==========
##            Outcome
## Arm         Better Worse Sum
##   Treatment     34    39  73
##   Control       29    27  56
##   Sum           63    66 129
##   Treatment event rate: 39/73 = 0.534
##   Control event rate:   27/56 = 0.482
##   Unadjusted RD:        0.052
## 
## ========== Risk quartile: Q2 (n = 92 ) ==========
##            Outcome
## Arm         Better Worse Sum
##   Treatment     24    27  51
##   Control       22    19  41
##   Sum           46    46  92
##   Treatment event rate: 27/51 = 0.529
##   Control event rate:   19/41 = 0.463
##   Unadjusted RD:        0.066
## 
## ========== Risk quartile: Q3 (n = 105 ) ==========
##            Outcome
## Arm         Better Worse Sum
##   Treatment     17    30  47
##   Control       17    41  58
##   Sum           34    71 105
##   Treatment event rate: 30/47 = 0.638
##   Control event rate:   41/58 = 0.707
##   Unadjusted RD:        -0.069
## 
## ========== Risk quartile: Q4 (Highest Risk) (n = 105 ) ==========
##            Outcome
## Arm         Better Worse Sum
##   Treatment     17    35  52
##   Control       13    40  53
##   Sum           30    75 105
##   Treatment event rate: 35/52 = 0.673
##   Control event rate:   40/53 = 0.755
##   Unadjusted RD:        -0.082

print_strata_summary(risk_full, y, z, "Full-sample (no treatment indicator)")

## 
## 
## ########## FULL-SAMPLE (NO TREATMENT INDICATOR) ##########
## 
## ========== Risk quartile: Q1 (Lowest Risk) (n = 118 ) ==========
##            Outcome
## Arm         Better Worse Sum
##   Treatment     30    32  62
##   Control       28    28  56
##   Sum           58    60 118
##   Treatment event rate: 32/62 = 0.516
##   Control event rate:   28/56 = 0.500
##   Unadjusted RD:        0.016
## 
## ========== Risk quartile: Q2 (n = 115 ) ==========
##            Outcome
## Arm         Better Worse Sum
##   Treatment     30    34  64
##   Control       18    33  51
##   Sum           48    67 115
##   Treatment event rate: 34/64 = 0.531
##   Control event rate:   33/51 = 0.647
##   Unadjusted RD:        -0.116
## 
## ========== Risk quartile: Q3 (n = 93 ) ==========
##            Outcome
## Arm         Better Worse Sum
##   Treatment     19    29  48
##   Control       23    22  45
##   Sum           42    51  93
##   Treatment event rate: 29/48 = 0.604
##   Control event rate:   22/45 = 0.489
##   Unadjusted RD:        0.115
## 
## ========== Risk quartile: Q4 (Highest Risk) (n = 105 ) ==========
##            Outcome
## Arm         Better Worse Sum
##   Treatment     13    36  49
##   Control       12    44  56
##   Sum           25    80 105
##   Treatment event rate: 36/49 = 0.735
##   Control event rate:   44/56 = 0.786
##   Unadjusted RD:        -0.051

print_strata_summary(risk_split, y_split, z_split, "Sample-split control-arm")

## 
## 
## ########## SAMPLE-SPLIT CONTROL-ARM ##########
## 
## ========== Risk quartile: Q1 (Lowest Risk) (n = 96 ) ==========
##            Outcome
## Arm         Better Worse Sum
##   Treatment     27    35  62
##   Control       13    21  34
##   Sum           40    56  96
##   Treatment event rate: 35/62 = 0.565
##   Control event rate:   21/34 = 0.618
##   Unadjusted RD:        -0.053
## 
## ========== Risk quartile: Q2 (n = 95 ) ==========
##            Outcome
## Arm         Better Worse Sum
##   Treatment     23    29  52
##   Control       13    30  43
##   Sum           36    59  95
##   Treatment event rate: 29/52 = 0.558
##   Control event rate:   30/43 = 0.698
##   Unadjusted RD:        -0.140
## 
## ========== Risk quartile: Q3 (n = 93 ) ==========
##            Outcome
## Arm         Better Worse Sum
##   Treatment     29    32  61
##   Control       19    13  32
##   Sum           48    45  93
##   Treatment event rate: 32/61 = 0.525
##   Control event rate:   13/32 = 0.406
##   Unadjusted RD:        0.118
## 
## ========== Risk quartile: Q4 (Highest Risk) (n = 85 ) ==========
##            Outcome
## Arm         Better Worse Sum
##   Treatment     13    35  48
##   Control        8    29  37
##   Sum           21    64  85
##   Treatment event rate: 35/48 = 0.729
##   Control event rate:   29/37 = 0.784
##   Unadjusted RD:        -0.055

3.3. CATE estimation as a function of the prognostic score

BART and BCF are refit using the prognostic score as the sole covariate. This evaluates whether the treatment effect varies continuously with baseline prognosis. We compare CATE estimates derived across all three prognostic score approaches.

fit_cate_by_risk <- function(risk_pred, y, z, label = "", seed = GLOBAL_SEED) {
  ## BART
  bart_res <- fit_bart_fixed(
    X = data.frame(prog_score = risk_pred), y = y, z = z,
    power = 2, base = 0.95, k = 2, ntree = 200, seed = seed
  )

  ## BCF
  X_bcf <- matrix(risk_pred, ncol = 1)
  colnames(X_bcf) <- "prog_score"
  pihat <- fitted(glm(z ~ risk_pred, family = binomial))

  set.seed(seed)
  sink(tempfile())
  bcf_fit <- bcf::bcf(
    y = y, z = z,
    x_control = X_bcf, x_moderate = X_bcf, pihat = pihat,
    nburn = 2000, nsim = 2000, n_chains = 4,
    random_seed = seed, no_output = TRUE
  )
  sink()

  cate_bcf_post <- bcf_fit$tau
  rm(bcf_fit)   # drop the multi-hundred-MB fit object before this is cached

  list(
    bart = bart_res,
    bcf_cate_mean = colMeans(cate_bcf_post),
    bcf_cate_lo   = apply(cate_bcf_post, 2, quantile, 0.025),
    bcf_cate_hi   = apply(cate_bcf_post, 2, quantile, 0.975),
    bcf_cate_post = cate_bcf_post,
    risk_pred     = risk_pred,
    label         = label
  )
}

cate_risk <- cached(paste0("cate_by_risk_", OUTCOME_TYPE), {
  list(
    ctrl  = fit_cate_by_risk(risk_ctrl,  y, z_num, "Control-arm only"),
    full  = fit_cate_by_risk(risk_full,  y, z_num, "Full-sample (no treatment indicator)"),
    split = fit_cate_by_risk(risk_split, y_split, z_split, "Sample-split control-arm")
  )
})

build_cate_df <- function(res) {
  bind_rows(
    tibble::tibble(approach = res$label, method = "BART",
                   prog_score = res$risk_pred, cate = res$bart$cate_mean,
                   lo = res$bart$cate_lo, hi = res$bart$cate_hi),
    tibble::tibble(approach = res$label, method = "Bayesian Causal Forest",
                   prog_score = res$risk_pred, cate = res$bcf_cate_mean,
                   lo = res$bcf_cate_lo, hi = res$bcf_cate_hi)
  )
}

cate_risk_df <- bind_rows(lapply(cate_risk, build_cate_df))

ggplot(cate_risk_df, aes(x = prog_score, y = cate, color = method, fill = method)) +
  geom_ribbon(aes(ymin = lo, ymax = hi), alpha = 0.12, color = NA) +
  geom_point(size = 0.6, alpha = 0.5) +
 # geom_smooth(method = "loess", se = FALSE, linewidth = 0.8) +
  geom_hline(yintercept = 0, linetype = "dotted", color = "red") +
  facet_grid(method ~ approach) +
  labs(x = "Predicted baseline prognosis",
       y = paste0("CATE (", EFFECT_LABEL, ")"),
       title = "CATE as a function of baseline prognosis",
       subtitle = "Rows: BART vs BCF. Columns: prognostic score approach.") +
  theme_bw(base_size = 14) +
  theme(legend.position = "none")

risk_cate_summary <- bind_rows(lapply(cate_risk, function(res) {
  tibble::tibble(
    approach           = res$label,
    BART_mean_CATE     = mean(res$bart$cate_mean),
    BART_sd_CATE       = sd(res$bart$cate_mean),
    BART_cor_risk_cate = cor(res$risk_pred, res$bart$cate_mean),
    BCF_mean_CATE      = mean(res$bcf_cate_mean),
    BCF_sd_CATE        = sd(res$bcf_cate_mean),
    BCF_cor_risk_cate  = cor(res$risk_pred, res$bcf_cate_mean)
  )
}))

knitr::kable(risk_cate_summary, digits = 4,
             caption = "CATE-by-baseline-prognosis summary across prognostic score approaches")

CATE-by-baseline-prognosis summary across prognostic score approaches

approach	BART_mean_CATE	BART_sd_CATE	BART_cor_risk_cate	BCF_mean_CATE	BCF_sd_CATE	BCF_cor_risk_cate
Control-arm only	-0.0023	0.0428	-0.5761	-0.0049	0.0076	-0.3754
Full-sample (no treatment indicator)	-0.0197	0.0133	0.6798	-0.0116	0.0087	0.8982
Sample-split control-arm	-0.0478	0.0346	0.8936	-0.0255	0.0116	0.9528

CART subgroup discovery by baseline prognosis

# The same `cart_fit()` from Part II, with `varimp = FALSE` and a custom title.
risk_scores <- list(ctrl = risk_ctrl, full = risk_full, split = risk_split)

par(mfrow = c(3, 2))

for (nm in names(cate_risk)) {
  res <- cate_risk[[nm]]
  cart_fit(point_est = res$bart$cate_mean,
           posterior = res$bart$cate_post,
           X = data.frame(prog_score = risk_scores[[nm]]),
           metric = EFFECT_LABEL, maxdepth = 2, varimp = FALSE,
           label = paste0("BART (", res$label, ")"))

  cart_fit(point_est = res$bcf_cate_mean,
           posterior = res$bcf_cate_post,
           X = data.frame(prog_score = risk_scores[[nm]]),
           metric = EFFECT_LABEL, maxdepth = 2, varimp = FALSE,
           label = paste0("BCF (", res$label, ")"))
}

## 
## ==== RISK DIFFERENCE SUBGROUPS ====
## 
## Node: 7   n = 77 
##   Mean Risk Difference : 0.05 
##   95% Posterior CrI for Subgroup Mean: [ -0.126 , 0.225 ]
## 
## Node: 3   n = 142 
##   Mean Risk Difference : -0.056 
##   95% Posterior CrI for Subgroup Mean: [ -0.195 , 0.08 ]
## 
## Node: 4   n = 46 
##   Mean Risk Difference : -0.009 
##   95% Posterior CrI for Subgroup Mean: [ -0.117 , 0.098 ]
## 
## Node: 6   n = 166 
##   Mean Risk Difference : 0.021 
##   95% Posterior CrI for Subgroup Mean: [ -0.115 , 0.15 ]

## 
## ==== RISK DIFFERENCE SUBGROUPS ====
## 
## Node: 7   n = 77 
##   Mean Risk Difference : 0.005 
##   95% Posterior CrI for Subgroup Mean: [ -0.087 , 0.118 ]
## 
## Node: 3   n = 171 
##   Mean Risk Difference : -0.013 
##   95% Posterior CrI for Subgroup Mean: [ -0.109 , 0.061 ]
## 
## Node: 4   n = 39 
##   Mean Risk Difference : 0 
##   95% Posterior CrI for Subgroup Mean: [ -0.17 , 0.186 ]
## 
## Node: 6   n = 144 
##   Mean Risk Difference : -0.002 
##   95% Posterior CrI for Subgroup Mean: [ -0.085 , 0.084 ]

## 
## ==== RISK DIFFERENCE SUBGROUPS ====
## 
## Node: 3   n = 44 
##   Mean Risk Difference : -0.038 
##   95% Posterior CrI for Subgroup Mean: [ -0.229 , 0.148 ]
## 
## Node: 6   n = 172 
##   Mean Risk Difference : -0.015 
##   95% Posterior CrI for Subgroup Mean: [ -0.145 , 0.115 ]
## 
## Node: 7   n = 49 
##   Mean Risk Difference : 0.004 
##   95% Posterior CrI for Subgroup Mean: [ -0.086 , 0.101 ]
## 
## Node: 4   n = 166 
##   Mean Risk Difference : -0.027 
##   95% Posterior CrI for Subgroup Mean: [ -0.158 , 0.105 ]

## 
## ==== RISK DIFFERENCE SUBGROUPS ====
## 
## Node: 4   n = 374 
##   Mean Risk Difference : -0.014 
##   95% Posterior CrI for Subgroup Mean: [ -0.096 , 0.053 ]
## 
## Node: 5   n = 48 
##   Mean Risk Difference : 0.012 
##   95% Posterior CrI for Subgroup Mean: [ -0.133 , 0.199 ]
## 
## Node: 3   n = 9 
##   Mean Risk Difference : -0.028 
##   95% Posterior CrI for Subgroup Mean: [ -0.247 , 0.116 ]

## 
## ==== RISK DIFFERENCE SUBGROUPS ====
## 
## Node: 4   n = 105 
##   Mean Risk Difference : -0.086 
##   95% Posterior CrI for Subgroup Mean: [ -0.248 , 0.082 ]
## 
## Node: 7   n = 214 
##   Mean Risk Difference : -0.021 
##   95% Posterior CrI for Subgroup Mean: [ -0.15 , 0.108 ]
## 
## Node: 6   n = 22 
##   Mean Risk Difference : -0.048 
##   95% Posterior CrI for Subgroup Mean: [ -0.202 , 0.104 ]
## 
## Node: 3   n = 28 
##   Mean Risk Difference : -0.106 
##   95% Posterior CrI for Subgroup Mean: [ -0.316 , 0.108 ]

## 
## ==== RISK DIFFERENCE SUBGROUPS ====
## 
## Node: 4   n = 123 
##   Mean Risk Difference : -0.038 
##   95% Posterior CrI for Subgroup Mean: [ -0.178 , 0.053 ]
## 
## Node: 6   n = 210 
##   Mean Risk Difference : -0.02 
##   95% Posterior CrI for Subgroup Mean: [ -0.119 , 0.065 ]
## 
## Node: 7   n = 29 
##   Mean Risk Difference : -0.002 
##   95% Posterior CrI for Subgroup Mean: [ -0.189 , 0.223 ]
## 
## Node: 3   n = 7 
##   Mean Risk Difference : -0.054 
##   95% Posterior CrI for Subgroup Mean: [ -0.313 , 0.083 ]

---

Part IV. Randomization-based inference for ITE quantiles

The fit-the-fit approach estimates how the average effect varies with the observed covariates. This section addresses HTE on a different axis: it gives randomization-based inference for the marginal distribution and quantiles of the individual treatment effects \[\tau_i = Y_i(1) - Y_i(0), \qquad F_n(c) = \tfrac1n\sum_{i=1}^n \mathbf 1(\tau_i \le c),\] for the \(n\) randomized patients.

1. Different axis of heterogeneity. A CATE captures heterogeneity explained by the observed covariates. The ITE distribution makes a statement about the individual effects themselves, rather than modeling how covariates modify those effects.

2. Different identification and inference regime. Part I is a model-based Bayesian estimate that leans on a prior and a fitted response surface. This section is randomization-based and finite-sample exact under the trial’s assignment mechanism. The trade-off is partial identification. Because each patient is observed under only one arm, the joint distribution of \((Y_i(1), Y_i(0))\), and hence the distribution of \(\tau_i\), is not identified. The method therefore returns confidence bounds, not point estimates.

3. Interpreting a null result. Because the estimand is partially identified and clinical trials are often small, these bounds can be wide or uninformative. An uninformative bound is a legitimate, reportable finding: it simply indicates that the randomized data alone, free of modeling assumptions, cannot definitively pin down a non-trivial benefiting fraction. That is not evidence of homogeneity.

Setup: shuffle the units

The shuffle matters: the rank statistic breaks ties by index order, so the method assumes units have been randomly permuted before analysis. Record the shuffle seed with any reported result.

ALPHA <- 0.1    # 1 - ALPHA confidence
NPERM <- 1e4    # permutation draws for the null distribution

## Randomly shuffle the units before analysis: the rank statistic breaks ties by index order 
set.seed(GLOBAL_SEED)
ord <- sample(length(z))
Z   <- as.numeric(z)[ord]
Y_larger_bad <- y[ord]
Y_larger_good <- y_cont[ord]
n   <- length(Z)

cat("n =", n, " treated =", sum(Z), " control =", sum(1 - Z), "\n")

## n = 431  treated = 223  control = 208

One-sided confidence intervals for effect quantiles

Sign convention

If a larger outcome is…	“Benefit” means	Read	In this trial
good	\(\tau_i > 0\)	lower limits	y_cont (function score, ventilator-free days)
bad	\(\tau_i < 0\)	upper limits	y (mortality, more ADL limitations)

Each point represents the \(1-\alpha\) lower (or upper) confidence limit for \(\tau_{(k)}\), the \(k\)-th smallest individual effect, and the inference is simultaneously valid across all \(k\). Quantiles whose lower (upper) confidence limit is \(-\infty\) (\(+\infty\)) are uninformative and are omitted from the plot.

# Function: Plot one-sided confidence limits for effect quantiles
plot_quantile_bound <- function(ci.limit, bound_type = "lower", k_start = NULL, 
                                caption = NULL, main = NULL, 
                                x_custom = FALSE, x_custom_range = c(-20, 20),
                                quantiles = seq(1, 0.1, -0.1), 
                                line = 3, fontsize = 1.2, numbersize = 1){
  n = length(ci.limit)
  ticks = ceiling(quantiles*n)
  
  # crop the y-axis
  num_finite <- sum(!is.nan(ci.limit) & is.finite(ci.limit))
  
  if (bound_type == "lower") {
    if (is.null(k_start)) {
      k_start <- n - num_finite
    }
    # Crop the bottom where values are -Inf, but keep the top at n
    ylim <- c(k_start - 1, n + 1)
    
  } else if (bound_type == "upper") {
    if (is.null(k_start)) {
      k_start <- 1
    }
    # Crop the top where values are +Inf, but keep the bottom at 1
    ylim <- c(k_start, num_finite + 1)
  }
  
  # Set the x-axis limits
  if(x_custom)
    xlim = x_custom_range
  else
    xlim = range(ci.limit[is.finite(ci.limit)]) * 1.1
  
  par(mar = c(4, 4, 2, 5) + 0.1)
 
  # Swap the label based on the requested bound type
  if (bound_type == "lower") {
  x_label = bquote(.(sprintf("%.0f%%", 100 * (1 - ALPHA))) ~ 
                   "simultaneous" ~ "one-sided" ~ "lower" ~ 
                   "confidence" ~ "interval" ~ "for" ~ tau[(k)])
} else {
  x_label = bquote(.(sprintf("%.0f%%", 100 * (1 - ALPHA))) ~ 
                   "simultaneous" ~ "one-sided" ~ "upper" ~ 
                   "confidence" ~ "interval" ~ "for" ~ tau[(k)])
}
  
  plot(NA, ylab = "k", xlab = x_label,  
       ylim = ylim, xlim = xlim, yaxt = "n",
       main = main, cex.lab=fontsize, cex.axis = numbersize)
  
  # Draw the grey intervals shooting in the correct direction
  for (k in 1:length(ci.limit)) {
    if (bound_type == "lower") {
      # Lower bounds: line shoots to the right
      lines(c(max(ci.limit[k], min(xlim) - 10 * diff(xlim), na.rm = TRUE), 
              max(ci.limit[is.finite(ci.limit)]) + 10 * diff(xlim)), 
            rep(k, 2), col = "grey")
    } else {
      # Upper bounds: line shoots from the left
      lines(c(min(ci.limit[is.finite(ci.limit)]) - 10 * diff(xlim), 
              min(ci.limit[k], max(xlim) + 10 * diff(xlim), na.rm = TRUE)), 
            rep(k, 2), col = "grey")
    }
  }
  
  points(ci.limit, c(1:length(ci.limit)), pch = 20)
  abline(v = 0, lty = 2)  # Draws a dashed (lty = 2) vertical reference line at 0.
  axis(side = 2, at = ticks, cex.axis = numbersize) # Draw the axes
  axis(side = 4, at = ticks, labels = paste0(quantiles * 100, '%'), cex = fontsize)
 if (!is.null(caption)) {
    mtext(caption, side = 1, line = caption_line, cex = 0.9, font = 1, col = "#555555")
  }
  mtext("quantile", side = 4, line = line, cex=fontsize)
}

# one-sided lower confidence interval
ci_one_sided_lower <- cached("ite_onesided_lower", {
  method_combine(Z = Z, Y = Y_larger_good, N = n, k_vec = 1:n, alpha = ALPHA, simul = TRUE,
                 treat.method.list   = list(name = "Stephenson", s = 6),
                 control.method.list = list(name = "Stephenson", s = 6),
                 nperm = NPERM)
})

informative <- ci_one_sided_lower$k[is.finite(ci_one_sided_lower$lower)]
cat(sprintf("smallest informative k = %s of %d (%.0f%% quantile)\n",
            if (length(informative)) min(informative) else NA, n,
            if (length(informative)) floor(100 * min(informative) / n) else NA))

## smallest informative k = 109 of 431 (25% quantile)

plot_quantile_bound(ci_one_sided_lower$lower, bound_type = "lower"
                    #, main = sprintf("%.0f%% simultaneous lower limits", 100 * (1 - ALPHA))
                    )

# one-sided upper confidence interval for the $k$-th quantile can be derived by switching the treatment labels and taking the opposite sign of the lower limit of the $(n-k+1)$-th quantile.
ci_one_sided_upper <- cached("ite_onesided_upper", {
  ub <- method_combine(Z = 1 - Z, Y = Y_larger_bad, N = n, k_vec = 1:n, alpha = ALPHA,
                       simul = TRUE, nperm = NPERM,
                       treat.method.list   = list(name = "Stephenson", s = 6),
                       control.method.list = list(name = "Stephenson", s = 6))
  data.frame(k = 1:n, lower = -Inf, upper = -ub$lower[n:1])
})

plot_quantile_bound(ci_one_sided_upper$upper, bound_type = "upper")

For a binary outcome \(\tau_i \in \{-1, 0, +1\}\), so the quantile function is a three-point step and the resulting bounds may carry limited information.

Two-sided confidence intervals

# Function for drawing two-sided confidence intervals for the effect quantiles
plot_two_sided_CIs <- function(lb, ub, k_start = NULL, main = NULL,
                               x_custom = FALSE, x_custom_range = c(-20, 20),
                               quantiles = c(1, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1),
                               line = 3, fontsize = 1.2, numbersize = 1){
  n = length(lb)
  ticks = ceiling(quantiles*n)
  if (is.null(k_start)){
    k_start = 1
  }
  ylim = c(k_start-1, n+1)
  if(x_custom)
    xlim = x_custom_range
  else
    xlim = range(c(lb[lb > -Inf], ub[ub<Inf])) * 1.1
  par(mar=c(4,4,2,5)+0.1)
  plot(NA, ylab = "k", xlab = bquote(.(sprintf("%.0f%%", 100 * (1 - ALPHA))) ~ 
                   "simultaneous" ~ "two-sided" ~  
                   "confidence" ~ "interval" ~ "for" ~ tau[(k)]),
       ylim = ylim, xlim = xlim, yaxt = "n",
       main = main, cex.lab=fontsize, cex.axis = numbersize)
  for (k in 1:length(lb)) {
    lines(c(max(lb[k], min(xlim)-10*diff(xlim)), min(ub[k], max(xlim) + 10*diff(xlim))), rep(k, 2), col = "grey")
  }
  points(lb, c(1:length(lb)), pch = 20)
  points(ub, c(1:length(ub)), pch = 20)
  abline(v = 0, lty = 2)
  axis(side = 2, at = ticks, cex.axis=numbersize)
  axis(side = 4, at = ticks, labels = paste0(floor(ticks*100/n), '%'), cex=fontsize )
  mtext("quantile", side = 4, line = line, cex=fontsize)
}

## Two-sided confidence intervals can be constructed by combining one-sided lower and upper intervals using a Bonferroni correction.
ci_two <- cached("ite_twosided", {
  lb <- method_combine(Z = Z, Y = Y_larger_good, N = n, k_vec = 1:n, alpha = ALPHA / 2,
                       simul = TRUE, nperm = NPERM,
                       treat.method.list   = list(name = "Stephenson", s = 6),
                       control.method.list = list(name = "Stephenson", s = 6))
  
  ub <- method_combine(Z = 1 - Z, Y = Y_larger_good, N = n, k_vec = 1:n, alpha = ALPHA / 2,
                       simul = TRUE, nperm = NPERM,
                       treat.method.list   = list(name = "Stephenson", s = 6),
                       control.method.list = list(name = "Stephenson", s = 6))
  data.frame(k = 1:n, lower = lb$lower, upper = -ub$lower[n:1])
})

lv  <- c(0.1, 0.25, 0.5, 0.75, 0.9, 1.0)
idx <- ceiling(lv * n)
knitr::kable(
  data.frame(`Quantile level` = paste0(100 * lv, "%"),
             k = idx,
             `Lower limit` = round(ci_two$lower[idx], 3),
             `Upper limit` = round(ci_two$upper[idx], 3),
             check.names = FALSE),
  caption = sprintf("%.0f%% two-sided confidence intervals for ITE quantiles",
                    100 * (1 - ALPHA)))

90% two-sided confidence intervals for ITE quantiles

Quantile level	k	Lower limit	Upper limit
10%	44	-Inf	5.999
25%	108	-Inf	11.999
50%	216	-21.999	22.999
75%	324	-11.000	Inf
90%	388	-5.000	Inf
100%	431	0.000	Inf

plot_two_sided_CIs(ci_two$lower, ci_two$upper
                  # main = sprintf("%.0f%% two-sided CIs for ITE quantiles", 100 * (1 - ALPHA))
                  )

What fraction of patients benefit?

If the lower limit satisfies \(L_k > c\) then \(\tau_{(k)} > c\), and because the effects are sorted, \(\tau_{(j)} > c\) for every \(j \ge k\). So at least \(n - k + 1\) patients have an effect above \(c\). The mirror argument on the upper limits bounds the number below \(c\). Both statements maintain \(1-\alpha\) coverage because the limits are simultaneously valid.

## Counted from the ONE-SIDED limits at level ALPHA, in the direction of benefit:
##   larger outcome GOOD -> benefit is tau_i > c, read from the LOWER limits
##   larger outcome BAD  -> benefit is tau_i < c, read from the UPPER limits

## (a) continuous score, larger is better
thresholds <- c(0, 2, 5)   # on the outcome scale; edit for clinical relevance
n_benefit <- sapply(thresholds, function(cc) {
  kk <- ci_one_sided_lower$k[ci_one_sided_lower$lower > cc]
  if (length(kk)) n - min(kk) + 1L else 0L })

knitr::kable(
  data.frame(
    `Threshold c`    = thresholds,
    `N with tau > c` = n_benefit,
    `% of patients`  = sprintf("%.1f%%", 100 * n_benefit / n),
    check.names = FALSE),
  align = c("l", "l", "r"),
  col.names = c("Threshold $c$", "$N$ with $\\tau_i > c$", "% of patients"),
  caption = sprintf(
    paste("%.0f%% lower bounds on the number of patients with ITE > c (n = %d)"),
    100 * (1 - ALPHA), n))

90% lower bounds on the number of patients with ITE > c (n = 431)

Threshold \(c\)	\(N\) with \(\tau_i > c\)	% of patients
0	0	0.0%
2	0	0.0%
5	0	0.0%

## (b) binary outcome, larger is worse: only c = 0 is meaningful, since
## tau_i is confined to {-1, 0, +1}
kk <- ci_one_sided_upper$k[ci_one_sided_upper$upper < 0]
n_helped <- if (length(kk)) max(kk) else 0L

knitr::kable(
  data.frame(
    `Threshold c`    = 0,
    `N with tau < 0` = n_helped,
    `% of patients`  = sprintf("%.1f%%", 100 * n_helped / n),
    check.names = FALSE),
  align = c("l", "l", "r"),
  col.names = c("Threshold $c$", "$N$ with $\\tau_i < 0$", "% of patients"),
  caption = sprintf(
    paste("%.0f%% lower bound on the number of patients with ITE < 0 (n = %d)"),
    100 * (1 - ALPHA), n))

90% lower bound on the number of patients with ITE < 0 (n = 431)

Threshold \(c\)	\(N\) with \(\tau_i < 0\)	% of patients
0	0	0.0%

Read a row as: with 90% confidence, at least this many of the 431 randomized patients had an individual effect beyond c in the direction of benefit.

An uninformative bound of zero simply means the data cannot certify a non-zero benefiting fraction; state this plainly rather than reporting it as evidence that no one benefits.

Worked example 2: the distracted-driving trial with a continuous outcome

The trial. A randomized trial of five interventions to reduce distracted driving. The outcome is Int_Percent_ActiveUse_TripSeconds, handheld phone use as a proportion of drive time during the intervention period. We compare Arm 5 (prize money) against Arm 1 (control).

Lower handheld use is better. So a treatment benefit corresponds to \(\tau_i < 0\), and statements regarding the number of benefiting drivers are read from the upper confidence limits.

DRIVE_PATH <- "continuous_outcome_driving_trial_data.csv"
drive <- read.csv(DRIVE_PATH)

## Arm 5 (prize money) vs Arm 1 (control).  
drive2 <- drive %>% filter(Arm %in% c("Arm  1", "Arm  5"))

z_drive <- drive2$Arm == "Arm  5"
y_drive <- drive2$Int_Percent_ActiveUse_TripSeconds

cat("Arm 5 vs Arm 1 analysis set:", nrow(drive2), "participants",
    "(treated:", sum(z_drive), " control:", sum(!z_drive), ")\n")

## Arm 5 vs Arm 1 analysis set: 582 participants (treated: 282  control: 300 )

# This is the near-tie-free endpoint: 573 distinct outcome values across 582 drivers, 
# against 40 distinct values across 431 patients for the continuous score and 2 for the binary endpoint.
cat("Outcome (proportion of drive time on a handheld phone):\n")

## Outcome (proportion of drive time on a handheld phone):

print(summary(y_drive))

##    Min. 1st Qu.  Median    Mean 3rd Qu.    Max. 
## 0.00000 0.03128 0.06352 0.08770 0.11404 0.51814

## Same shuffle-then-analyse pattern as Section 3.1.
set.seed(GLOBAL_SEED)
ord_dr <- sample(length(z_drive))
Z_dr   <- as.numeric(z_drive)[ord_dr]
Y_dr   <- as.numeric(y_drive)[ord_dr]
n_dr   <- length(Z_dr)

cat("n =", n_dr, " treated =", sum(Z_dr), " control =", sum(1 - Z_dr), "\n")

## n = 582  treated = 282  control = 300

Confidence limits for the effect quantiles

# one-sided upper confidence limits
ci_ub_dr <- cached("ite_onesided_upper_drive_allrows", {
  method_combine(Z = 1 - Z_dr, Y = Y_dr, N = n_dr, k_vec = 1:n_dr, alpha = ALPHA, simul = TRUE,
                 treat.method.list   = list(name = "Stephenson", s = 6),
                 control.method.list = list(name = "Stephenson", s = 6),
                 nperm = NPERM)
})

plot_quantile_bound(-ci_ub_dr$lower[n_dr:1], bound_type = "upper")

# two-sided confidence intervals
ci_dr <- cached("ite_twosided_drive", {
  lb <- method_combine(Z = Z_dr, Y = Y_dr, N = n_dr, k_vec = 1:n_dr,
                       alpha = ALPHA / 2, simul = TRUE, nperm = NPERM,
                       treat.method.list   = list(name = "Stephenson", s = 6),
                       control.method.list = list(name = "Stephenson", s = 6))
  ub <- method_combine(Z = 1 - Z_dr, Y = Y_dr, N = n_dr, k_vec = 1:n_dr,
                       alpha = ALPHA / 2, simul = TRUE, nperm = NPERM,
                       treat.method.list   = list(name = "Stephenson", s = 6),
                       control.method.list = list(name = "Stephenson", s = 6))
  data.frame(k = 1:n_dr, lower = lb$lower, upper = -ub$lower[n_dr:1])
})

plot_two_sided_CIs(ci_dr$lower, ci_dr$upper)

lv_dr  <- c(0.1, 0.25, 0.5, 0.75, 0.9, 1.0)
idx_dr <- ceiling(lv_dr * n_dr)
knitr::kable(
  data.frame(`Quantile level` = paste0(100 * lv_dr, "%"),
             k = idx_dr,
             `Lower limit` = round(ci_dr$lower[idx_dr], 3),
             `Upper limit` = round(ci_dr$upper[idx_dr], 3),
             check.names = FALSE),
  caption = sprintf(
    "%.0f%% two-sided CIs for ITE quantiles, Arm 5 vs Arm 1 (n = %d)",
    100 * (1 - ALPHA), n_dr))

90% two-sided CIs for ITE quantiles, Arm 5 vs Arm 1 (n = 582)

Quantile level	k	Lower limit	Upper limit
10%	59	-Inf	0.007
25%	146	-Inf	0.035
50%	291	-0.140	0.088
75%	437	-0.078	Inf
90%	524	-0.044	Inf
100%	582	-0.019	Inf

How many drivers actually reduced their handheld use?

Because lower values indicate improvement, the relevant column is “N with \(\tau_i < c\)”: the number of drivers whose handheld use fell by at least -c of drive time. The thresholds below are on the proportion scale, so -0.02 is a two percentage point reduction in the share of drive time spent on a handheld phone.

drive_thresholds <- c(0, -0.02, -0.05)

upper_dr <- -ci_ub_dr$lower[n_dr:1]

below_dr <- sapply(drive_thresholds, function(cc) {
  kk <- which(upper_dr < cc); if (length(kk)) max(kk) else 0L })

knitr::kable(
  data.frame(
    `Threshold c`    = drive_thresholds,
    `N with tau < c` = below_dr,
    `% of drivers`   = sprintf("%.1f%%", 100 * below_dr / n_dr),
    check.names = FALSE),
   align = c("l", "l", "r"),
  col.names = c("Threshold $c$", "$N$ with $\\tau_i < c$", "% of drivers"),
  caption = sprintf(
    paste("%.0f%% lower bounds on the NUMBER of drivers with ITE < c (n = %d)"),
    100 * (1 - ALPHA), n_dr))

90% lower bounds on the NUMBER of drivers with ITE < c (n = 582)

Threshold \(c\)	\(N\) with \(\tau_i < c\)	% of drivers
0.00	37	6.4%
-0.02	0	0.0%
-0.05	0	0.0%

Read the c = 0 row as: with 90% confidence, at least 37 of the 582 drivers reduced their handheld use at all. The rows below it certify progressively larger reductions for progressively fewer drivers.

---

Appendix.

Cache status for this knit

Which results were recomputed on this run and which were reloaded from cache/.

knitr::kable(cache_report(),
             caption = paste0("Cache status (REFRESH_ALL = ", REFRESH_ALL, ")"))

Cache status (REFRESH_ALL = FALSE)

Result	Status	Seconds
bcf_fit_binary	reloaded	0
cate_by_risk_binary	reloaded	0
causal_forest_binary	reloaded	0
ite_onesided_lower	reloaded	0
ite_onesided_upper	reloaded	0
ite_onesided_upper_drive_allrows	reloaded	0
ite_twosided	reloaded	0
ite_twosided_drive	reloaded	0
prior_results_binary	reloaded	0
step1_bart_binary	reloaded	0
vivi_vsup_rd	reloaded	0
vivi_vsup_rr	reloaded	0
xgb_slearner_binary	reloaded	0
xgb_tlearner_binary	reloaded	0