simple_honesty_compare.R

rm(list=ls())
library(grf)
set.seed(123)

# Data with known treatment propensities = 1/2 (W.hat)
# and expected outcomes = 0 (Y.hat)
get_data = function(n) {
  p = 5
  X = matrix(runif(n * p), nrow = n, ncol = p)
  W.hat = rep(1/2, n)
  Y.hat = rep(0, n)
  W = rbinom(n, 1, 0.5)
  sigma = function(x) 1 / (1 + exp(-4 * (x - 0.5)))
  tau = sigma(X[, 1]) * sigma(X[, 2])
  Y = (W - 1/2) * tau + rnorm(n)

  list(X=X, Y=Y, W=W, Y.hat=Y.hat, W.hat=W.hat, tau=tau)
}

params = list(
  honesty = c(TRUE, FALSE),
  honesty.fraction = c(0.5, 0.25, 0.1),
  prune = c(TRUE, FALSE)
)
arguments = expand.grid(params)
arguments[arguments$honesty == FALSE, 'honesty.fraction'] = NA
arguments[arguments$honesty == FALSE, 'prune'] = FALSE
arguments = arguments[!duplicated(arguments), ]

run_experiment = function(n, args) {
  df = get_data(n)

  out = sapply(1:nrow(args), function (i) {
    honesty.fraction = args[i, 'honesty.fraction']
    if (is.na(honesty.fraction))
      honesty.fraction = NULL

    cf = causal_forest(df$X, df$Y, df$W, Y.hat = df$Y.hat, W.hat = df$W.hat,
                       honesty = args[i, 'honesty'],
                       honesty.fraction = honesty.fraction,
                       prune = args[i, 'prune'])

    tau.hat = predict(cf)$predictions
    mse = mean((df$tau - tau.hat)^2)

    mse
  })

  out
}

nsim = 1000

cat("running sim1\n")
n = 100
a100 = replicate(n = nsim, run_experiment(n = n, arguments))
m = rowMeans(a100)
se = apply(a100, 1, function(x) sd(x) / sqrt(nsim))
df100 = cbind(data.frame(mse=m, mse_se=se, n=n, nsim=nsim), arguments)

cat("running sim2\n")
n = 1500
a1500 = replicate(n = nsim, run_experiment(n = n, arguments))
m = rowMeans(a1500)
se = apply(a1500, 1, function(x) sd(x) / sqrt(nsim))
df1500 = cbind(data.frame(mse=m, mse_se=se, n=n, nsim=nsim), arguments)

out  = rbind(df100, df1500)
write.csv(out, 'honesty.csv', row.names = FALSE)