Tutorial

Running the HTE pipeline on your own trial

This walks through the pipeline from a standing start: what to install, what your data has to look like, and what each part produces.

1. What you need

Package installation

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

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

Source files

Both sit next to the .Rmd and are loaded by it.

File	What it provides
`clusterfunctions.R`	`lbart.cluster()` and `wbart.cluster()` — run several BART chains in parallel and collect the posterior draws across chains into one array
`helper_functions.R`	`method_combine()` and the rest of the ITE-quantile machinery

Download both into your working directory:

## clusterfunctions.R — from this site
download.file("https://zhe-chen-1999.github.io/hte-pipeline-site/code/clusterfunctions.R",
              destfile = "clusterfunctions.R")

## helper_functions.R — from the ITE-quantile paper's repository
download.file(
  paste0("https://raw.githubusercontent.com/Zhe-Chen-1999/",
         "Enhanced_inference_for_ITE_quantiles/main/helper_functions.R"),
  destfile = "helper_functions.R")

2. What your data has to look like

One row per randomized patient, with no missing values in the analysis variables:

Object	Type	Meaning
`y`	numeric	the outcome — `0/1` for binary, any numeric for continuous
`z`	logical	treatment assignment, `TRUE` = treated
`X`	data.frame	baseline covariates

Parts of the pipeline convert these internally: z_num <- as.numeric(z) is created for Bayesian Causal Forest, causal forest, and the propensity models, which require a numeric 0/1 indicator, and model.matrix() expands X into a numeric design matrix for XGBoost and BCF.

3. The one block you edit

Everything downstream is driven by a single configuration chunk near the top:

## ============================ EDIT THIS BLOCK =============================
## ---- Outcome type ----
OUTCOME_TYPE <- "binary"       # "binary" or "continuous"

## ---- Execution & Caching Parameters ----
GLOBAL_SEED <- 123
CACHE_DIR   <- "cache"   # expensive results are cached here
REFRESH_ALL <- FALSE     # TRUE recomputes every cached result on this knit

## ---- 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: Heterogeneity analysis by prognosis score
## =========================================================================

What each one does:

Setting	Effect
`OUTCOME_TYPE`	Governs Parts I–III: the BART likelihood (logit vs. Gaussian), the XGBoost objective, the prognostic-score family, and the summary tables. Risk ratio output appears only for a binary outcome
`GLOBAL_SEED`	Seeds the BART chains, the cross-validation folds, and the unit shuffle in Part IV
`CACHE_DIR`	Where cached results are written
`REFRESH_ALL`	`TRUE` recomputes everything on this knit, ignoring the cache
`RUN_PRIOR_GRID`	`FALSE` skips §2.1, the eight-configuration BART prior grid
`RUN_CROSS_METHOD`	`FALSE` skips §2.2, the BCF / causal forest / XGBoost comparison
`RUN_BASELINE_RISK`	`FALSE` skips all of Part III

How the cache works

A full run takes hours. To iterate on prose quickly, each expensive result is wrapped in a small helper:

step1.res <- cached("step1_bart_binary", {
  bart_fit_binary_outcome(X, y, z, ...)     # only runs if the file is absent
})

The first time through, the block runs and the result is written to cache/step1_bart_binary.rds. Every later knit finds that file and reloads it in seconds. The appendix prints a table of what was computed and what was reloaded on that run.

The one thing to watch

The cache is keyed on the name, not on the inputs. It has no idea what went into producing the result. Change your data or a seed and the old result is silently reloaded without warning.

So after any such change, set REFRESH_ALL <- TRUE, knit once, then set it back.

To recompute one result rather than all of them, either pass refresh = TRUE to that call:

step1.res <- cached("step1_bart_binary", { ... }, refresh = TRUE)

or delete its file, which has the same effect:

rm cache/step1_bart_binary.rds

And if you change an analysis in a way that makes the old result meaningless, give it a new name so the stale file cannot be loaded by accident.

4. Part I — a “fit-the-fit” approach

§1.1 fits BART to the outcome and returns a posterior conditional average treatment effect for every patient:

step1.res <- bart_fit_binary_outcome(
  X, y, z,
  Power  = 3, Base = 0.25, Ntrees = 50,   # selected by 10-fold cross-validation
  seed = GLOBAL_SEED, nchains = 4,
  ndpost = 1000, nskip = 500
)
cate.res <- step1.res$cate_results

Passing vectors rather than scalars for Power/Base/Ntrees triggers the cross-validation search.

This section also reports MCMC diagnostics for each chain, including autocorrelation of the fitted response surface to assess mixing and Geweke statistics to assess stationarity. Slow decay in the autocorrelation function (ACF) indicates strong serial dependence and potentially poor mixing, and may motivate increasing the thinning interval (BART_KEEPEVERY). Large absolute Geweke Z-statistics indicate potential non-stationarity or lack of convergence and may motivate increasing the burn-in period (BART_NSKIP).

§1.2 fits a classification and regression tree to those estimates, which is where the interpretable subgroups come from:

cart_fit(point_est = cate.res$cate,
         posterior = step1.res$cate_posterior,
         X = X_unscaled, metric = EFFECT_LABEL,
         maxdepth = 3, cp = 0.01)

§1.3 produces VIVI–VSUP heatmaps: the diagonal is a posterior split-inclusion proportion for each covariate, and the off-diagonal is pairwise interaction strength. Colour intensity encodes magnitude and desaturation encodes posterior uncertainty.

5. Part II — sensitivity analyses and cross-method comparison

Two checks on whether the Part I picture survives scrutiny.

§2.1 — BART Prior Sensitivity. BART is refit across a panel of prior configurations spanning under-regularized, default, CV-selected, and over-regularized regimes, by varying tree depth (power, base), leaf shrinkage (k), and ensemble size (ntree).

§2.2 — Cross-Method Comparison. Bayesian Causal Forest, XGBoost T- and S-learners, and causal forest are fitted to the same data for a direct side-by-side comparison of their CATE distributions.

6. Part III — treatment effect heterogeneity by prognostic score

Effect heterogeneity sometimes tracks a patient’s overall prognosis rather than any individual covariate. The prognostic score is the expected untreated outcome given baseline covariates, \[S_0(X) = E[Y(0) \mid X],\] which for a binary outcome is the baseline predicted risk \(P(Y(0)=1 \mid X)\).

6.1 Estimating the score

If a high-quality, externally developed prognostic model exists for your population, use it: it is estimated independently of your trial and sidesteps everything below. When none is available, the pipeline offers three internal approaches and reports all three side by side.

Approach	Estimates	The catch
1. Control arm only	\(E[Y(0)\mid X]\) — the natural prognostic quantity	A prognostic model is fit using control-arm participants only and then applied to predict the untreated outcome for all participants. Control outcomes are used both to fit the score and in the downstream HTE analysis, so the score is in-sample for controls but out-of-sample for treated patients. That asymmetry can manufacture a CATE-by-risk gradient where none exists, especially with many covariates relative to the control-arm size
2. Full sample, no treatment indicator	\(E[Y \mid X]\)	Fits 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 proportion that depends on the trial’s allocation ratio, not baseline prognosis under no treatment
3. Sample splitting	\(E[Y(0)\mid X]\), honestly	Reserves part of the control arm to fit the score and excludes those patients from the HTE analysis. Prevents outcome reuse and makes score estimation and HTE analysis independent, at the cost of a smaller sample for both stages. Best suited to larger trials

## 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")

PROG_FAMILY follows OUTCOME_TYPE — logistic for a binary outcome, Gaussian for a continuous one.

6.2 Effects within prognostic strata

Patients are grouped into quartiles of the score and the observed treatment effect is reported within each, with bootstrap confidence intervals.

6.3 CATE as a function of the prognostic score

BART and BCF are refit using the prognostic score as the sole covariate, assessing whether treatment effect heterogeneity is driven by baseline prognosis. A CART model is also fitted to discover subgroups defined by baseline prognosis alone.

7. Part IV — randomization-based inference for ITE quantiles

The method assumes units have been randomly permuted before analysis.

For each \(k \in\) k_vec, method_combine() yields a one-sided lower confidence interval for \(\tau_{(k)}\), the \(k\)-th smallest individual effect, and the inference is simultaneously valid across all \(k\).

## Shuffle first: the rank statistic breaks ties by index order
set.seed(GLOBAL_SEED)
ord <- sample(length(z))
Z <- as.numeric(z)[ord]; Y <- as.numeric(y_cont)[ord]; n <- length(Z)

## One-sided lower confidence limits for every effect quantile
ci_lower <- method_combine(
  Z = Z, Y = Y, 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)
)

plot_quantile_bound(ci_lower$lower, bound_type = "lower")

One-sided upper confidence intervals are obtained by re-running with the treatment labels switched, then negating and reversing the resulting bounds:

ub <- method_combine(Z = 1 - Z, Y = Y, 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))
upper <- -ub$lower[n:1]

Two-sided intervals combine a lower and an upper limit by Bonferroni correction.

Turning limits into a count

If the lower limit satisfies \(L_k > c\) then \(\tau_{(k)} > c\), and because the effects are sorted, so does every larger one — hence at least \(n - k + 1\) patients have an effect above \(c\):

n_benefit <- sapply(thresholds, function(cc) {
  kk <- which(ci_lower$lower > cc)
  if (length(kk)) n - min(kk) + 1L else 0L
})

Read the direction from the right limit

If a larger outcome is good, benefit is \(\tau_i > 0\) — read it from the lower limits. If a larger outcome is bad, benefit is \(\tau_i < 0\) — read it from the upper limits.

Next steps

See it run end to end on real trial data.
Get in touch if you would like help applying it to your trial.