|
| 1 | +#!/usr/bin/env Rscript |
| 2 | +# Generate R bacondecomp parity goldens for diff-diff BaconDecomposition. |
| 3 | +# |
| 4 | +# Requires: install.packages("bacondecomp") (CRAN; main function is bacon()) |
| 5 | +# install.packages("jsonlite") |
| 6 | +# Output: ../data/r_bacondecomp_golden.json |
| 7 | +# |
| 8 | +# The diff-diff BaconDecomposition implementation (`diff_diff/bacon.py`) with |
| 9 | +# the default ``weights="exact"`` is expected to match the values in this JSON |
| 10 | +# to atol=1e-6 on the per-component (treated, control, type) tuples, and to |
| 11 | +# match the TWFE coefficient to the same tolerance. The ``weights="approximate"`` |
| 12 | +# path is a library-only optimization and is NOT covered by this parity harness. |
| 13 | +# |
| 14 | +# Three fixtures: |
| 15 | +# 1. uniform_3groups_with_never_treated — 3 timing groups + never-treated U; |
| 16 | +# exercises all three comparison types (treated/never, earlier/later, |
| 17 | +# later/earlier). |
| 18 | +# 2. two_groups_no_never_treated — 2 timing groups only; tests the |
| 19 | +# timing-only decomposition where the s_{kU} terms drop. |
| 20 | +# 3. always_treated_remapped — 3 timing groups + 1 always-treated cohort |
| 21 | +# (first_treat = 1). Validates that Python's warn+remap of t_i < 1 into |
| 22 | +# U matches R bacondecomp's native behavior. |
| 23 | +# |
| 24 | +# Run: |
| 25 | +# cd benchmarks/R && Rscript generate_bacon_golden.R |
| 26 | + |
| 27 | +suppressPackageStartupMessages({ |
| 28 | + library(bacondecomp) |
| 29 | + library(jsonlite) |
| 30 | +}) |
| 31 | + |
| 32 | +stopifnot(packageVersion("bacondecomp") >= "0.1.0") |
| 33 | + |
| 34 | +# --------------------------------------------------------------------------- |
| 35 | +# DGP helpers |
| 36 | +# --------------------------------------------------------------------------- |
| 37 | + |
| 38 | +# Build a balanced panel with absorbing treatment. |
| 39 | +# n_units : units per timing group (excluding never-treated) |
| 40 | +# n_periods : panel length (1..T) |
| 41 | +# cohort_times : vector of first-treatment times, one per cohort |
| 42 | +# always_treated_count : optional cohort treated at first_treat = 1 |
| 43 | +# (i.e., always-treated for the observable window) |
| 44 | +# never_treated_count : units with first_treat = 0 |
| 45 | +# true_effect : constant ATT |
| 46 | +# seed : reproducibility |
| 47 | +build_panel <- function(n_units_per_cohort, n_periods, cohort_times, |
| 48 | + always_treated_count = 0L, never_treated_count = 0L, |
| 49 | + true_effect = 2.0, seed = 42L) { |
| 50 | + set.seed(seed) |
| 51 | + units <- list() |
| 52 | + uid <- 1L |
| 53 | + |
| 54 | + # Always-treated cohort (first_treat = 1; treated in every observable period) |
| 55 | + if (always_treated_count > 0L) { |
| 56 | + for (i in seq_len(always_treated_count)) { |
| 57 | + units[[length(units) + 1L]] <- data.frame( |
| 58 | + unit = uid, time = seq_len(n_periods), first_treat = 1L |
| 59 | + ) |
| 60 | + uid <- uid + 1L |
| 61 | + } |
| 62 | + } |
| 63 | + |
| 64 | + # Never-treated U |
| 65 | + if (never_treated_count > 0L) { |
| 66 | + for (i in seq_len(never_treated_count)) { |
| 67 | + units[[length(units) + 1L]] <- data.frame( |
| 68 | + unit = uid, time = seq_len(n_periods), first_treat = 0L |
| 69 | + ) |
| 70 | + uid <- uid + 1L |
| 71 | + } |
| 72 | + } |
| 73 | + |
| 74 | + # Treated cohorts |
| 75 | + for (g in cohort_times) { |
| 76 | + for (i in seq_len(n_units_per_cohort)) { |
| 77 | + units[[length(units) + 1L]] <- data.frame( |
| 78 | + unit = uid, time = seq_len(n_periods), first_treat = as.integer(g) |
| 79 | + ) |
| 80 | + uid <- uid + 1L |
| 81 | + } |
| 82 | + } |
| 83 | + |
| 84 | + df <- do.call(rbind, units) |
| 85 | + |
| 86 | + # Treatment indicator: D_{it} = 1 iff first_treat in {1,..,T} AND time >= first_treat. |
| 87 | + df$D <- as.integer(df$first_treat > 0L & df$time >= df$first_treat) |
| 88 | + |
| 89 | + # Outcome: unit FE + linear time + constant treatment effect + noise. |
| 90 | + unit_fe <- rnorm(uid - 1L, sd = 2.0) |
| 91 | + df$y <- unit_fe[df$unit] + |
| 92 | + 0.1 * df$time + |
| 93 | + true_effect * df$D + |
| 94 | + rnorm(nrow(df), sd = 0.5) |
| 95 | + |
| 96 | + df |
| 97 | +} |
| 98 | + |
| 99 | +# --------------------------------------------------------------------------- |
| 100 | +# Extract bacondecomp::bacon() output into a fixture-shaped list. |
| 101 | +# --------------------------------------------------------------------------- |
| 102 | + |
| 103 | +extract_bacon <- function(df, fixture_name) { |
| 104 | + # bacondecomp::bacon takes the OUTCOME ~ TREATMENT formula plus id_var/time_var. |
| 105 | + # It returns a data.frame with columns: treated, untreated, estimate, weight, |
| 106 | + # plus a `type` column (e.g. "Both Treated", "Treated vs Untreated"), and an |
| 107 | + # attribute beta_hat_w (the weighted sum, which equals the TWFE coefficient). |
| 108 | + res <- bacondecomp::bacon( |
| 109 | + formula = y ~ D, |
| 110 | + data = df, |
| 111 | + id_var = "unit", |
| 112 | + time_var = "time" |
| 113 | + ) |
| 114 | + |
| 115 | + # When the data contains a never-treated group, bacon() returns a list with |
| 116 | + # $two_by_twos (the per-component table) and $Omega (the variance-weighted |
| 117 | + # contributions). Without never-treated, it returns the data.frame directly. |
| 118 | + if (is.list(res) && !is.data.frame(res)) { |
| 119 | + components_df <- res$two_by_twos |
| 120 | + twfe_coef <- as.numeric(attr(res, "beta_hat_w")) |
| 121 | + # Fallback: re-derive TWFE from the components if attr is missing. |
| 122 | + if (is.null(twfe_coef) || length(twfe_coef) == 0L) { |
| 123 | + twfe_coef <- sum(components_df$estimate * components_df$weight) |
| 124 | + } |
| 125 | + } else { |
| 126 | + components_df <- res |
| 127 | + twfe_coef <- sum(components_df$estimate * components_df$weight) |
| 128 | + } |
| 129 | + |
| 130 | + # Components vary across bacondecomp versions; normalize the column names. |
| 131 | + cols <- names(components_df) |
| 132 | + treated_col <- if ("treated" %in% cols) "treated" else "g1" |
| 133 | + untreated_col <- if ("untreated" %in% cols) "untreated" else "g2" |
| 134 | + estimate_col <- if ("estimate" %in% cols) "estimate" else "Estimate" |
| 135 | + weight_col <- if ("weight" %in% cols) "weight" else "Weight" |
| 136 | + type_col <- if ("type" %in% cols) "type" else NA_character_ |
| 137 | + |
| 138 | + components <- lapply(seq_len(nrow(components_df)), function(i) { |
| 139 | + list( |
| 140 | + treated_group = as.numeric(components_df[[treated_col]][i]), |
| 141 | + control_group = as.numeric(components_df[[untreated_col]][i]), |
| 142 | + estimate = as.numeric(components_df[[estimate_col]][i]), |
| 143 | + weight = as.numeric(components_df[[weight_col]][i]), |
| 144 | + type = if (!is.na(type_col)) |
| 145 | + as.character(components_df[[type_col]][i]) |
| 146 | + else NA_character_ |
| 147 | + ) |
| 148 | + }) |
| 149 | + |
| 150 | + weights_sum <- sum(sapply(components, function(c) c$weight)) |
| 151 | + |
| 152 | + list( |
| 153 | + panel = list( |
| 154 | + unit = as.integer(df$unit), |
| 155 | + time = as.integer(df$time), |
| 156 | + y = as.numeric(df$y), |
| 157 | + first_treat = as.integer(df$first_treat), |
| 158 | + treated = as.integer(df$D) |
| 159 | + ), |
| 160 | + r_twfe_coef = twfe_coef, |
| 161 | + r_components = components, |
| 162 | + r_weights_sum = weights_sum, |
| 163 | + n_components = length(components) |
| 164 | + ) |
| 165 | +} |
| 166 | + |
| 167 | +# --------------------------------------------------------------------------- |
| 168 | +# Fixtures |
| 169 | +# --------------------------------------------------------------------------- |
| 170 | + |
| 171 | +cat("Building fixture 1: uniform_3groups_with_never_treated...\n") |
| 172 | +df1 <- build_panel( |
| 173 | + n_units_per_cohort = 30L, |
| 174 | + n_periods = 6L, |
| 175 | + cohort_times = c(3L, 4L, 5L), |
| 176 | + always_treated_count = 0L, |
| 177 | + never_treated_count = 30L, |
| 178 | + true_effect = 2.0, |
| 179 | + seed = 101L |
| 180 | +) |
| 181 | +fixture_1 <- extract_bacon(df1, "uniform_3groups_with_never_treated") |
| 182 | + |
| 183 | +cat("Building fixture 2: two_groups_no_never_treated...\n") |
| 184 | +df2 <- build_panel( |
| 185 | + n_units_per_cohort = 30L, |
| 186 | + n_periods = 6L, |
| 187 | + cohort_times = c(3L, 5L), |
| 188 | + always_treated_count = 0L, |
| 189 | + never_treated_count = 0L, |
| 190 | + true_effect = 2.0, |
| 191 | + seed = 202L |
| 192 | +) |
| 193 | +fixture_2 <- extract_bacon(df2, "two_groups_no_never_treated") |
| 194 | + |
| 195 | +cat("Building fixture 3: always_treated_remapped...\n") |
| 196 | +# 3 timing-cohorts + 5 always-treated units (first_treat = 1, i.e., treated |
| 197 | +# in every observable period) + 30 never-treated. R's bacondecomp natively |
| 198 | +# groups the first_treat=1 cohort with U (since they are treated throughout |
| 199 | +# every observable period and never serve as a within-window control), which |
| 200 | +# matches what diff-diff's warn+remap does in Python. |
| 201 | +df3 <- build_panel( |
| 202 | + n_units_per_cohort = 25L, |
| 203 | + n_periods = 6L, |
| 204 | + cohort_times = c(3L, 4L, 5L), |
| 205 | + always_treated_count = 5L, |
| 206 | + never_treated_count = 25L, |
| 207 | + true_effect = 2.0, |
| 208 | + seed = 303L |
| 209 | +) |
| 210 | +fixture_3 <- extract_bacon(df3, "always_treated_remapped") |
| 211 | + |
| 212 | +# --------------------------------------------------------------------------- |
| 213 | +# Write JSON |
| 214 | +# --------------------------------------------------------------------------- |
| 215 | + |
| 216 | +out <- list( |
| 217 | + meta = list( |
| 218 | + generated_at = format(Sys.Date()), |
| 219 | + bacondecomp_version = as.character(packageVersion("bacondecomp")), |
| 220 | + r_version = R.version.string, |
| 221 | + description = paste( |
| 222 | + "Goodman-Bacon (2021) decomposition parity goldens for diff-diff", |
| 223 | + "BaconDecomposition. Parity target: atol=1e-6 on per-component", |
| 224 | + "(treated, control, type) tuples plus the TWFE coefficient." |
| 225 | + ) |
| 226 | + ), |
| 227 | + uniform_3groups_with_never_treated = fixture_1, |
| 228 | + two_groups_no_never_treated = fixture_2, |
| 229 | + always_treated_remapped = fixture_3 |
| 230 | +) |
| 231 | + |
| 232 | +out_path <- "../data/r_bacondecomp_golden.json" |
| 233 | +write_json(out, out_path, pretty = TRUE, digits = NA, auto_unbox = TRUE) |
| 234 | +cat(sprintf("Wrote %s\n", out_path)) |
0 commit comments