## ----setup, include = FALSE-------------------------------------------------------------
knitr::opts_chunk$set(collapse = TRUE, comment = "#>", message = FALSE)
library(estimatr)
options(digits = 4)

## ---------------------------------------------------------------------------------------
CHECKS <- list()

check <- function(label, ours, theirs, tol = 1e-10) {
  gap <- max(abs(ours - theirs) / pmax(abs(theirs), 1))
  # Two jobs: record the gap in the running list for the final table, and
  # return a one-row data frame so the calling chunk prints its own result.
  CHECKS[[label]] <<- gap
  data.frame(gap = sprintf("%.1e", gap), holds = gap < tol)
}

## ---------------------------------------------------------------------------------------
set.seed(20260826)
N <- 100

d <- data.frame(
  z  = rbinom(N, 1, 0.5),
  x  = rnorm(N),
  w  = runif(N, 0.5, 2),
  g  = rep(1:10, each = 10),
  cl = rep(1:20, each = 5),
  inst = rnorm(N)
)
d$en <- d$inst + rnorm(N, 0, 0.5)
d$y <- 1 + 0.5 * d$z + 2 * d$x + d$en + rnorm(N)

## ---------------------------------------------------------------------------------------
check("lm_robust()",
      coef(lm_robust(y ~ z + x, data = d)),
      coef(lm(y ~ z + x, data = d)))

## ---------------------------------------------------------------------------------------
check("lm_robust(weights = )",
      coef(lm_robust(y ~ z + x, data = d, weights = w)),
      coef(lm(y ~ z + x, data = d, weights = w)))

## ---------------------------------------------------------------------------------------
hc_vcov <- function(fit, type) {
  X <- model.matrix(fit)
  e <- residuals(fit)
  bread <- solve(crossprod(X))
  h <- rowSums((X %*% bread) * X)
  n <- nrow(X)
  k <- ncol(X)
  adj <- switch(type,
    HC0 = e^2,
    HC1 = e^2 * n / (n - k),
    HC2 = e^2 / (1 - h),
    HC3 = e^2 / (1 - h)^2
  )
  bread %*% crossprod(X * sqrt(adj)) %*% bread
}

## ---------------------------------------------------------------------------------------
fit_lm <- lm(y ~ z + x, data = d)

check("lm_robust(se_type = 'classical')",
      lm_robust(y ~ z + x, data = d, se_type = "classical")$vcov,
      vcov(fit_lm))

do.call(rbind, lapply(c("HC0", "HC1", "HC2", "HC3"), function(ty) {
  cbind(se_type = ty,
        check(paste0("lm_robust(se_type = '", ty, "')"),
              lm_robust(y ~ z + x, data = d, se_type = ty)$vcov,
              hc_vcov(fit_lm, ty)))
}))

## ---------------------------------------------------------------------------------------
cr_vcov <- function(fit, cluster, stata = FALSE) {
  X <- model.matrix(fit)
  e <- residuals(fit)
  bread <- solve(crossprod(X))
  meat <- Reduce(`+`, lapply(split(seq_len(nrow(X)), cluster), function(i) {
    tcrossprod(crossprod(X[i, , drop = FALSE], e[i]))
  }))
  v <- bread %*% meat %*% bread
  if (!stata) return(v)
  S <- length(unique(cluster))
  v * (S / (S - 1)) * ((nrow(X) - 1) / (nrow(X) - ncol(X)))
}

## ---------------------------------------------------------------------------------------
check("lm_robust(clusters = )",
      lm_robust(y ~ z + x, data = d, clusters = cl, se_type = "CR0")$vcov,
      cr_vcov(fit_lm, d$cl))

check("lm_robust(se_type = 'stata')",
      lm_robust(y ~ z + x, data = d, clusters = cl, se_type = "stata")$vcov,
      cr_vcov(fit_lm, d$cl, stata = TRUE))

## ---------------------------------------------------------------------------------------
absorbed <- lm_robust(y ~ z + x, data = d, fixed_effects = ~ g)
dummies  <- lm_robust(y ~ z + x + factor(g), data = d)
keep <- c("z", "x")

check("lm_robust(fixed_effects = )",
      c(coef(absorbed)[keep], absorbed$std.error[keep]),
      c(coef(dummies)[keep], dummies$std.error[keep]))

## ---------------------------------------------------------------------------------------
lm_robust(mpg ~ hp, data = mtcars, weights = wt, se_type = "HC2")$std.error

## ---------------------------------------------------------------------------------------
dd <- d
dd$x_c <- dd$x - mean(dd$x)
lin  <- lm_lin(y ~ z, covariates = ~ x, data = d)
byhand <- lm_robust(y ~ z * x_c, data = dd)

check("lm_lin()",
      c(coef(lin)[["z"]], lin$std.error[["z"]]),
      c(coef(byhand)[["z"]], byhand$std.error[["z"]]))

## ---------------------------------------------------------------------------------------
tsls_coef <- function(y, X, Z) {
  xhat <- Z %*% solve(crossprod(Z), crossprod(Z, X))
  as.vector(solve(crossprod(xhat), crossprod(xhat, y)))
}

check("iv_robust()",
      unname(coef(iv_robust(y ~ en + x | inst + x, data = d))),
      tsls_coef(d$y, model.matrix(~ en + x, d), model.matrix(~ inst + x, d)))

## ---------------------------------------------------------------------------------------
fit <- lm_robust(y ~ z + x, data = d)
lh <- lh_robust(y ~ z + x, data = d, linear_hypothesis = "z + x = 0")
a <- c(0, 1, 1)

check("lh_robust()",
      c(lh$lh$coefficients[[1]], lh$lh$std.error[[1]]),
      c(sum(a * coef(fit)), sqrt(drop(t(a) %*% fit$vcov %*% a))))

## ---------------------------------------------------------------------------------------
dim_fit <- difference_in_means(y ~ z, data = d)
ols <- lm_robust(y ~ z, data = d, se_type = "HC2")

check("difference_in_means()",
      c(dim_fit$coefficients[["z"]], dim_fit$std.error[["z"]]),
      c(ols$coefficients[["z"]], ols$std.error[["z"]]))

## ---------------------------------------------------------------------------------------
blocked <- data.frame(bl = rep(1:10, each = 10),
                      z = rep(rep(0:1, each = 5), times = 10))
blocked$y <- rnorm(100) + 0.3 * blocked$z
difference_in_means(y ~ z, data = blocked, blocks = bl)$design

pairs <- data.frame(bl = rep(1:50, each = 2), z = rep(c(0, 1), 50))
pairs$y <- rnorm(100) + 0.3 * pairs$z
difference_in_means(y ~ z, data = pairs, blocks = bl)$design

# Both kinds of block in one design: 1.x applied the matched-pairs estimator
# to all of it, after a warning.
hybrid <- rbind(blocked, transform(pairs, bl = bl + 100))
difference_in_means(y ~ z, data = hybrid, blocks = bl)$design

## ---------------------------------------------------------------------------------------
ht_estimate <- function(y, z, pr) {
  mean(y * z / pr) - mean(y * (1 - z) / (1 - pr))
}

pr <- rep(0.5, N)
check("horvitz_thompson()",
      horvitz_thompson(y ~ z, data = d, condition_prs = pr)$coefficients[[1]],
      ht_estimate(d$y, d$z, pr))

## ---------------------------------------------------------------------------------------
Y1 <- d$y[d$z == 1] / 0.5
Y0 <- d$y[d$z == 0] / 0.5

check("horvitz_thompson() variance, simple randomization",
      horvitz_thompson(y ~ z, data = d, condition_prs = pr)$std.error[[1]],
      sqrt((sum(Y1^2) + sum(Y0^2)) / N^2))

## ----eval = requireNamespace("randomizr", quietly = TRUE), warning = FALSE--------------
library(randomizr)
set.seed(2)
decl <- declare_ra(blocks = rep(c("a", "b", "c", "d"), each = 50), prob = 0.4)
Z <- conduct_ra(decl)
dat_ht <- data.frame(Y = rnorm(200) + 0.5 * Z, Z = Z)

# The design-aware variance
horvitz_thompson(Y ~ Z, data = dat_ht, condition_prs = decl)$std.error

# The conservative bound, from the marginals alone
horvitz_thompson(Y ~ Z, data = dat_ht,
                 condition_prs = c("0" = 0.6, "1" = 0.4))$std.error

## ----echo = FALSE-----------------------------------------------------------------------
knitr::kable(
  data.frame(
    Promise = names(CHECKS),
    `Largest relative gap` = sprintf("%.1e", unlist(CHECKS)),
    Holds = unlist(CHECKS) < 1e-10,
    check.names = FALSE
  ),
  row.names = FALSE
)

## ---------------------------------------------------------------------------------------
stopifnot(all(unlist(CHECKS) < 1e-10))