|
| 1 | +# Generate cross-language weighted-2SLS parity fixture for HAD Phase 4.5 B |
| 2 | +# (mass-point + weights). |
| 3 | +# |
| 4 | +# Purpose: validate ``_fit_mass_point_2sls(..., weights=...)`` against |
| 5 | +# ``estimatr::iv_robust(y ~ d | Z, weights=w, se_type=...)`` bit-exactly. |
| 6 | +# estimatr's HC1 sandwich and Stata-style CR1 under pweights match the |
| 7 | +# Wooldridge 2010 Ch. 12 / Angrist-Pischke 4.1.3 pweight convention |
| 8 | +# that the Python port implements (w² in the HC1 meat, w·u in the CR1 |
| 9 | +# cluster score, weighted bread Z'WX). estimatr's classical SE uses a |
| 10 | +# different DOF / projection convention and is skipped in parity tests |
| 11 | +# (documented deviation; diverges by O(1/n) at non-uniform weights). |
| 12 | +# |
| 13 | +# Usage: |
| 14 | +# Rscript benchmarks/R/generate_estimatr_iv_robust_golden.R |
| 15 | +# |
| 16 | +# Output: |
| 17 | +# benchmarks/data/estimatr_iv_robust_golden.json |
| 18 | +# |
| 19 | +# Phase 4.5 B of HeterogeneousAdoptionDiD (de Chaisemartin et al. 2026). |
| 20 | +# Python test loader: tests/test_estimatr_iv_robust_parity.py. |
| 21 | + |
| 22 | +library(jsonlite) |
| 23 | +library(estimatr) |
| 24 | + |
| 25 | +stopifnot(packageVersion("estimatr") >= "1.0") |
| 26 | + |
| 27 | +# ------------------------------------------------------------------------- |
| 28 | +# DGP builders |
| 29 | +# ------------------------------------------------------------------------- |
| 30 | + |
| 31 | +dgp_mass_point <- function(n, seed, weight_pattern = "uniform", |
| 32 | + include_cluster = FALSE, d_lower = 0.3) { |
| 33 | + # Mass-point dose: a fraction at d_lower, rest uniform(d_lower, 1). |
| 34 | + set.seed(seed) |
| 35 | + n_mass <- round(0.15 * n) |
| 36 | + d_mass <- rep(d_lower, n_mass) |
| 37 | + d_cont <- runif(n - n_mass, d_lower, 1.0) |
| 38 | + d <- c(d_mass, d_cont) |
| 39 | + # Reshuffle to avoid ordered-by-dose artifacts. |
| 40 | + perm <- sample.int(n) |
| 41 | + d <- d[perm] |
| 42 | + |
| 43 | + # True DGP: dy = 2 * d + 0.3 * d^2 + eps |
| 44 | + dy <- 2.0 * d + 0.3 * d^2 + rnorm(n, sd = 0.4) |
| 45 | + |
| 46 | + # Weights |
| 47 | + w <- switch(weight_pattern, |
| 48 | + "uniform" = rep(1.0, n), |
| 49 | + "mild" = 1.0 + 0.3 * rnorm(n), |
| 50 | + "informative" = 1.0 + 2.0 * abs(d - 0.5) + runif(n, 0, 0.2), |
| 51 | + "heavy_tailed" = pmax(0.1, rlnorm(n, meanlog = 0, sdlog = 0.5)) |
| 52 | + ) |
| 53 | + # Clip to positive. |
| 54 | + w <- pmax(w, 0.01) |
| 55 | + |
| 56 | + cluster <- if (include_cluster) sample.int(max(4, n %/% 20), n, replace = TRUE) else NULL |
| 57 | + |
| 58 | + list(d = d, dy = dy, w = w, cluster = cluster, d_lower = d_lower) |
| 59 | +} |
| 60 | + |
| 61 | +# ------------------------------------------------------------------------- |
| 62 | +# Fit weighted 2SLS with estimatr at specified se_type. |
| 63 | +# ------------------------------------------------------------------------- |
| 64 | + |
| 65 | +fit_iv_robust <- function(dgp, se_type, use_cluster = FALSE) { |
| 66 | + d <- dgp$d |
| 67 | + dy <- dgp$dy |
| 68 | + w <- dgp$w |
| 69 | + Z <- as.integer(d > dgp$d_lower) |
| 70 | + df <- data.frame(d = d, dy = dy, Z = Z, w = w) |
| 71 | + if (use_cluster) df$cluster <- dgp$cluster |
| 72 | + |
| 73 | + fit <- if (use_cluster) { |
| 74 | + iv_robust(dy ~ d | Z, data = df, weights = w, clusters = cluster, |
| 75 | + se_type = se_type) |
| 76 | + } else { |
| 77 | + iv_robust(dy ~ d | Z, data = df, weights = w, se_type = se_type) |
| 78 | + } |
| 79 | + |
| 80 | + list( |
| 81 | + beta = as.numeric(coef(fit)["d"]), |
| 82 | + se = as.numeric(fit$std.error["d"]), |
| 83 | + # Intercept for manual sandwich verification. |
| 84 | + alpha = as.numeric(coef(fit)["(Intercept)"]), |
| 85 | + se_intercept = as.numeric(fit$std.error["(Intercept)"]), |
| 86 | + n = as.integer(nobs(fit)), |
| 87 | + se_type = se_type |
| 88 | + ) |
| 89 | +} |
| 90 | + |
| 91 | +# ------------------------------------------------------------------------- |
| 92 | +# Build the DGP × se_type fixture grid. |
| 93 | +# ------------------------------------------------------------------------- |
| 94 | + |
| 95 | +# Each DGP × se_type combination becomes one fixture entry. DGPs vary |
| 96 | +# sample size, weight informativeness, and cluster structure so the |
| 97 | +# Python test exercises all three sandwich variants (HC1, classical, CR1). |
| 98 | +fixtures <- list() |
| 99 | + |
| 100 | +dgps <- list( |
| 101 | + list(n = 200, seed = 42, weight = "uniform", cluster = FALSE, name = "uniform_n200"), |
| 102 | + list(n = 500, seed = 123, weight = "mild", cluster = FALSE, name = "mild_n500"), |
| 103 | + list(n = 500, seed = 7, weight = "informative", cluster = FALSE, name = "informative_n500"), |
| 104 | + list(n = 1000, seed = 321, weight = "heavy_tailed", cluster = FALSE, name = "heavy_n1000"), |
| 105 | + list(n = 600, seed = 99, weight = "informative", cluster = TRUE, name = "informative_cluster_n600") |
| 106 | +) |
| 107 | + |
| 108 | +# For the non-clustered DGPs, emit HC1 + classical entries (Python |
| 109 | +# parity tests target HC1; classical deviates by O(1/n) and is recorded |
| 110 | +# as a reference only). For the clustered DGP, emit the Stata-style CR1 |
| 111 | +# entry (matches `diff_diff/had.py::_fit_mass_point_2sls` CR1 convention |
| 112 | +# bit-exactly; see Gate #0 verification in the Phase 4.5 B plan). |
| 113 | +for (dgp_spec in dgps) { |
| 114 | + dgp <- dgp_mass_point( |
| 115 | + n = dgp_spec$n, |
| 116 | + seed = dgp_spec$seed, |
| 117 | + weight_pattern = dgp_spec$weight, |
| 118 | + include_cluster = dgp_spec$cluster |
| 119 | + ) |
| 120 | + |
| 121 | + if (dgp_spec$cluster) { |
| 122 | + entry <- list( |
| 123 | + name = dgp_spec$name, |
| 124 | + n = dgp_spec$n, |
| 125 | + d_lower = dgp$d_lower, |
| 126 | + weight_pattern = dgp_spec$weight, |
| 127 | + seed = dgp_spec$seed, |
| 128 | + d = dgp$d, |
| 129 | + dy = dgp$dy, |
| 130 | + w = dgp$w, |
| 131 | + cluster = dgp$cluster, |
| 132 | + cr1 = fit_iv_robust(dgp, se_type = "stata", use_cluster = TRUE) |
| 133 | + ) |
| 134 | + } else { |
| 135 | + entry <- list( |
| 136 | + name = dgp_spec$name, |
| 137 | + n = dgp_spec$n, |
| 138 | + d_lower = dgp$d_lower, |
| 139 | + weight_pattern = dgp_spec$weight, |
| 140 | + seed = dgp_spec$seed, |
| 141 | + d = dgp$d, |
| 142 | + dy = dgp$dy, |
| 143 | + w = dgp$w, |
| 144 | + cluster = NULL, |
| 145 | + hc1 = fit_iv_robust(dgp, se_type = "HC1", use_cluster = FALSE), |
| 146 | + classical = fit_iv_robust(dgp, se_type = "classical", use_cluster = FALSE) |
| 147 | + ) |
| 148 | + } |
| 149 | + fixtures[[dgp_spec$name]] <- entry |
| 150 | +} |
| 151 | + |
| 152 | +# ------------------------------------------------------------------------- |
| 153 | +# Serialize |
| 154 | +# ------------------------------------------------------------------------- |
| 155 | + |
| 156 | +out <- list( |
| 157 | + metadata = list( |
| 158 | + description = "estimatr::iv_robust weighted 2SLS parity fixture for HAD Phase 4.5 B", |
| 159 | + estimatr_version = as.character(packageVersion("estimatr")), |
| 160 | + r_version = as.character(getRversion()), |
| 161 | + n_dgps = length(fixtures), |
| 162 | + hc1_atol = 1e-10, |
| 163 | + cr1_atol = 1e-10, |
| 164 | + notes = paste( |
| 165 | + "HC1 (se_type='HC1') and CR1 (se_type='stata') under pweights match", |
| 166 | + "Python's _fit_mass_point_2sls bit-exactly at atol=1e-10. Classical", |
| 167 | + "(se_type='classical') uses estimatr's projection-form DOF convention", |
| 168 | + "(n-k + X_hat'WX_hat bread) which differs from Python's sandwich form", |
| 169 | + "(sum(w)-k + Z'W^2Z meat); included as a reference without parity", |
| 170 | + "assertion.", |
| 171 | + sep = " " |
| 172 | + ) |
| 173 | + ), |
| 174 | + fixtures = fixtures |
| 175 | +) |
| 176 | + |
| 177 | +# Ensure output directory exists. |
| 178 | +out_dir <- "benchmarks/data" |
| 179 | +if (!dir.exists(out_dir)) dir.create(out_dir, recursive = TRUE) |
| 180 | +out_path <- file.path(out_dir, "estimatr_iv_robust_golden.json") |
| 181 | +write_json(out, path = out_path, digits = 17, auto_unbox = TRUE, null = "null") |
| 182 | +message(sprintf("Wrote %d DGP fixtures to %s", length(fixtures), out_path)) |
0 commit comments