if (!require("grf")) install.packages("grf", dependencies = TRUE); library(grf)
if (!require("hdm")) install.packages("hdm", dependencies = TRUE); library(hdm)
if (!require("OutcomeWeights")) install.packages("OutcomeWeights", dependencies = TRUE); library(OutcomeWeights)
set.seed(1234)
data(pension)
D = pension$p401
Z = pension$e401
Y = pension$net_tfa
X = model.matrix(~ 0 + age + db + educ + fsize + hown + inc + male + marr + pira + twoearn, data = pension)
var_nm = c("Age","Benefit pension","Education","Family size","Home owner","Income","Male","Married","IRA","Two earners")
colnames(X) = var_nm
rf_Y.hat = regression_forest(X,Y)
Y.hat = predict(rf_Y.hat)$predictions
S = get_forest_weights(rf_Y.hat)
cf = causal_forest(X,Y,D,Y.hat=Y.hat)
D.hat = cf$W.hat
S.tau = get_outcome_weights(cf, S = S)
all.equal(as.numeric(S.tau$omega %*% Y),
as.numeric(predict(cf)$predictions))
# ATE
## Package output
average_treatment_effect(cf, target.sample = "all") # Ignore the warning, which is overly cautious and paternalistic
## Rebuild using weights following Appendix A.3.2
lambda1 = D / D.hat
lambda0 = (1-D) / (1-D.hat)
N = length(Y)
ones = matrix(1,N,1)
### Horribly slow but close to formula (do not run)
# T_ate_slow = S.tau$omega + diag(lambda1 - lambda0) %*% (diag(N) - S - diag(as.numeric(D-D.hat)) %*% S.tau$omega)
# omega_ate_slow = t(ones) %*% T_ate_slow / N
# omega_ate_slow %*% Y
### Faster avoiding slow matrix multiplications
scaled_S.tau = (D - D.hat) * S.tau$omega # Element-wise multiplication
S_adjusted = diag(N) - S - scaled_S.tau
T_ate = S.tau$omega + (lambda1 - lambda0) * S_adjusted # Element-wise multiplication
omega_ate = t(ones) %*% T_ate / N
## Check numerical equivalence also with package weights
omega_ate_pkg = get_outcome_weights(cf, target = "ATE", S = S, S.tau = S.tau$omega)
all.equal(as.numeric(omega_ate %*% Y),
as.numeric(average_treatment_effect(cf, target.sample = "all")[1]),
as.numeric(omega_ate_pkg$omega %*% Y))
# ATT
## Package output
average_treatment_effect(cf, target.sample = "treated")
## Rebuild using weights
## Check numerical equivalence
# ATU
## Package output
average_treatment_effect(cf, target.sample = "control") # Ignore the warning, which is overly cautious and paternalistic
## Rebuild using weights
## Check numerical equivalence
# Overlap
## Package output
average_treatment_effect(cf, target.sample = "overlap") # Ignore the warning, this is overly cautious and paternalistic
## Rebuild using weights
## Check numerical equivalence
# LATE
ivf = instrumental_forest(X,Y,D,Z,Y.hat=Y.hat)
omega_if = get_outcome_weights(ivf, S = S)
all.equal(as.numeric(omega_if$omega %*% Y),
as.numeric(predict(ivf)$predictions))
## Package output
average_treatment_effect(ivf)
## Rebuild using weights
## Check numerical equivalence
