Using the tidytreatment package with bartCause

Joshua J Bon

2025-01-10

This vignette demonstrates an example workflow for heterogeneous treatment effect models using the BART package for fitting Bayesian Additive Regression Trees and tidytreatment for investigating the output of such models. The tidytreatment package can also be used with bartMachine models, support for bcf is coming soon (see branch bcf-hold on github).

Simulate data

Below we load packages and simulate data using the scheme described by J. Hill and Su (2013) with the additional of 1 categorical variable. It it implemented in the function simulate_hill_su_data():


# load packages
library(bartCause)
library(stan4bart)
library(tidytreatment)
library(dplyr)
library(tidybayes)
library(ggplot2)

# set seed so vignette is reproducible
set.seed(101)

# simulate data
sim <- simulate_su_hill_data(n = 100, treatment_linear = FALSE,  omega = 0, add_categorical = TRUE,
                             n_subjects = 10, sd_subjects = 0.1,
                             coef_categorical_treatment = c(0,0,1),
                             coef_categorical_nontreatment = c(-1,0,-1)
)

Now we can take a look at some data summaries.


# non-treated vs treated counts:
table(sim$data$z)
#> 
#>  0  1 
#> 48 52

dat <- sim$data
# a selection of data
dat %>% select(y, z, c1, x1:x3) %>% head()
#>            y z c1         x1          x2         x3
#> 1  7.4280183 1  3  0.1820194 -0.06353599  0.3524953
#> 2  0.5275992 0  1 -2.3814015  0.20423903  0.8935155
#> 3  3.5616838 0  2 -0.5518608 -1.10040685  1.9034043
#> 4  9.3384301 1  2  0.7267062 -0.11821280 -0.5202530
#> 5  3.7328197 1  1 -1.8133403  0.73768934  0.2913854
#> 6 -1.7501469 0  2 -0.8688956 -0.38677915 -0.7522927

# repeated observation counts for subjects:
table(sim$data$subject_id)
#> 
#>  1  2  3  4  5  6  7  8  9 10 
#>  8 10  6 14 14 10  9  8 13  8

Fit the regression model

Run the model to be used to assess treatment effects. Here we will use the bartCause, for causal inference with Bayesian additive regression trees (BART) (J. L. Hill 2011). For more on BART see Chipman, George, and E. McCulloch (2010) and Sparapani et al. (2016). The package can be found on CRAN.

We are following the procedure in Hahn, Murray, and Carvalho (2020) (albeit without their causal forest model) where we estimate a propensity score for being assigned to the treatment regime, which improves estimation properties. This is done automatically in bartCause. The procedure is roughly as follows:

  1. Fit ‘variable selection’ model (VS): Regress response variable against all potential confounders (i.e. no treatment variable)
  2. Select a subset of confounders from the VS model which are most associated with the response variable
  3. Fit a ‘propensity score’ model (PS): A binary response model estimating the treatment assignment propensity score for using only the variables selected in step 2
  4. Fit the treatment effect model (TE): Using the all potential confounders and propensity score from step 3

Note: Using the automatic treatment assignment model in bartCause ignores step 1-2 uses all potential confounders to estimate the PS model.

  
# STEP 1 VS Model: Regress y ~ covariates
vs_bart <- stan4bart(y ~ bart(. - subject_id - z) + (1|subject_id), 
                             data = dat, iter = 5000, verbose = -1)

# STEP 2: Variable selection
  # select most important vars from y ~ covariates model
  # very simple selection mechanism. Should use cross-validation in practice
covar_ranking <- covariate_importance(vs_bart)
var_select <- covar_ranking %>% 
  filter(avg_inclusion > mean(avg_inclusion) - sd(avg_inclusion)) %>% # at minimum: within 1 sd of mean inclusion
  pull(variable)

# change categorical variables to just one variable
var_select <- unique(gsub("c1.[1-3]$","c1", var_select))

var_select
#>  [1] "c1"  "x1"  "x2"  "x3"  "x4"  "x5"  "x6"  "x7"  "x8"  "x9"  "x10"
# includes all covariates

# STEP 3 PS Model: Regress z ~ selected covariates
ps_bart <- stan4bart(z ~ bart(. - subject_id - y) + (1|subject_id), 
                             data = dat, iter = 5000, verbose = -1)

# store propensity score in data
prop_score <- fitted(ps_bart)

# Step 4 TE Model: Regress y ~ z + covariates + propensity score
te_bart <- bartc(response = y, treatment = z, confounders = x1 + x2 + x3 + x4 + x5 + x6 + x7 + x8 + x9 + x10,  
                 parametric = (1|subject_id), data = dat, method.trt = prop_score, 
                 iter = 5000, bart_args = list(keepTrees = TRUE))
#> fitting response model via method 'bart'

#* The posterior samples are kept small to manage size on CRAN

Extract the posterior (tidy style)

Methods for extracting the posterior in a tidy format is included in the tidytreatment.


# get model parameters (excluding BART paramaters)
posterior_params <- tidy_draws(te_bart)

posterior_fitted <- epred_draws(te_bart, value = "fitted")

# Function to tidy predicted draws (adds predicted noise to fitted values)
posterior_pred <- predicted_draws(te_bart, value = "predicted")

Use some plotting functions from the tidybayes package

Since tidytreatment follows the tidybayes output specifications, functions from tidybayes should work.


treatment_var_and_c1 <- 
  dat %>% 
  select(z,c1) %>%
  mutate(.row = 1:n(), z = as.factor(z))

posterior_fitted %>%
  left_join(treatment_var_and_c1, by = ".row") %>%
  ggplot() + 
  stat_halfeye(aes(x = z, y = fitted)) + 
  facet_wrap(~c1, labeller = as_labeller( function(x) paste("c1 =",x) ) ) +
  xlab("Treatment (z)") + ylab("Posterior predicted value") +
  theme_bw() + ggtitle("Effect of treatment with 'c1' on posterior fitted values")

Calculate Treatment Effects

Posterior conditional (average) treatment effects can be calculated using the treatment_effects function. This function finds the posterior values of \[ \tau(x) = \text{E}(y ~ \vert~ T = 1, X = x) - \text{E}(y ~ \vert~ T = 0, X = x) \] for each unit of measurement, \(i\), (e.g. subject) in the data sample.

Some histogram summaries are presented below.


# sample based (using data from fit) conditional treatment effects, posterior draws
posterior_treat_eff <- treatment_effects(te_bart)

# check lines up with summary results...

# Histogram of treatment effect (all draws)
posterior_treat_eff %>% 
  ggplot() +
  geom_histogram(aes(x = icate), binwidth = 0.1, colour = "white") + 
  theme_bw() + ggtitle("Histogram of treatment effect (all draws)")


# Histogram of treatment effect (median for each subject)
posterior_treat_eff %>% summarise(cte_hat = median(icate)) %>%
  ggplot() +
  geom_histogram(aes(x = cte_hat), binwidth = 0.1, colour = "white") + 
  theme_bw() + ggtitle("Histogram of treatment effect (median for each subject)")

References

Chipman, Hugh A., Edward I. George, and Robert E. McCulloch. 2010. “BART: Bayesian Additive Regression Trees.” The Annals of Applied Statistics 4 (1): 266–98. https://doi.org/10.1214/09-AOAS285.
Hahn, P. Richard, Jared S. Murray, and Carlos M. Carvalho. 2020. “Bayesian Regression Tree Models for Causal Inference: Regularization, Confounding, and Heterogeneous Effects (with Discussion).” Bayesian Analysis 15 (3): 965–1056. https://doi.org/10.1214/19-BA1195.
Hill, Jennifer L. 2011. “Bayesian Nonparametric Modeling for Causal Inference.” Journal of Computational and Graphical Statistics 20 (1): 217–40. https://doi.org/10.1198/jcgs.2010.08162.
Hill, Jennifer, and Yu-Sung Su. 2013. “Assessing Lack of Common Support in Causal Inference Using Bayesian Nonparametrics: Implications for Evaluating the Effect of Breastfeeding on Children’s Cognitive Outcomes.” The Annals of Applied Statistics 7 (3): 1386–1420. https://doi.org/10.1214/13-AOAS630.
Sparapani, Rodney A, Brent R Logan, Robert E McCulloch, and Purushottam W Laud. 2016. “Nonparametric Survival Analysis Using Bayesian Additive Regression Trees (BART).” Statistics in Medicine 35 (16): 2741–53. https://doi.org/10.1002/sim.6893.