examples-predict.lm_robust.R

library(estimatr)
set.seed(60637)

N <- 40
dat <- data.frame(
  x = rnorm(N, mean = 2.3),
  x2 = rpois(N, lambda = 2),
  x3 = runif(N)
)

dat$y0 <- rnorm(N) + dat$x
dat$y1 <- dat$y0 + 0.35
dat$y2 <- dat$y0 + 0.55

dat$z_multi <- sample(0:2, size = nrow(dat), replace = TRUE)
dat$z_bin <- 1*(dat$z_multi>0)
dat$y <- (dat$z_multi == 0)*dat$y0 + (dat$z_multi == 1)*dat$y1 + (dat$z_multi == 2)*dat$y2

# Binary 01 treatments with lm_lin; estimation + prediction work as expected
lmlin_bin <- lm_lin(y ~ z_bin, covariates = ~ x + x2 + x3, data = dat)
predict(lmlin_bin, newdata = dat)
max(abs(predict(lmlin_bin, newdata = dat)-lmlin_bin$fitted.values))

# Multi-valued numeric treatments with lm_lin; estimation works as expected
lmlin_mult <- lm_lin(y ~ z_multi, covariates = ~ x, data = dat)
# prediction does not
predict(lmlin_mult, newdata = dat)
# Error in X[, !beta_na, drop = FALSE] :
#   (subscript) logical subscript too long

# Binary factorial treatment with lm_lin; estimation works,
lmlin_bin_f <- lm_lin(y ~ as.factor(z_bin), covariates = ~ x + x2 + x3, data = dat)
# prediction breaks
predict(lmlin_bin_f, newdata = dat)
# Error in X[, treat_name] : subscript out of bounds

# Multi-valued factorial treatments with lm_lin; estimation works
lmlin_mult_f <- lm_lin(y ~ as.factor(z_multi), covariates = ~ x, data = dat)
# prediction breaks
predict(lmlin_mult_f, newdata = dat)
# Error in X[, treat_name] : subscript out of bounds