Using the tidytreatment package with stochtree

This vignette demonstrates an example workflow for heterogeneous treatment effect models using the stochtree package for fitting Bayesian Additive Regression Trees and tidytreatment for investigating the output of such models. The tidytreatment package can also be used with other BART packages, see other vignettes for details.

Simulate data

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


# load packages
library(stochtree)
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 stochtree::bcf, for causal inference with Bayesian additive regression trees (BART) (Hill 2011). For more on BART see Chipman et al. (2010) and Sparapani et al. (2016). The package can be found on CRAN.

We are following the procedure in Hahn et al. (2020) (though there an automated version is possible with stochtree::bcf) where we estimate a propensity score for being assigned to the treatment regime, which improves estimation properties. The procedure is roughly as follows:

Fit ‘variable selection’ model (VS): Regress response variable against all potential confounders (i.e. no treatment variable)
Select a subset of confounders from the VS model which are most associated with the response variable
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
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
dat_covariates <- dat |> select(-y,-z, -subject_id)

vs_bart <- stochtree::bart(X_train = dat_covariates, y_train = dat$y, num_burnin = 5000, num_mcmc = 1000, 
                rfx_group_ids_train = dat$subject_id, 
                random_effects_params = list(model_spec = 'intercept_only')
                )

# subject_id as random effect!

# STEP 2: Variable selection
  # select most important vars from y ~ covariates model
  # Note: This is an overly simple selection mechanism.
  # See package {bartMan} vignettes and https://doi.org/10.52933/jdssv.v4i1.79
  # for discussion of variable importance methods.
covar_ranking <- covariate_importance(vs_bart, X_train = dat_covariates)

var_select <- covar_ranking |> 
  filter(avg_inclusion > max(avg_inclusion) - 2*sd(avg_inclusion)) %>% # at minimum: within 2 sd of max inclusion
  pull(variable)

var_select
#> [1] "c1"  "x1"  "x10" "x5"  "x6"  "x7"

# STEP 3 PS Model: Regress z ~ selected covariates
ps_bart <- bart(general_params = list(outcome_model = OutcomeModel(outcome = "binary")),
                X_train = dat_covariates, y_train = as.logical(dat$z), num_burnin = 5000, 
                num_mcmc = 1000, 
                rfx_group_ids_train = dat$subject_id, 
                random_effects_params = list(model_spec = 'intercept_only')
                )
#> Warning in bart(general_params = list(outcome_model = OutcomeModel(outcome =
#> "binary")), : Global error variance will not be sampled with a probit link as
#> it is fixed at 1

# store propensity score in data
prop_score <- epred_draws(ps_bart) |> 
  summarise(prop = mean(.value))

# Step 4 TE Model: Regress y ~ z + covariates + propensity score
te_bart <- bcf(X_train = dat_covariates, y_train = dat$y, 
               Z_train = dat$z,
               propensity_train = pull(prop_score, prop),
               rfx_group_ids_train = dat$subject_id, 
               num_burnin = 5000, num_mcmc = 1000, 
               random_effects_params = list(model_spec = 'intercept_only'))

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 = cte), 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(cte)) %>%
  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–420. 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.