08b_placebo_test_cito.R

library(tidyverse)
library(pensynth)
library(digest)
library(cli)

schools_df <- read_rds("processed_data/school_data.rds")
synth_mats <- read_rds("processed_data/synth_mats_cito.rds")

placebo_err <- function(mats, N_placebo = 250, verbose = TRUE) {# subsampling for computational efficiency
  N_donor <- ncol(mats$X0)
  N_placebo <- min(N_donor, N_placebo) 
  placebo_idx <- sample(N_donor, N_placebo)
  
  EY <- matrix(NA, nrow = nrow(mats$Y1), ncol = N_placebo)
  EZ <- matrix(NA, nrow = nrow(mats$Z1), ncol = N_placebo)
  
  for (i in seq_along(placebo_idx)) {
    if (verbose) cat("Estimating model", i, "/", N_placebo, "...\r")
    pidx <- placebo_idx[i]
    
    X1p  <- mats$X0[,pidx]
    X0p  <- mats$X0[,-pidx]
    Y1p  <- mats$Y0[,pidx]
    Y0p  <- mats$Y0[,-pidx]
    Z1p  <- mats$Z0[,pidx]
    Z0p  <- mats$Z0[,-pidx]
    
    # hold-out validated version is more correct, but takes way longer and 
    # yields basically the same results
    # pmod <- cv_pensynth(
    #   X1 = X1p,
    #   X0 = X0p,
    #   Z1 = Z1p,
    #   Z0 = Z0p,
    #   nlambda = 20,
    #   opt_pars = clarabel::clarabel_control(verbose = FALSE)
    # )

    # Instead, we use approx. the avg. optimal lambda from the main analysis
    pmod <- pensynth(
      X1 = X1p,
      X0 = X0p,
      lambda = exp(-2),
      opt_pars = clarabel::clarabel_control(verbose = FALSE)
    )
    
    Y1phat <- predict(pmod, Y0p)
    Z1phat <- predict(pmod, Z0p)
    EY[,i] <- Y1p - Y1phat
    EZ[,i] <- Z1p - Z1phat
  }
  if (verbose) cat("\n")
  return(list(EY = EY, EZ = EZ))
}


# use caching to avoid recomputing if the donor pool is the same as seen before
cache <- list()
res <- list()
for (i_trt in seq_along(synth_mats)) {
  mats <- synth_mats[[i_trt]]
  trt_name <- names(synth_mats)[i_trt]
  # check cache
  hash <- digest(mats$X0)
  if (hash %in% names(cache)) {
    cli_alert_info(paste0("Cache hit for unit ", i_trt, " (", trt_name, ")"))
    res[[trt_name]] <- cache[[hash]]
    next
  }
  cli_alert_info(paste0("Computing placebo values for unit ", i_trt, " (", trt_name, ")"))
  cache[[hash]] <- res[[trt_name]] <- placebo_err(mats)
}

cat("Saving placebo values to processed_data/placebo_cito.rds")
write_rds(res, "processed_data/placebo_cito.rds")

# Now, actually compute the placebo p-values
placebo_res <- read_rds("processed_data/placebo_cito.rds")
psynth_list <- read_rds("processed_data/psynth_list_cito.rds")
synth_mats <- read_rds("processed_data/synth_mats_cito.rds")


# Create an average ACE over units, per permutation
null_vals <- colMeans(Reduce(rbind, map(placebo_res, \(i) i$EY)))
null_dist <- ecdf(null_vals)
test_stat <- mean(unlist(imap(synth_mats, \(m, i) m$Y1 - predict(psynth_list[[i]], m$Y0))))
placebo_pval <- 1 - null_dist(test_stat)
df_ace <- tibble(type = "ACE", statistic = c("test", rep("null", length(null_vals))), value = c(test_stat, null_vals))

png("img/placebo_cito.png", width = 3200, height = 1800, res = 300)
dd <- density(null_vals, bw = 0.7)
plot(
  x = dd, 
  main = "Placebo test for CITO score", 
  xlab = "Post-intervention Average Causal Effect"
)
polygon(c(dd$x[dd$x>test_stat], test_stat), c(dd$y[dd$x>test_stat], 0), col = "#aaeeff88", border = NA)
abline(v = test_stat, lwd = 2, col = "red")
rug(null_vals, ticksize = 0.02, lwd = 1)
mtext(paste0("Placebo p-value: ", round(placebo_pval, 3)), line = .5)
mtext("Averaged over treated units and years", 1, line = 4, cex = 0.8)
legend(
  x = "topleft", 
  legend = c("Null ACE", "Obs. ACE"), 
  lwd = c(0, 2), 
  col = c(NA, "red"), 
  fill = c("#aaeeff88", NA), 
  border = c(TRUE, FALSE), 
  seg.len = 1, 
  merge = TRUE
)
dev.off()

# just to check, the visual distortion of the KDE
# do the p-value based on the KDE-estimated density
print(placebo_pval)
integrate(approxfun(dd, yleft = 0, yright = 0), test_stat, Inf)
# this is close enough to the quantile-based p-value

# now, do 2016, CE averaged over units, per permutation
null_vals <- colMeans(Reduce(rbind, map(placebo_res, \(i) i$EY[3,,drop=FALSE])))
null_dist <- ecdf(null_vals)
test_stat <- mean(unlist(imap(synth_mats, \(m, i) m$Y1[3,] - predict(psynth_list[[i]], m$Y0)[3,])))
placebo_pval <- 1 - null_dist(test_stat)
df_2016 <- tibble(type = "2016", statistic = c("test", rep("null", length(null_vals))), value = c(test_stat, null_vals))

png("img/placebo_cito_2016.png", width = 3200, height = 1800, res = 300)
dd <- density(null_vals, bw = 0.7)
plot(
  x = dd, 
  main = "Placebo test for CITO score in 2016", 
  xlab = "Causal Effect in 2016"
)
polygon(c(dd$x[dd$x>test_stat], test_stat), c(dd$y[dd$x>test_stat], 0), col = "#aaeeff88", border = NA)
abline(v = test_stat, lwd = 2, col = "red")
rug(null_vals, ticksize = 0.02, lwd = 1)
mtext(paste0("Placebo p-value: ", round(placebo_pval, 3)), line = .5)
mtext("Averaged over treated units", 1, line = 4, cex = 0.8)
legend(
  x = "topleft", 
  legend = c("Null ACE", "Obs. ACE"), 
  lwd = c(0, 2), 
  col = c(NA, "red"), 
  fill = c("#aaeeff88", NA), 
  border = c(TRUE, FALSE), 
  seg.len = 1, 
  merge = TRUE
)
dev.off()

# just to check, the visual distortion of the KDE
# do the p-value based on the KDE-estimated density
print(placebo_pval)
integrate(approxfun(dd, yleft = 0, yright = 0), test_stat, Inf)
# this is close enough to the quantile-based p-value

write_rds(bind_rows(df_ace, df_2016) |> mutate(n = length(placebo_res)), "processed_data/placebo_values_cito.rds")