# \donttest{
if (require("grf", quietly = TRUE)) {
# Simulate data with two treatment arms (k = 1, 2) and a control arm (k = 0).
n <- 3000
p <- 5
X <- matrix(runif(n * p), n, p)
W <- as.factor(sample(c("0", "1", "2"), n, replace = TRUE))
Y <- X[, 1] + X[, 2] * (W == "1") + 1.5 * X[, 3] * (W == "2") + rnorm(n)
# Fit a CATE estimator on a training sample.
train <- sample(1:n, n/2)
tau.forest <- grf::multi_arm_causal_forest(X[train, ], Y[train], W[train])
# Predict CATEs on held out evaluation data.
test <- -train
tau.hat <- predict(tau.forest, X[test, ], drop = TRUE)$predictions
# Form costs.
cost <- cbind(X[test, 4] / 4, X[test, 5])
# Estimate nuisance components for test set AIPW scores.
X.test <- X[test, ]
Y.test <- Y[test]
W.test <- W[test]
# Fit models for E[Y | W = k, X], k = 0, 1, 2, using for example separate random forests.
Y0.forest <- grf::regression_forest(X.test[W.test == 0, ], Y.test[W.test == 0])
Y1.forest <- grf::regression_forest(X.test[W.test == 1, ], Y.test[W.test == 1])
Y2.forest <- grf::regression_forest(X.test[W.test == 2, ], Y.test[W.test == 2])
mu.hat = cbind(
mu0 = predict(Y0.forest, X.test)$predictions,
mu1 = predict(Y1.forest, X.test)$predictions,
mu2 = predict(Y2.forest, X.test)$predictions
)
# If unknown, estimate the propensity scores E[W = k | X].
W.hat <- predict(grf::probability_forest(X.test, W.test))$predictions
# Form doubly robust scores.
DR.scores <- get_aipw_scores(Y.test, W.test, mu.hat, W.hat)
# Fit a Qini curve estimated with forest-based AIPW.
qini <- maq(tau.hat, cost, DR.scores, R = 200)
plot(qini)
}
# }
Run the code above in your browser using DataLab