Causal ML - Bonus Assignment

Note: If you run the causal forest code chunks and your computer struggles already with n=100, you probably have the issue raised here. You can solve it by installing this daily build RStudio. Or just focus on the drawing part ;-)

Causal Christmas Tree Challenge

I have drawn a (stylized) christmas tree using two potential outcomes functions and investigate how well causal forests approximate the resulting CATE function for different sample sizes.

library(tidyverse)
library(grf)
library(patchwork)

set.seed(1234)

# Define the functions
cw = 0.05
ch = 0.6
m1 = function(x){(x < -0.8) * (3*x + 3) + ((x >= -0.8) & (x < (-0.8+cw))) * (-0.8*3 + 3 + ch) + ((x>(-0.8+cw)) & (x < -0.5)) * (3*x + 3) + ((x >= -0.5) & (x < (-0.5+cw))) * (-0.5*3 + 3 + ch) + ((x>(-0.5+cw)) & (x < -0.2)) * (3*x + 3) + ((x >= -0.2) & (x < (-0.2+cw))) * (-0.2*3 + 3 + ch) + ((x>(-0.2+cw)) & (x < 0)) * (3*x + 3) + (x >= 0 & x < 0.2) * (-3*x + 3) + ((x >= 0.2) & (x < (0.2+cw))) * (-0.2*3 + 3 + ch) + ((x>(0.2+cw)) & (x < 0.5)) * (-3*x + 3) + ((x >= 0.5) & (x < (0.5+cw))) * (-0.5*3 + 3 + ch) + ((x>(0.5+cw)) & (x < 0.8)) * (-3*x + 3) + ((x >= 0.8) & (x < (0.8+cw))) * (-0.8*3 + 3 + ch) + ((x>(0.8+cw))) * (-3*x + 3)  }
m0 = function(x){0 * ((x<(-0.7)) + (x>(0.07))) 
                - 0.8 *  ((x>(-0.07) & x<(0.07))) }
tau = function(x){m1(x) - m0(x)}

# Plot the two potential outcome fcts
g2 = data.frame(x = c(-1, 1)) %>% ggplot(aes(x)) + stat_function(fun=m1,size=1,colour="forestgreen") + 
  stat_function(fun=m0,size=1,colour="forestgreen") + ylab("Y(w)") + xlab("X1") + theme_bw()
# Plot CATE function
g3 = data.frame(x = c(-1, 1)) %>% ggplot(aes(x)) + stat_function(fun=tau,size=1) + ylab("CATE") + xlab("X1") + theme_bw()
g2

g3

Causal Forest meets causal christmas tree

I wrote a little function that takes the two potential outcome functions as inputs, draws a random treatment \(W \sim Bernoulli(1/2)\) and adds standard normal noise to the outcome. Then I run it for sample sizes 100/1000/10000/100000 and observe how the approximation of the very hard function improves (runs about 15 minutes on my laptop).

cf_estimation = function(m1,m0,n,p=2,...) {
  # Define the functions
  cw = 0.05
  ch = 0.6
  
  # Get CATE function
  tau = function(x){m1(x) - m0(x)}
  
  # Draw sample
  X = matrix(runif(n*p,-1,1),ncol=p)
  W = rbinom(n,1,1/2)
  Y = W*m1(X[,1]) + (1-W)*m0(X[,1]) + rnorm(n,0,1)
  
  # Run CF
  cf = causal_forest(X, Y, W, ...)
  cates = predict(cf)$predictions

  # Plot
  g = data.frame(x=X[,1],y=cates) %>% ggplot() + geom_point(aes(x=x,y=y),shape="square",color="blue") +
    stat_function(fun=tau,size=1) + ylab("CATE") + ggtitle(paste0("n=",toString(n)))
  
  # RMSE
  rmse = sqrt(mean((cates - tau(X[,1]))^2))

  # Return results
  list("g" = g,"RMSE" = rmse)
}

n100 = cf_estimation(m1,m0,100,tune.parameters = "all")
n1000 = cf_estimation(m1,m0,1000,tune.parameters = "all")
n10000 = cf_estimation(m1,m0,10000,tune.parameters = "all")
n100000 = cf_estimation(m1,m0,100000,tune.parameters = "all")

(n100$g | n1000$g) / (n10000$g | n100000$g)

data.frame(RMSE = c(n100$RMSE,n1000$RMSE,n10000$RMSE,n100000$RMSE),
   n = factor(c("n=100","n=1000","n=10000","n=100000"))) %>%
  ggplot(aes(x=n,y=RMSE)) + geom_point() + theme_bw() + geom_hline(yintercept = 0)

Bonus assignment: Your turn (5P)

Draw also something using two potential outcome functions (either something different or beautify my tree) and check how well causal forests can approximate the CATE function resulting from your drawing. You can reuse and/or modify the code above for this purpose. Send me you solutions by 22.12.