Add revised CSF experiments (#1083)

Author: erikcs

experiments/csf/README.md

@@ -2,12 +2,12 @@ _This folder has replication files for the paper "Estimating Heterogeneous Treat
 
 * Figure 1: `prediction_comparison.R`
 
-* Table 1 and Table 2, Figure 5 and Figure 6: `simulation_mse_output.R`
+* MSE and classification error simulations: `simulation_mse_output.R`
 
-* Table 3: `simulation_coverage_output.R`
+* 95 % CI coverage table: `simulation_coverage_output.R`
 
-* Figure 2: `simulation_blp.R`
+* Best linear projection simulation: `simulation_blp.R`
 
-* Figure 3, Figure 4, and Table 4: `hiv.R`
+* HIV application: `hiv.R`
 
 These scripts were run using R version 3.5. In addition to `grf` they rely on the additional packages: `"ggplot2", "randomForestSRC", "speff2trial", "texreg", "xtable"`.

experiments/csf/hiv.R

@@ -1,8 +1,11 @@
+# The following script reproduces the HIV application from
+# the manuscript.
 rm(list = ls())
 library(ggplot2)
 library(texreg)
 library(speff2trial) # "ACTG175" data set.
 library(grf)
+library(xtable)
 set.seed(123)
 
 data = ACTG175[ACTG175$arms == 1 | ACTG175$arms == 3, ]
@@ -13,10 +16,10 @@ Y = data$days
 W = as.numeric(data$arms == 1) # W = 0 : ddI, W = 1: ZDV+ddI
 D = data$cens
 
-# Figure 3 - histogram overlaid
+# Overlaid histogram with T.max
 ggplot(data.frame(Y, Censored = factor(D, labels = c("Yes", "No"))), aes(x = Y, fill = Censored)) +
         geom_histogram(alpha = 0.5) +
-        geom_vline(xintercept = 1000, linetype = 2, col = "red") +
+        geom_vline(xintercept = 1000, linetype = 1, col = "red") +
         xlab("Survival time (days)") +
         ylab("Frequency") +
         theme_bw() +
@@ -27,21 +30,37 @@ ggsave("HIV_histogram.pdf", width = 6, height = 5)
 # Truncate Y at Y.max
 Y.max = 1000
 
-cs.forest = causal_survival_forest(X, Y, W, D, horizon = Y.max)
+cs.forest = causal_survival_forest(X, Y, W, D, horizon = Y.max, num.trees = 10000, ci.group.size = 12)
+
+# Estimates and SEs for a random subset of individuals
+idx = sample(nrow(X), 10)
+pp = predict(cs.forest, estimate.variance = TRUE)
+vimp = variable_importance(cs.forest)
+colnames(X)[order(vimp)[1:4]]
+df = data.frame(
+        CATE = pp$predictions[idx],
+        CATE.se = sqrt(pp$variance.estimates[idx]),
+        hemophilia = ifelse(X[idx, "hemo"] == 1, "Yes", "No"),
+        gender = ifelse(X[idx, "gender"] == 1, "Male", "Female"),
+        homosexual.activity = ifelse(X[idx, "homo"] == 1, "Yes", "No"),
+        antiretroviral.history = ifelse(X[idx, "preanti"] == 1, "Experienced", "Naive")
+)
+print(xtable(df[order(df$CATE), ]
+             ), include.rownames = FALSE)
 
 # BLP
 full = best_linear_projection(cs.forest, X)
 age = best_linear_projection(cs.forest, X[, "age", drop = F])
 
 # Same names as in paper
-varnames = c("Constant", "age", "weight", "Karnofsky score",
-             "CD4 count", "CD8 count", "gender", "homosexual activity",
-             "race", "symptomatic status", "intravenous drug use",
-             "hemophilia", "antiretroviral history",
+varnames = c("Constant", "Age", "Weight", "Karnofsky score",
+             "CD4 count", "CD8 count", "Gender", "Homosexual activity",
+             "Race", "Symptomatic status", "Intravenous drug use",
+             "Hemophilia", "Antiretroviral history",
              "CD4 count 20+/-5 weeks", "CD8 count 20+/-5 weeks"
              )
 
-# Table 4
+# BLP Table
 texreg(list(full, age),
        custom.model.names = c("All covariates", "Age only"),
        table = FALSE,
@@ -51,14 +70,18 @@ texreg(list(full, age),
        custom.coef.names = varnames
 )
 
-# Figure 4
+# CATE plot
 X.median <- apply(X, 2, median)
 age.test = seq(min(X$age), max(X$age))
 X.test = matrix(rep(X.median, length(age.test)), length(age.test), byrow = TRUE)
 X.test[, 1] = age.test
-cs.pred = predict(cs.forest, X.test)
+cs.pred = predict(cs.forest, X.test, estimate.variance = TRUE)
 pt = cs.pred$predictions
+ub = pt + sqrt(cs.pred$variance.estimates) * qnorm(0.975)
+lb = pt - sqrt(cs.pred$variance.estimates) * qnorm(0.975)
 pdf("HIV_data.pdf")
-plot(X.test[, 1], pt, type = 'l', xlab = "Age (years)", ylab ="CATE (days)")
+plot(X.test[, 1], pt, type = 'l', xlab = "Age (years)", ylab = "CATE (days)", ylim = c(min(lb), max(ub)))
+lines(X.test[, 1], ub, lty = 2)
+lines(X.test[, 1], lb, lty = 2)
 grid()
 dev.off()

experiments/csf/prediction_comparison.R

@@ -15,15 +15,15 @@ estimators = list(SRC = SRC1,
                   CSF = CSF)
 
 out = list()
-n = 2000
-p = 5
+n = 5000
+p = 15
 n.test = 2000
 dgp = "type2"
 # dgp = "type3"
-n.mc = 100000
 
-data = generate_causal_survival_data(n = n, p = p, dgp = dgp, n.mc = 10)
-data.test = generate_causal_survival_data(n = n.test, p = p, dgp = dgp, n.mc = n.mc)
+data = generate_causal_survival_data(n = n, p = p, dgp = dgp, n.mc = 1)
+data$Y = round(data$Y, 2)
+data.test = generate_causal_survival_data(n = n.test, p = p, dgp = dgp, n.mc = 100000)
 true.cate = data.test$cate
 for (j in 1:length(estimators)) {
   estimator = names(estimators)[j]
@@ -49,11 +49,11 @@ out$label = factor(out$label, levels = unique(out$label)[c(2, 1, 3)])
 
 ggplot(out, aes(y = predictions, x = true.cate)) +
   geom_point(size = 0.1) +
-  geom_abline(intercept = 0, slope = 1, col = "red", lty = 3) +
+  geom_abline(intercept = 0, slope = 1, col = "red", lty = 1) +
   facet_wrap(. ~ label, ncol = 3) +
   theme_bw() +
-  xlab("True effect") +
-  # xlab("") +
+  # xlab("True effect") +
+  xlab("") +
   ylab("Estimated effect")
 
 ggsave(paste0("prediction_comparsion_", dgp, ".pdf"), width = 6, height = 3)

experiments/csf/simulation_blp.R

@@ -4,13 +4,12 @@ library(grf)
 set.seed(123)
 
 n = 2000
-p = 5
+p = 15
 dgp = "type3"
-nreps = 200
+nreps = 500
 
 # ground truth
-n.test = 50000
-data.test = generate_causal_survival_data(n.test, p, dgp=dgp, n.mc = 50000)
+data.test = generate_causal_survival_data(50000, p, dgp=dgp, n.mc = 50000)
 df = data.frame(cate=data.test$cate, x=data.test$X)
 lm1 = coeftest(lm(cate ~ x.1 + x.2, df))
 true = lm1[2, 1]
@@ -60,15 +59,15 @@ res = replicate(nreps, {
 # Figure 2:
 res.cov = round(rowMeans(res), 2)
 pdf("blp_simulation.pdf")
-breaks = 7
+breaks = 20
 par(mfrow = c(2, 2))
 hist(res["blp.cate", ], breaks = breaks, main = "BLP (CATE)", xlab = paste("coverage: ", res.cov["cov.blp.cate"]))
-abline(v=true, col = "red", lty = 2)
+abline(v=true, col = "red", lty = 1)
 hist(res["blp.dr", ], breaks = breaks, main = "BLP (DR)", xlab = paste("coverage: ", res.cov["cov.blp.dr"]))
-abline(v=true, col = "red", lty = 2)
+abline(v=true, col = "red", lty = 1)
 
 hist(res["ate", ], breaks = breaks, main = "ATE (CATE)" , xlab = paste("coverage: ", res.cov["cov.ate"]))
-abline(v=true.ate, col = "red", lty = 2)
+abline(v=true.ate, col = "red", lty = 1)
 hist(res["ate.dr", ], breaks = breaks, main = "ATE (DR)", xlab = paste("coverage: ", res.cov["cov.ate.dr"]))
-abline(v=true.ate, col = "red", lty = 2)
+abline(v=true.ate, col = "red", lty = 1)
 dev.off()

experiments/csf/simulation_coverage.R

@@ -4,10 +4,9 @@ set.seed(123)
 
 out = list()
 n.sim = 1000
-n.mc = 100000
-p = 5
-X.test = matrix(c(0.2, 0.4, 0.6, 0.8), 4, p)
 grid = expand.grid(n = c(2000),
+                   p = 15,
+                   rho = c(0, 0.5),
                    num.trees = c(10000),
                    dgp = c("type1", "type2", "type3", "type4"),
                    stringsAsFactors = FALSE)
@@ -17,16 +16,19 @@ for (i in 1:nrow(grid)) {
   print(paste("grid", i, "of", nrow(grid)))
   print(grid[i, ])
   n = grid$n[i]
+  p = grid$p[i]
   dgp = grid$dgp[i]
+  rho = grid$rho[i]
   num.trees = grid$num.trees[i]
+  X.test = matrix(c(0.2, 0.4, 0.6, 0.8), 4, p)
 
+  data.test = generate_causal_survival_data(n = nrow(X.test), p = p, X = X.test, dgp = dgp, rho = rho, n.mc = 100000)
+  cate.true = data.test$cate
+  cate.true.prob = data.test$cate.prob
   for (sim in 1:n.sim) {
     print(paste("sim", sim))
-    data = generate_causal_survival_data(n = n, p = p, dgp = dgp, n.mc = 10)
+    data = generate_causal_survival_data(n = n, p = p, dgp = dgp, rho = rho, n.mc = 1)
     data$Y = round(data$Y, 2)
-    data.test = generate_causal_survival_data(n = nrow(X.test), p = p, X = X.test, dgp = dgp, n.mc = n.mc)
-    cate.true = data.test$cate
-    cate.true.prob = data.test$cate.prob
     forest.W = regression_forest(data$X, data$W, num.trees = 500, ci.group.size = 1)
     W.hat = predict(forest.W)$predictions
 
@@ -52,7 +54,9 @@ for (i in 1:nrow(grid)) {
       coverage = c(coverage, coverage.prob),
       width = c(width, width.prob),
       n = n,
+      p = p,
       dgp = dgp,
+      rho = rho,
       num.trees = num.trees,
       sim = sim,
       X.test = X.test[, 1]

experiments/csf/simulation_coverage_output.R

@@ -1,19 +1,41 @@
 # Run `simulation_coverage.R` to produce `coverage.csv.gz`.
-# Table 3 is produced below.
+# 95 % CI coverage table is produced below.
 
 rm(list = ls())
 library(xtable)
 df = read.csv("coverage.csv.gz")
-apply(df[c("target", "n", "num.trees", "dgp", "X.test")], 2, unique)
+apply(df[c("target", "n", "p", "rho", "num.trees", "dgp", "X.test")], 2, unique)
 
-tab = aggregate(list(coverage = df$coverage),
+tab = aggregate(list(coverage = df$coverage, width = df$width),
                 by = list(target = df$target,
                           dgp = df$dgp,
                           Xi = df$X.test,
+                          p = df$p,
+                          rho = df$rho,
                           n.train = df$n,
                           num.trees = df$num.trees),
                 FUN = mean)
 
-# Table 3
-options(digits = 2)
-xtabs(coverage ~ dgp + Xi + target, tab)
+# Table coverage and CI length
+# RMST
+print(xtable(
+  cbind(xtabs(coverage ~ dgp + Xi, tab, subset = target == "RMST" & rho == 0),
+      xtabs(width ~ dgp + Xi, tab, subset = target == "RMST" & rho == 0))
+))
+# SP
+print(xtable(
+  cbind(xtabs(coverage ~ dgp + Xi, tab, subset = target == "survival.probability" & rho == 0),
+        xtabs(width ~ dgp + Xi, tab, subset = target == "survival.probability" & rho == 0))
+))
+
+# w correlated X's
+# RMST
+print(xtable(
+  cbind(xtabs(coverage ~ dgp + Xi, tab, subset = target == "RMST" & rho == 0.5),
+        xtabs(width ~ dgp + Xi, tab, subset = target == "RMST" & rho == 0.5))
+))
+# SP
+print(xtable(
+  cbind(xtabs(coverage ~ dgp + Xi, tab, subset = target == "survival.probability" & rho == 0.5),
+        xtabs(width ~ dgp + Xi, tab, subset = target == "survival.probability" & rho == 0.5))
+))

experiments/csf/simulation_mse.R

@@ -14,9 +14,9 @@ estimators = list(SRC1 = SRC1,
 # *** Setup ***
 out = list()
 n.sim = 250
-n.mc = 100000
 grid = expand.grid(n = c(500, 1000, 2000, 5000),
-                   p = 5,
+                   p = 15,
+                   rho = c(0, 0.5),
                    n.test = 2000,
                    dgp = c("type1", "type2", "type3", "type4"),
                    stringsAsFactors = FALSE)
@@ -29,15 +29,16 @@ for (i in 1:nrow(grid)) {
   p = grid$p[i]
   n.test = grid$n.test[i]
   dgp = grid$dgp[i]
+  rho = grid$rho[i]
 
+  data.test = generate_causal_survival_data(n = n.test, p = p, dgp = dgp, rho = rho, n.mc = 100000)
+  true.cate = data.test$cate
+  true.cate.prob = data.test$cate.prob
+  true.cate.sign = data.test$cate.sign
   for (sim in 1:n.sim) {
     print(paste("sim", sim))
-    data = generate_causal_survival_data(n = n, p = p, dgp = dgp, n.mc = 10)
+    data = generate_causal_survival_data(n = n, p = p, dgp = dgp, rho = rho, n.mc = 1)
     data$Y = round(data$Y, 2)
-    data.test = generate_causal_survival_data(n = n.test, p = p, dgp = dgp, n.mc = n.mc)
-    true.cate = data.test$cate
-    true.cate.prob = data.test$cate.prob
-    true.cate.sign = data.test$cate.sign
     estimator.output = list()
     for (j in 1:length(estimators)) {
       estimator = names(estimators)[j]
@@ -57,6 +58,7 @@ for (i in 1:nrow(grid)) {
     df$p = p
     df$n.test = n.test
     df$dgp = dgp
+    df$rho = rho
     df$sim = sim
 
     out = c(out, list(df))