nonparam.R

nonparam.R

rm(list = ls())
library(grf)

coeffs = replicate(250, {
  # setup from 
  # https://tompepinsky.com/2022/06/07/illustrating-contamination-bias-with-implications-for-intersectional-description/

  effect1 = 1
  effect2 = 2
  effect3 = 3
  n_w = 500
  
  t <- c(
    sample(c(0, 1, 2), size = n_w, replace=TRUE, prob = c(.1,.4,.5)),
    sample(c(0, 1, 2), size = n_w, replace=TRUE, prob = c(1/3,1/3,1/3)),
    sample(c(0, 1, 2), size = n_w, replace=TRUE, prob = c(.5,.4,.1))
  )
  
  modelmat <- data.frame(rep(NA,3*n_w),t,c(rep(1,n_w),rep(2,n_w),rep(3,n_w)))
  names(modelmat) <- c("y","t","w")
  
  modelmat$y[modelmat$w==1 & modelmat$t==0] <- 0 + rnorm(length(modelmat$y[modelmat$w==1 & modelmat$t==0]),0,3)
  modelmat$y[modelmat$w==1 & modelmat$t==1] <- 0 + rnorm(length(modelmat$y[modelmat$w==1 & modelmat$t==1]),0,3)
  modelmat$y[modelmat$w==1 & modelmat$t==2] <- effect1 + rnorm(length(modelmat$y[modelmat$w==1 & modelmat$t==2]),0,3)
  
  modelmat$y[modelmat$w==2 & modelmat$t==0] <- 0 + rnorm(length(modelmat$y[modelmat$w==2 & modelmat$t==0]),0,3)
  modelmat$y[modelmat$w==2 & modelmat$t==1] <- 0 + rnorm(length(modelmat$y[modelmat$w==2 & modelmat$t==1]),0,3)
  modelmat$y[modelmat$w==2 & modelmat$t==2] <- effect2 + rnorm(length(modelmat$y[modelmat$w==2 & modelmat$t==2]),0,3)
  
  modelmat$y[modelmat$w==3 & modelmat$t==0] <- 0 + rnorm(length(modelmat$y[modelmat$w==3 & modelmat$t==0]),0,3)
  modelmat$y[modelmat$w==3 & modelmat$t==1] <- 0 + rnorm(length(modelmat$y[modelmat$w==3 & modelmat$t==1]),0,3)
  modelmat$y[modelmat$w==3 & modelmat$t==2] <- effect3 + rnorm(length(modelmat$y[modelmat$w==3 & modelmat$t==2]),0,3)
  
  ols = summary(lm(y~factor(t)+factor(w), data=modelmat))$coefficients
  
  cf = multi_arm_causal_forest(modelmat[, "w", drop = F], modelmat[, "y"], as.factor(modelmat[, "t"]))
  ate = average_treatment_effect(cf)

  c(
    ols = ols[2, 1],
    cf = ate[1, 1], # Contrast arm 1 - 0 should be zero.
    cf.se = ate[1, 2]
  )

})

par(mfrow = c(2, 1))

hist(coeffs["ols",], main = "OLS")
abline(v = 0, col='red')
hist(coeffs["cf",], main = "GRF (Multi-arm Causal Forest)")
abline(v = 0, col='red')

sd(coeffs["cf",])
# [1] 0.2142025
mean(coeffs["cf.se", ])
# [1] 0.2180095