Updated R vignettes

Author: andrewherren

vignettes/CausalInference.Rmd

@@ -134,9 +134,9 @@ plot(rowMeans(bcf_model_warmstart$tau_hat_test), tau_test,
      xlab = "predicted", ylab = "actual", main = "Treatment effect")
 abline(0,1,col="red",lty=3,lwd=3)
 sigma_observed <- var(y-E_XZ)
-plot_bounds <- c(min(c(bcf_model_warmstart$sigma2_samples, sigma_observed)), 
-                 max(c(bcf_model_warmstart$sigma2_samples, sigma_observed)))
-plot(bcf_model_warmstart$sigma2_samples, ylim = plot_bounds, 
+plot_bounds <- c(min(c(bcf_model_warmstart$sigma2_global_samples, sigma_observed)), 
+                 max(c(bcf_model_warmstart$sigma2_global_samples, sigma_observed)))
+plot(bcf_model_warmstart$sigma2_global_samples, ylim = plot_bounds, 
      ylab = "sigma^2", xlab = "Sample", main = "Global variance parameter")
 abline(h = sigma_observed, lty=3, lwd = 3, col = "blue")
 ```
@@ -184,9 +184,9 @@ plot(rowMeans(bcf_model_root$tau_hat_test), tau_test,
      xlab = "predicted", ylab = "actual", main = "Treatment effect")
 abline(0,1,col="red",lty=3,lwd=3)
 sigma_observed <- var(y-E_XZ)
-plot_bounds <- c(min(c(bcf_model_root$sigma2_samples, sigma_observed)), 
-                 max(c(bcf_model_root$sigma2_samples, sigma_observed)))
-plot(bcf_model_root$sigma2_samples, ylim = plot_bounds, 
+plot_bounds <- c(min(c(bcf_model_root$sigma2_global_samples, sigma_observed)), 
+                 max(c(bcf_model_root$sigma2_global_samples, sigma_observed)))
+plot(bcf_model_root$sigma2_global_samples, ylim = plot_bounds, 
      ylab = "sigma^2", xlab = "Sample", main = "Global variance parameter")
 abline(h = sigma_observed, lty=3, lwd = 3, col = "blue")
 ```
@@ -303,9 +303,9 @@ plot(rowMeans(bcf_model_warmstart$tau_hat_test), tau_test,
      xlab = "predicted", ylab = "actual", main = "Treatment effect")
 abline(0,1,col="red",lty=3,lwd=3)
 sigma_observed <- var(y-E_XZ)
-plot_bounds <- c(min(c(bcf_model_warmstart$sigma2_samples, sigma_observed)), 
-                 max(c(bcf_model_warmstart$sigma2_samples, sigma_observed)))
-plot(bcf_model_warmstart$sigma2_samples, ylim = plot_bounds, 
+plot_bounds <- c(min(c(bcf_model_warmstart$sigma2_global_samples, sigma_observed)), 
+                 max(c(bcf_model_warmstart$sigma2_global_samples, sigma_observed)))
+plot(bcf_model_warmstart$sigma2_global_samples, ylim = plot_bounds, 
      ylab = "sigma^2", xlab = "Sample", main = "Global variance parameter")
 abline(h = sigma_observed, lty=3, lwd = 3, col = "blue")
 ```
@@ -353,9 +353,9 @@ plot(rowMeans(bcf_model_root$tau_hat_test), tau_test,
      xlab = "predicted", ylab = "actual", main = "Treatment effect")
 abline(0,1,col="red",lty=3,lwd=3)
 sigma_observed <- var(y-E_XZ)
-plot_bounds <- c(min(c(bcf_model_root$sigma2_samples, sigma_observed)), 
-                 max(c(bcf_model_root$sigma2_samples, sigma_observed)))
-plot(bcf_model_root$sigma2_samples, ylim = plot_bounds, 
+plot_bounds <- c(min(c(bcf_model_root$sigma2_global_samples, sigma_observed)), 
+                 max(c(bcf_model_root$sigma2_global_samples, sigma_observed)))
+plot(bcf_model_root$sigma2_global_samples, ylim = plot_bounds, 
      ylab = "sigma^2", xlab = "Sample", main = "Global variance parameter")
 abline(h = sigma_observed, lty=3, lwd = 3, col = "blue")
 ```
@@ -472,9 +472,9 @@ plot(rowMeans(bcf_model_warmstart$tau_hat_test), tau_test,
      xlab = "predicted", ylab = "actual", main = "Treatment effect")
 abline(0,1,col="red",lty=3,lwd=3)
 sigma_observed <- var(y-E_XZ)
-plot_bounds <- c(min(c(bcf_model_warmstart$sigma2_samples, sigma_observed)), 
-                 max(c(bcf_model_warmstart$sigma2_samples, sigma_observed)))
-plot(bcf_model_warmstart$sigma2_samples, ylim = plot_bounds, 
+plot_bounds <- c(min(c(bcf_model_warmstart$sigma2_global_samples, sigma_observed)), 
+                 max(c(bcf_model_warmstart$sigma2_global_samples, sigma_observed)))
+plot(bcf_model_warmstart$sigma2_global_samples, ylim = plot_bounds, 
      ylab = "sigma^2", xlab = "Sample", main = "Global variance parameter")
 abline(h = sigma_observed, lty=3, lwd = 3, col = "blue")
 ```
@@ -522,9 +522,9 @@ plot(rowMeans(bcf_model_root$tau_hat_test), tau_test,
      xlab = "predicted", ylab = "actual", main = "Treatment effect")
 abline(0,1,col="red",lty=3,lwd=3)
 sigma_observed <- var(y-E_XZ)
-plot_bounds <- c(min(c(bcf_model_root$sigma2_samples, sigma_observed)), 
-                 max(c(bcf_model_root$sigma2_samples, sigma_observed)))
-plot(bcf_model_root$sigma2_samples, ylim = plot_bounds, 
+plot_bounds <- c(min(c(bcf_model_root$sigma2_global_samples, sigma_observed)), 
+                 max(c(bcf_model_root$sigma2_global_samples, sigma_observed)))
+plot(bcf_model_root$sigma2_global_samples, ylim = plot_bounds, 
      ylab = "sigma^2", xlab = "Sample", main = "Global variance parameter")
 abline(h = sigma_observed, lty=3, lwd = 3, col = "blue")
 ```
@@ -639,9 +639,9 @@ plot(rowMeans(bcf_model_warmstart$tau_hat_test), tau_test,
      xlab = "predicted", ylab = "actual", main = "Treatment effect")
 abline(0,1,col="red",lty=3,lwd=3)
 sigma_observed <- var(y-E_XZ)
-plot_bounds <- c(min(c(bcf_model_warmstart$sigma2_samples, sigma_observed)), 
-                 max(c(bcf_model_warmstart$sigma2_samples, sigma_observed)))
-plot(bcf_model_warmstart$sigma2_samples, ylim = plot_bounds, 
+plot_bounds <- c(min(c(bcf_model_warmstart$sigma2_global_samples, sigma_observed)), 
+                 max(c(bcf_model_warmstart$sigma2_global_samples, sigma_observed)))
+plot(bcf_model_warmstart$sigma2_global_samples, ylim = plot_bounds, 
      ylab = "sigma^2", xlab = "Sample", main = "Global variance parameter")
 abline(h = sigma_observed, lty=3, lwd = 3, col = "blue")
 ```
@@ -689,9 +689,9 @@ plot(rowMeans(bcf_model_root$tau_hat_test), tau_test,
      xlab = "predicted", ylab = "actual", main = "Treatment effect")
 abline(0,1,col="red",lty=3,lwd=3)
 sigma_observed <- var(y-E_XZ)
-plot_bounds <- c(min(c(bcf_model_root$sigma2_samples, sigma_observed)), 
-                 max(c(bcf_model_root$sigma2_samples, sigma_observed)))
-plot(bcf_model_root$sigma2_samples, ylim = plot_bounds, 
+plot_bounds <- c(min(c(bcf_model_root$sigma2_global_samples, sigma_observed)), 
+                 max(c(bcf_model_root$sigma2_global_samples, sigma_observed)))
+plot(bcf_model_root$sigma2_global_samples, ylim = plot_bounds, 
      ylab = "sigma^2", xlab = "Sample", main = "Global variance parameter")
 abline(h = sigma_observed, lty=3, lwd = 3, col = "blue")
 ```
@@ -806,9 +806,9 @@ plot(rowMeans(bcf_model_warmstart$rfx_preds_test), rfx_term_test,
      xlab = "predicted", ylab = "actual", main = "Random effects terms")
 abline(0,1,col="red",lty=3,lwd=3)
 sigma_observed <- var(y-E_XZ-rfx_term)
-plot_bounds <- c(min(c(bcf_model_warmstart$sigma2_samples, sigma_observed)), 
-                 max(c(bcf_model_warmstart$sigma2_samples, sigma_observed)))
-plot(bcf_model_warmstart$sigma2_samples, ylim = plot_bounds, 
+plot_bounds <- c(min(c(bcf_model_warmstart$sigma2_global_samples, sigma_observed)), 
+                 max(c(bcf_model_warmstart$sigma2_global_samples, sigma_observed)))
+plot(bcf_model_warmstart$sigma2_global_samples, ylim = plot_bounds, 
      ylab = "sigma^2", xlab = "Sample", main = "Global variance parameter")
 abline(h = sigma_observed, lty=3, lwd = 3, col = "blue")
 ```
@@ -919,9 +919,9 @@ plot(rowMeans(bcf_model_mcmc$y_hat_test), y_test,
      xlab = "predicted", ylab = "actual", main = "Outcome")
 abline(0,1,col="red",lty=3,lwd=3)
 sigma_observed <- var(y-E_XZ)
-plot_bounds <- c(min(c(bcf_model_mcmc$sigma2_samples, sigma_observed)), 
-                 max(c(bcf_model_mcmc$sigma2_samples, sigma_observed)))
-plot(bcf_model_mcmc$sigma2_samples, ylim = plot_bounds, 
+plot_bounds <- c(min(c(bcf_model_mcmc$sigma2_global_samples, sigma_observed)), 
+                 max(c(bcf_model_mcmc$sigma2_global_samples, sigma_observed)))
+plot(bcf_model_mcmc$sigma2_global_samples, ylim = plot_bounds, 
      ylab = "sigma^2", xlab = "Sample", main = "Global variance parameter")
 abline(h = sigma_observed, lty=3, lwd = 3, col = "blue")
 ```
@@ -981,9 +981,9 @@ plot(rowMeans(bcf_model_mcmc$y_hat_test), y_test,
      xlab = "predicted", ylab = "actual", main = "Outcome")
 abline(0,1,col="red",lty=3,lwd=3)
 sigma_observed <- var(y-E_XZ)
-plot_bounds <- c(min(c(bcf_model_mcmc$sigma2_samples, sigma_observed)), 
-                 max(c(bcf_model_mcmc$sigma2_samples, sigma_observed)))
-plot(bcf_model_mcmc$sigma2_samples, ylim = plot_bounds, 
+plot_bounds <- c(min(c(bcf_model_mcmc$sigma2_global_samples, sigma_observed)), 
+                 max(c(bcf_model_mcmc$sigma2_global_samples, sigma_observed)))
+plot(bcf_model_mcmc$sigma2_global_samples, ylim = plot_bounds, 
      ylab = "sigma^2", xlab = "Sample", main = "Global variance parameter")
 abline(h = sigma_observed, lty=3, lwd = 3, col = "blue")
 ```
@@ -1043,9 +1043,9 @@ plot(rowMeans(bcf_model_warmstart$y_hat_test), y_test,
      xlab = "predicted", ylab = "actual", main = "Outcome")
 abline(0,1,col="red",lty=3,lwd=3)
 sigma_observed <- var(y-E_XZ)
-plot_bounds <- c(min(c(bcf_model_warmstart$sigma2_samples, sigma_observed)), 
-                 max(c(bcf_model_warmstart$sigma2_samples, sigma_observed)))
-plot(bcf_model_warmstart$sigma2_samples, ylim = plot_bounds, 
+plot_bounds <- c(min(c(bcf_model_warmstart$sigma2_global_samples, sigma_observed)), 
+                 max(c(bcf_model_warmstart$sigma2_global_samples, sigma_observed)))
+plot(bcf_model_warmstart$sigma2_global_samples, ylim = plot_bounds, 
      ylab = "sigma^2", xlab = "Sample", main = "Global variance parameter")
 abline(h = sigma_observed, lty=3, lwd = 3, col = "blue")
 ```
@@ -1105,9 +1105,9 @@ plot(rowMeans(bcf_model_warmstart$y_hat_test), y_test,
      xlab = "predicted", ylab = "actual", main = "Outcome")
 abline(0,1,col="red",lty=3,lwd=3)
 sigma_observed <- var(y-E_XZ)
-plot_bounds <- c(min(c(bcf_model_warmstart$sigma2_samples, sigma_observed)), 
-                 max(c(bcf_model_warmstart$sigma2_samples, sigma_observed)))
-plot(bcf_model_warmstart$sigma2_samples, ylim = plot_bounds, 
+plot_bounds <- c(min(c(bcf_model_warmstart$sigma2_global_samples, sigma_observed)), 
+                 max(c(bcf_model_warmstart$sigma2_global_samples, sigma_observed)))
+plot(bcf_model_warmstart$sigma2_global_samples, ylim = plot_bounds, 
      ylab = "sigma^2", xlab = "Sample", main = "Global variance parameter")
 abline(h = sigma_observed, lty=3, lwd = 3, col = "blue")
 ```
@@ -1133,6 +1133,193 @@ test_outcome_mean <- rowMeans(bcf_model_warmstart$y_hat_test)
 sqrt(mean((y_test - test_outcome_mean)^2))
 ```
 
+## Demo 7: Probit Outcome Model, Heterogeneous Treatment Effect
+
+We consider a modified version of a data generating process from @hahn2020bayesian:
+
+\begin{equation*}
+\begin{aligned}
+y &= \mathbb{1}\left(w > 0\right)\\
+w &= \mu(X) + \tau(X) Z + \epsilon\\
+\epsilon &\sim N\left(0,1\right)\\
+\mu(X) &= 1 + g(X) + 6 \lvert X_3 - 1 \rvert\\
+\tau(X) &= 1 + 2 X_2 X_4\\
+g(X) &= \mathbb{I}(X_5=1) \times 2 - \mathbb{I}(X_5=2) \times 1 - \mathbb{I}(X_5=3) \times 4\\
+s_{\mu} &= \sqrt{\mathbb{V}(\mu(X))}\\
+\pi(X) &= 0.8 \phi\left(\frac{3\mu(X)}{s_{\mu}}\right) - \frac{X_1}{2} + \frac{2U+1}{20}\\
+X_1,X_2,X_3 &\sim N\left(0,1\right)\\
+X_4 &\sim \text{Bernoulli}(1/2)\\
+X_5 &\sim \text{Categorical}(1/3,1/3,1/3)\\
+U &\sim \text{Uniform}\left(0,1\right)\\
+Z &\sim \text{Bernoulli}\left(\pi(X)\right)
+\end{aligned}
+\end{equation*}
+
+### Simulation
+
+We draw from the DGP defined above
+
+```{r}
+n <- 2000
+x1 <- rnorm(n)
+x2 <- rnorm(n)
+x3 <- rnorm(n)
+x4 <- as.numeric(rbinom(n,1,0.5))
+x5 <- as.numeric(sample(1:3,n,replace=TRUE))
+X <- cbind(x1,x2,x3,x4,x5)
+p <- ncol(X)
+mu_x <- mu1(X)
+tau_x <- tau2(X)
+pi_x <- 0.8*pnorm((3*mu_x/sd(mu_x)) - 0.5*X[,1]) + 0.05 + runif(n)/10
+Z <- rbinom(n,1,pi_x)
+E_XZ <- mu_x + Z*tau_x
+w <- E_XZ + rnorm(n, 0, 1)
+y <- 1*(w > 0)
+delta_x <- pnorm(mu_x + tau_x) - pnorm(mu_x)
+X <- as.data.frame(X)
+X$x4 <- factor(X$x4, ordered = TRUE)
+X$x5 <- factor(X$x5, ordered = TRUE)
+
+# Split data into test and train sets
+test_set_pct <- 0.2
+n_test <- round(test_set_pct*n)
+n_train <- n - n_test
+test_inds <- sort(sample(1:n, n_test, replace = FALSE))
+train_inds <- (1:n)[!((1:n) %in% test_inds)]
+X_test <- X[test_inds,]
+X_train <- X[train_inds,]
+pi_test <- pi_x[test_inds]
+pi_train <- pi_x[train_inds]
+Z_test <- Z[test_inds]
+Z_train <- Z[train_inds]
+w_test <- w[test_inds]
+w_train <- w[train_inds]
+y_test <- y[test_inds]
+y_train <- y[train_inds]
+mu_test <- mu_x[test_inds]
+mu_train <- mu_x[train_inds]
+tau_test <- tau_x[test_inds]
+tau_train <- tau_x[train_inds]
+delta_x_test <- delta_x[test_inds]
+delta_x_train <- delta_x[train_inds]
+```
+
+### Sampling and Analysis
+
+#### Warmstart
+
+We first simulate from an ensemble model of $y \mid X$ using "warm-start" 
+initialization samples (@krantsevich2023stochastic). This is the default in 
+`stochtree`.
+
+```{r}
+num_gfr <- 10
+num_burnin <- 0
+num_mcmc <- 100
+num_samples <- num_gfr + num_burnin + num_mcmc
+general_params <- list(keep_every = 5, 
+                       probit_outcome_model = T, 
+                       sample_sigma2_global = F)
+prognostic_forest_params <- list(sample_sigma2_leaf = F)
+treatment_effect_forest_params <- list(sample_sigma2_leaf = F)
+bcf_model_warmstart <- bcf(
+    X_train = X_train, Z_train = Z_train, y_train = y_train, propensity_train = pi_train, 
+    X_test = X_test, Z_test = Z_test, propensity_test = pi_test, 
+    num_gfr = num_gfr, num_burnin = num_burnin, num_mcmc = num_mcmc, 
+    general_params = general_params, prognostic_forest_params = prognostic_forest_params, 
+    treatment_effect_forest_params = treatment_effect_forest_params
+)
+```
+
+Inspect the BART samples that were initialized with an XBART warm-start
+
+```{r}
+mu_hat_test <- rowMeans(bcf_model_warmstart$mu_hat_test)
+plot(mu_hat_test, mu_test, xlab = "predicted", 
+     ylab = "actual", main = "Prognostic function")
+abline(0,1,col="red",lty=3,lwd=3)
+tau_hat_test <- rowMeans(bcf_model_warmstart$tau_hat_test)
+plot(tau_hat_test, tau_test, xlab = "predicted", 
+     ylab = "actual", main = "Treatment effect")
+abline(0,1,col="red",lty=3,lwd=3)
+delta_x_hat_test <- pnorm(mu_hat_test+tau_hat_test) - pnorm(mu_hat_test)
+plot(delta_x_hat_test, delta_x_test, 
+     xlab = "predicted", ylab = "actual", main = "Distributional treatment\neffect")
+abline(0,1,col="red",lty=3,lwd=3)
+```
+
+Examine test set interval coverage
+
+```{r}
+test_lb <- apply(
+    pnorm(bcf_model_warmstart$mu_hat_test + bcf_model_warmstart$tau_hat_test) - 
+        pnorm(bcf_model_warmstart$mu_hat_test), 1, quantile, 0.025)
+test_ub <- apply(
+    pnorm(bcf_model_warmstart$mu_hat_test + bcf_model_warmstart$tau_hat_test) - 
+        pnorm(bcf_model_warmstart$mu_hat_test), 1, quantile, 0.975)
+cover <- (
+    (test_lb <= delta_x_test) & 
+    (test_ub >= delta_x_test)
+)
+mean(cover)
+```
+
+#### BART MCMC without Warmstart
+
+Next, we simulate from this ensemble model without any warm-start initialization. 
+
+```{r}
+num_gfr <- 0
+num_burnin <- 2000
+num_mcmc <- 100
+num_samples <- num_gfr + num_burnin + num_mcmc
+general_params <- list(keep_every = 5, 
+                       probit_outcome_model = T, 
+                       sample_sigma2_global = F)
+prognostic_forest_params <- list(sample_sigma2_leaf = F)
+treatment_effect_forest_params <- list(sample_sigma2_leaf = F)
+bcf_model_root <- bcf(
+    X_train = X_train, Z_train = Z_train, y_train = y_train, propensity_train = pi_train, 
+    X_test = X_test, Z_test = Z_test, propensity_test = pi_test, 
+    num_gfr = num_gfr, num_burnin = num_burnin, num_mcmc = num_mcmc, 
+    general_params = general_params, prognostic_forest_params = prognostic_forest_params, 
+    treatment_effect_forest_params = treatment_effect_forest_params
+)
+```
+
+Inspect the BART samples that were initialized with an XBART warm-start
+
+```{r}
+mu_hat_test <- rowMeans(bcf_model_root$mu_hat_test)
+plot(mu_hat_test, mu_test, xlab = "predicted", 
+     ylab = "actual", main = "Prognostic function")
+abline(0,1,col="red",lty=3,lwd=3)
+tau_hat_test <- rowMeans(bcf_model_root$tau_hat_test)
+plot(tau_hat_test, tau_test, xlab = "predicted", 
+     ylab = "actual", main = "Treatment effect")
+abline(0,1,col="red",lty=3,lwd=3)
+delta_x_hat_test <- pnorm(mu_hat_test+tau_hat_test) - pnorm(mu_hat_test)
+plot(delta_x_hat_test, delta_x_test, 
+     xlab = "predicted", ylab = "actual", main = "Distributional treatment\neffect")
+abline(0,1,col="red",lty=3,lwd=3)
+```
+
+Examine test set interval coverage
+
+```{r}
+test_lb <- apply(
+    pnorm(bcf_model_root$mu_hat_test + bcf_model_root$tau_hat_test) - 
+        pnorm(bcf_model_root$mu_hat_test), 1, quantile, 0.025)
+test_ub <- apply(
+    pnorm(bcf_model_root$mu_hat_test + bcf_model_root$tau_hat_test) - 
+        pnorm(bcf_model_root$mu_hat_test), 1, quantile, 0.975)
+cover <- (
+    (test_lb <= delta_x_test) & 
+    (test_ub >= delta_x_test)
+)
+mean(cover)
+```
+
 # Continuous Treatment
 
 ## Demo 1: Nonlinear Outcome Model, Heterogeneous Treatment Effect
@@ -1230,9 +1417,9 @@ plot(rowMeans(bcf_model_warmstart$tau_hat_test), tau_test,
      xlab = "predicted", ylab = "actual", main = "Treatment effect")
 abline(0,1,col="red",lty=3,lwd=3)
 sigma_observed <- var(y-E_XZ)
-plot_bounds <- c(min(c(bcf_model_warmstart$sigma2_samples, sigma_observed)), 
-                 max(c(bcf_model_warmstart$sigma2_samples, sigma_observed)))
-plot(bcf_model_warmstart$sigma2_samples, ylim = plot_bounds, 
+plot_bounds <- c(min(c(bcf_model_warmstart$sigma2_global_samples, sigma_observed)), 
+                 max(c(bcf_model_warmstart$sigma2_global_samples, sigma_observed)))
+plot(bcf_model_warmstart$sigma2_global_samples, ylim = plot_bounds, 
      ylab = "sigma^2", xlab = "Sample", main = "Global variance parameter")
 abline(h = sigma_observed, lty=3, lwd = 3, col = "blue")
 ```
@@ -1280,9 +1467,9 @@ plot(rowMeans(bcf_model_root$tau_hat_test), tau_test,
      xlab = "predicted", ylab = "actual", main = "Treatment effect")
 abline(0,1,col="red",lty=3,lwd=3)
 sigma_observed <- var(y-E_XZ)
-plot_bounds <- c(min(c(bcf_model_root$sigma2_samples, sigma_observed)), 
-                 max(c(bcf_model_root$sigma2_samples, sigma_observed)))
-plot(bcf_model_root$sigma2_samples, ylim = plot_bounds, 
+plot_bounds <- c(min(c(bcf_model_root$sigma2_global_samples, sigma_observed)), 
+                 max(c(bcf_model_root$sigma2_global_samples, sigma_observed)))
+plot(bcf_model_root$sigma2_global_samples, ylim = plot_bounds, 
      ylab = "sigma^2", xlab = "Sample", main = "Global variance parameter")
 abline(h = sigma_observed, lty=3, lwd = 3, col = "blue")
 ```