From cfea355a07a5e026ae78b1cfedbce6215c17fbd7 Mon Sep 17 00:00:00 2001 From: igerber Date: Mon, 6 Jul 2026 21:26:57 -0400 Subject: [PATCH] =?UTF-8?q?feat(callaway):=20allow=5Funbalanced=5Fpanel=20?= =?UTF-8?q?=E2=80=94=20RC-on-panel=20parity=20with=20R=20did?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds CallawaySantAnna(allow_unbalanced_panel=True) matching R did::att_gt(allow_unbalanced_panel=TRUE) on unbalanced panels. When set and the input panel is unbalanced, the pooled observations route through the repeated-cross-section levels estimator (DRDID::reg_did_rc) instead of the default within-cell panel differencing (a different estimand on unbalanced data), and the per-observation influence function is clustered by the original unit. ATT matches R bit-for-bit (cells AND dynamic aggregation): the RC cohort mass and the aggregation pg use fixed UNIT cohort mass (not observation count), and the per-unit WIF divides the per-observation weight-influence-function by each unit's observation count so the unit-clustered sum is not over-counted. The analytical SE matches R up to the documented CR1 sqrt(G/(G-1)) finite-sample factor (diff-diff's cluster path applies the Bessel correction R's att_gt getSE omits). The bootstrap resamples by unit via the same PSU. Inert on balanced panels (byte-identical to the default). Independently, the default path now emits a UserWarning on unbalanced input (previously silent). survey_design= with the flag raises NotImplementedError (deferred). Routing engages only when the panel is actually unbalanced (route-as-RC via an internal flag threaded through the six self.panel sites; the unique-unit guard is naturally bypassed since panel stays True). Verified vs R did 2.5.1 (benchmarks/data/cs_unbalanced_golden.json + TestAllowUnbalancedPanel); 499 CS tests pass (RC / survey / panel paths unchanged). The TODO's "weighting" framing was a mis-diagnosis: the dominant unbalanced divergence is the estimator choice. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_0178r8ZK6VRiXbaBiknrWjQB --- CHANGELOG.md | 12 + TODO.md | 2 +- benchmarks/R/generate_cs_unbalanced_golden.R | 96 ++++ benchmarks/data/cs_unbalanced_golden.json | 37 ++ diff_diff/staggered.py | 222 +++++++++- diff_diff/staggered_aggregation.py | 85 +++- diff_diff/staggered_bootstrap.py | 54 ++- diff_diff/staggered_results.py | 9 + docs/methodology/REGISTRY.md | 44 +- tests/test_csdid_ported.py | 437 +++++++++++++++++++ 10 files changed, 926 insertions(+), 72 deletions(-) create mode 100644 benchmarks/R/generate_cs_unbalanced_golden.R create mode 100644 benchmarks/data/cs_unbalanced_golden.json diff --git a/CHANGELOG.md b/CHANGELOG.md index 400da3db..6c0cc4ca 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -50,6 +50,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `UserWarning`. `leave_one_out` is recorded on `ImputationDiDResults` (and `to_dict()` / `summary()`); replicate-weight survey designs raise `NotImplementedError` (their variance bypasses the influence-function path). Default `False` preserves R `didimputation` parity. +- **`CallawaySantAnna(allow_unbalanced_panel=True)`** — parity with R + `did::att_gt(allow_unbalanced_panel=TRUE)` on unbalanced panels. When set and the input panel is + unbalanced (some units unobserved in some periods), the pooled observations are routed through the + repeated-cross-section levels estimator (`DRDID::reg_did_rc`), replacing the default within-cell + panel differencing (a different estimand on unbalanced data), and the per-observation influence + function is clustered by the original unit. **ATT matches R bit-for-bit** — per-cell AND dynamic + aggregation (fixed unit-cohort-mass `pg` reweighting + a per-unit WIF correction); the analytical + SE matches up to the documented CR1 `sqrt(G/(G-1))` finite-sample factor. **Inert on balanced + panels** (byte-identical to the default). Independently, the default path now emits a `UserWarning` + on unbalanced input (previously silent) pointing to the flag. `survey_design=` with the flag raises + `NotImplementedError` (deferred). Verified against R `did` 2.5.1 + (`benchmarks/data/cs_unbalanced_golden.json`). - **`ContinuousDiD` lowest-dose-as-control** (`control_group="lowest_dose"`, CGBS 2024 Remark 3.1) for settings with no untreated group (`P(D=0) = 0`): the lowest-dose group `d_L` becomes the comparison and the estimand is `ATT(d) − ATT(d_L)` (with `ATT(d_L) = 0` the omitted reference). It is a diff --git a/TODO.md b/TODO.md index 17291ace..53d10180 100644 --- a/TODO.md +++ b/TODO.md @@ -128,7 +128,7 @@ Doable in principle, but no current caller and/or explicitly out of paper scope. | Issue | Location | PR | Priority | |-------|----------|----|----------| -| CallawaySantAnna **unbalanced-panel event-study weighting**: a dynamic horizon pooling >1 (g,t) cell weights each cell by its per-cell aggregation weight (valid `n_treated` / `agg_weight`), but R `did::aggte` weights by the fixed cohort probability `pg = n_g/N` (or survey mass) from the group column. Balanced panels coincide exactly (verified ~1e-5, incl. universal zero-reference dilution); unbalanced panels can differ at multi-cell horizons (e.g. a two-real-cell horizon `-0.065` vs R `-0.136`). Pre-existing + general (independent of universal reference cells; documented Deviation from R in REGISTRY). Fix = weight every event-study cell by fixed cohort mass (touches `_aggregate_event_study` + the bootstrap bucket + WIF `pg` consistency); gate on balanced byte-identity + fresh unbalanced R parity. | `staggered_aggregation.py`, `staggered_bootstrap.py` | SE-audit D3 | Low | +| CallawaySantAnna **unbalanced-panel R parity — LANDED** via `allow_unbalanced_panel=True` (matches R `did::att_gt(allow_unbalanced_panel=TRUE)` / `DRDID::reg_did_rc`: ATT bit-exact on cells AND dynamic aggregation via fixed unit-cohort-mass `pg` + a per-unit WIF; SE up to the documented CR1 `sqrt(G/(G-1))` factor). The earlier "weighting" framing was a mis-diagnosis — on unbalanced panels the dominant divergence from R is the *estimator* (within-cell differencing vs RC-on-pooled-obs), not only the weighting; both are resolved by the flag. The DEFAULT path keeps within-cell differencing as a documented design choice and now emits a `UserWarning` on unbalanced input (no-silent-failures). **Remaining deferred:** `survey_design=` × `allow_unbalanced_panel=` (per-obs vs per-unit weight resolution — currently fail-closed `NotImplementedError`); and covariate / ipw / dr × the flag R-parity verification (the RC path supports them; the committed golden covers `reg` no-cov). | `staggered.py`, `staggered_aggregation.py` | SE-audit D3 | Low | | CallawaySantAnna event-study bucket/weight construction is duplicated between the analytical aggregator (`staggered_aggregation.py::_aggregate_event_study`) and the multiplier bootstrap (`staggered_bootstrap.py`): both group (g,t) by `e = t - g`, apply the finite/NaN/reference masks, and read cohort weights. Both already consume the same source-materialized universal reference cells (so they agree), but the bucket logic is copy-pasted. Extract one shared helper returning per-event-time buckets (finite cells, NaN cells, reference flags, cohort weights, combined-IF inputs) used by both. Pure refactor; gate on byte-identical analytical + bootstrap output. | `staggered_aggregation.py`, `staggered_bootstrap.py` | SE-audit D3 | Low | | `StackedDiD` survey re-resolution intra-file dedup (raw-weight extraction ×3, compose-normalize ×3, resolve-on-stacked ×2). The cross-estimator ContinuousDiD/EfficientDiD panel-to-unit collapse consolidation LANDED (#226 shared helpers `ResolvedSurveyDesign.subset_to_units_by_row_idx` / `build_unit_first_row_index`); StackedDiD is deliberately NOT on that path (control units are duplicated across sub-experiments, so it re-resolves at stacked granularity rather than collapsing to one row per unit). The residual is stacked-specific, low value, and touches the numerically-sensitive composed-weight renormalization. Post-filter re-resolution / metadata-recompute unification across the three estimators was assessed and is not warranted — they use genuinely different mechanisms and already delegate to shared `_resolve_survey_for_fit` / `compute_survey_metadata`. | `stacked_did.py` | #226 | Low | | `solve_ols` residual memory floor is inherent to SVD-based lstsq after the PR-E marshalling slim (rust-side allocator high-water 15.32 -> 7.81 GB on the 2.4M x 130 clustered solve; remaining = one fused equilibrated input copy + thin-SVD U transient + vcov scores block + LAPACK gelsd internals on the python path). Further reduction requires the tall-skinny-QR algorithm change - same trade-off as the parked QR-reuse row below (library-wide last-digit perturbation + loses the second independent collinearity detector); the two land together if ever reopened. Diagnosis + measurements in docs/performance-plan.md. | `rust/src/linalg.rs::solve_ols`, `diff_diff/linalg.py` | PR-E | Low | diff --git a/benchmarks/R/generate_cs_unbalanced_golden.R b/benchmarks/R/generate_cs_unbalanced_golden.R new file mode 100644 index 00000000..8a164ed2 --- /dev/null +++ b/benchmarks/R/generate_cs_unbalanced_golden.R @@ -0,0 +1,96 @@ +#!/usr/bin/env Rscript +# Golden values for CallawaySantAnna(allow_unbalanced_panel=True) parity vs +# R did::att_gt(allow_unbalanced_panel=TRUE) + aggte(type="dynamic"). +# +# Usage: Rscript benchmarks/R/generate_cs_unbalanced_golden.R +# Output: benchmarks/data/cs_unbalanced_golden.json +# +# R's allow_unbalanced_panel=TRUE sets panel=FALSE and runs the repeated-cross- +# section levels estimator (DRDID::reg_did_rc) on the pooled observations. +# diff-diff matches the ATT bit-for-bit; the analytical SE matches up to the +# CR1 G/(G-1) finite-sample factor diff-diff's cluster path applies (see the +# test + REGISTRY note). The panel is deliberately UNBALANCED (attrition in +# cohort 3 at t>=4) so a cell's valid-unit count < its cohort mass, exercising +# the obs-vs-unit pg weighting and the RC-on-panel estimator. +suppressMessages({ + library(did) + library(jsonlite) +}) + +set.seed(20260706) + +make_panel <- function() { + rows <- list() + k <- 1L + add_unit <- function(u, g, drop_late) { + ufe <- rnorm(1) + for (t in 1:5) { + if (drop_late && t >= 4 && runif(1) < 0.40) next + post <- if (g != 0 && t >= g) 1 else 0 + eff <- if (g == 3 && post) 1.0 * (t - g + 1) else if (g == 4 && post) 2.0 * (t - g + 1) else 0 + rows[[k]] <<- data.frame(unit = u, period = t, g = g, y = ufe + 0.3 * t + eff + rnorm(1, 0, 0.5)) + k <<- k + 1L + } + } + u <- 0L + for (i in 1:50) { add_unit(u, 3, TRUE); u <- u + 1L } + for (i in 1:50) { add_unit(u, 4, FALSE); u <- u + 1L } + for (i in 1:100) { add_unit(u, 0, FALSE); u <- u + 1L } + do.call(rbind, rows) +} + +df <- make_panel() + +out <- att_gt( + yname = "y", tname = "period", idname = "unit", gname = "g", data = df, + control_group = "nevertreated", est_method = "reg", + allow_unbalanced_panel = TRUE, bstrap = FALSE, cband = FALSE +) +agg <- aggte(out, type = "dynamic", na.rm = TRUE, bstrap = FALSE, cband = FALSE) +agg_simple <- aggte(out, type = "simple", na.rm = TRUE, bstrap = FALSE, cband = FALSE) +agg_group <- aggte(out, type = "group", na.rm = TRUE, bstrap = FALSE, cband = FALSE) + +golden <- list( + meta = list( + did_version = as.character(packageVersion("did")), + n_units = length(unique(df$unit)), + note = paste0( + "allow_unbalanced_panel=TRUE (panel=FALSE -> DRDID::reg_did_rc on pooled ", + "obs). ATT bit-exact vs diff-diff; SE parity up to the CR1 G/(G-1) factor." + ) + ), + data = list( + unit = as.numeric(df$unit), + period = as.numeric(df$period), + first_treat = as.numeric(df$g), + outcome = as.numeric(df$y) + ), + cells = list( + g = as.numeric(out$group), + t = as.numeric(out$t), + att = as.numeric(out$att), + se = as.numeric(out$se) + ), + event_study = list( + egt = as.numeric(agg$egt), + att = as.numeric(agg$att.egt), + se = as.numeric(agg$se.egt), + overall_att = as.numeric(agg$overall.att), + overall_se = as.numeric(agg$overall.se) + ), + simple = list( + overall_att = as.numeric(agg_simple$overall.att), + overall_se = as.numeric(agg_simple$overall.se) + ), + group = list( + egt = as.numeric(agg_group$egt), + att = as.numeric(agg_group$att.egt), + se = as.numeric(agg_group$se.egt), + overall_att = as.numeric(agg_group$overall.att), + overall_se = as.numeric(agg_group$overall.se) + ) +) + +out_path <- file.path("benchmarks", "data", "cs_unbalanced_golden.json") +writeLines(toJSON(golden, auto_unbox = TRUE, digits = 16, pretty = TRUE), out_path) +cat("Wrote", out_path, "\n") diff --git a/benchmarks/data/cs_unbalanced_golden.json b/benchmarks/data/cs_unbalanced_golden.json new file mode 100644 index 00000000..5c5fd104 --- /dev/null +++ b/benchmarks/data/cs_unbalanced_golden.json @@ -0,0 +1,37 @@ +{ + "meta": { + "did_version": "2.5.1", + "n_units": 200, + "note": "allow_unbalanced_panel=TRUE (panel=FALSE -> DRDID::reg_did_rc on pooled obs). ATT bit-exact vs diff-diff; SE parity up to the CR1 G/(G-1) factor." + }, + "data": { + "unit": [0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 7, 7, 7, 7, 8, 8, 8, 8, 9, 9, 9, 9, 9, 10, 10, 10, 10, 11, 11, 11, 11, 12, 12, 12, 12, 13, 13, 13, 13, 13, 14, 14, 14, 14, 15, 15, 15, 15, 15, 16, 16, 16, 16, 17, 17, 17, 17, 17, 18, 18, 18, 18, 18, 19, 19, 19, 20, 20, 20, 20, 21, 21, 21, 21, 22, 22, 22, 23, 23, 23, 23, 23, 24, 24, 24, 24, 24, 25, 25, 25, 25, 26, 26, 26, 26, 27, 27, 27, 27, 27, 28, 28, 28, 28, 29, 29, 29, 29, 29, 30, 30, 30, 30, 31, 31, 31, 31, 32, 32, 32, 32, 33, 33, 33, 33, 34, 34, 34, 34, 35, 35, 35, 35, 36, 36, 36, 36, 36, 37, 37, 37, 38, 38, 38, 38, 38, 39, 39, 39, 39, 40, 40, 40, 40, 41, 41, 41, 41, 42, 42, 42, 42, 43, 43, 43, 43, 44, 44, 44, 44, 45, 45, 45, 45, 46, 46, 46, 46, 47, 47, 47, 47, 47, 48, 48, 48, 48, 48, 49, 49, 49, 50, 50, 50, 50, 50, 51, 51, 51, 51, 51, 52, 52, 52, 52, 52, 53, 53, 53, 53, 53, 54, 54, 54, 54, 54, 55, 55, 55, 55, 55, 56, 56, 56, 56, 56, 57, 57, 57, 57, 57, 58, 58, 58, 58, 58, 59, 59, 59, 59, 59, 60, 60, 60, 60, 60, 61, 61, 61, 61, 61, 62, 62, 62, 62, 62, 63, 63, 63, 63, 63, 64, 64, 64, 64, 64, 65, 65, 65, 65, 65, 66, 66, 66, 66, 66, 67, 67, 67, 67, 67, 68, 68, 68, 68, 68, 69, 69, 69, 69, 69, 70, 70, 70, 70, 70, 71, 71, 71, 71, 71, 72, 72, 72, 72, 72, 73, 73, 73, 73, 73, 74, 74, 74, 74, 74, 75, 75, 75, 75, 75, 76, 76, 76, 76, 76, 77, 77, 77, 77, 77, 78, 78, 78, 78, 78, 79, 79, 79, 79, 79, 80, 80, 80, 80, 80, 81, 81, 81, 81, 81, 82, 82, 82, 82, 82, 83, 83, 83, 83, 83, 84, 84, 84, 84, 84, 85, 85, 85, 85, 85, 86, 86, 86, 86, 86, 87, 87, 87, 87, 87, 88, 88, 88, 88, 88, 89, 89, 89, 89, 89, 90, 90, 90, 90, 90, 91, 91, 91, 91, 91, 92, 92, 92, 92, 92, 93, 93, 93, 93, 93, 94, 94, 94, 94, 94, 95, 95, 95, 95, 95, 96, 96, 96, 96, 96, 97, 97, 97, 97, 97, 98, 98, 98, 98, 98, 99, 99, 99, 99, 99, 100, 100, 100, 100, 100, 101, 101, 101, 101, 101, 102, 102, 102, 102, 102, 103, 103, 103, 103, 103, 104, 104, 104, 104, 104, 105, 105, 105, 105, 105, 106, 106, 106, 106, 106, 107, 107, 107, 107, 107, 108, 108, 108, 108, 108, 109, 109, 109, 109, 109, 110, 110, 110, 110, 110, 111, 111, 111, 111, 111, 112, 112, 112, 112, 112, 113, 113, 113, 113, 113, 114, 114, 114, 114, 114, 115, 115, 115, 115, 115, 116, 116, 116, 116, 116, 117, 117, 117, 117, 117, 118, 118, 118, 118, 118, 119, 119, 119, 119, 119, 120, 120, 120, 120, 120, 121, 121, 121, 121, 121, 122, 122, 122, 122, 122, 123, 123, 123, 123, 123, 124, 124, 124, 124, 124, 125, 125, 125, 125, 125, 126, 126, 126, 126, 126, 127, 127, 127, 127, 127, 128, 128, 128, 128, 128, 129, 129, 129, 129, 129, 130, 130, 130, 130, 130, 131, 131, 131, 131, 131, 132, 132, 132, 132, 132, 133, 133, 133, 133, 133, 134, 134, 134, 134, 134, 135, 135, 135, 135, 135, 136, 136, 136, 136, 136, 137, 137, 137, 137, 137, 138, 138, 138, 138, 138, 139, 139, 139, 139, 139, 140, 140, 140, 140, 140, 141, 141, 141, 141, 141, 142, 142, 142, 142, 142, 143, 143, 143, 143, 143, 144, 144, 144, 144, 144, 145, 145, 145, 145, 145, 146, 146, 146, 146, 146, 147, 147, 147, 147, 147, 148, 148, 148, 148, 148, 149, 149, 149, 149, 149, 150, 150, 150, 150, 150, 151, 151, 151, 151, 151, 152, 152, 152, 152, 152, 153, 153, 153, 153, 153, 154, 154, 154, 154, 154, 155, 155, 155, 155, 155, 156, 156, 156, 156, 156, 157, 157, 157, 157, 157, 158, 158, 158, 158, 158, 159, 159, 159, 159, 159, 160, 160, 160, 160, 160, 161, 161, 161, 161, 161, 162, 162, 162, 162, 162, 163, 163, 163, 163, 163, 164, 164, 164, 164, 164, 165, 165, 165, 165, 165, 166, 166, 166, 166, 166, 167, 167, 167, 167, 167, 168, 168, 168, 168, 168, 169, 169, 169, 169, 169, 170, 170, 170, 170, 170, 171, 171, 171, 171, 171, 172, 172, 172, 172, 172, 173, 173, 173, 173, 173, 174, 174, 174, 174, 174, 175, 175, 175, 175, 175, 176, 176, 176, 176, 176, 177, 177, 177, 177, 177, 178, 178, 178, 178, 178, 179, 179, 179, 179, 179, 180, 180, 180, 180, 180, 181, 181, 181, 181, 181, 182, 182, 182, 182, 182, 183, 183, 183, 183, 183, 184, 184, 184, 184, 184, 185, 185, 185, 185, 185, 186, 186, 186, 186, 186, 187, 187, 187, 187, 187, 188, 188, 188, 188, 188, 189, 189, 189, 189, 189, 190, 190, 190, 190, 190, 191, 191, 191, 191, 191, 192, 192, 192, 192, 192, 193, 193, 193, 193, 193, 194, 194, 194, 194, 194, 195, 195, 195, 195, 195, 196, 196, 196, 196, 196, 197, 197, 197, 197, 197, 198, 198, 198, 198, 198, 199, 199, 199, 199, 199], + "period": [1, 2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 5, 1, 2, 3, 5, 1, 2, 3, 4, 5, 1, 2, 3, 5, 1, 2, 3, 5, 1, 2, 3, 4, 1, 2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4, 5, 1, 2, 3, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 1, 2, 3, 5, 1, 2, 3, 4, 1, 2, 3, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 5, 1, 2, 3, 5, 1, 2, 3, 4, 5, 1, 2, 3, 5, 1, 2, 3, 4, 5, 1, 2, 3, 5, 1, 2, 3, 5, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4, 5, 1, 2, 3, 1, 2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 5, 1, 2, 3, 5, 1, 2, 3, 4, 1, 2, 3, 5, 1, 2, 3, 5, 1, 2, 3, 4, 1, 2, 3, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5], + "first_treat": [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + "outcome": [0.55163821453080797, 0.39065007574849858, 2.7978193128668969, 3.3680896884283404, 4.4924673189838158, 0.28932853056245339, 0.77578550457491868, 0.81170043772903122, 3.2910754022319657, 1.69505300786762, 1.1715149381629328, 1.8392777005695331, 3.6086446401405277, -1.2021438575152441, -0.55188730107246231, 0.61582299497411275, 3.3057602011875251, 3.4626355063777998, -1.1779251502267203, -0.6339778306930236, 1.9960641042864611, 3.1248874167054357, 0.34679757724695159, 2.0739995526599042, 3.6241587918989402, 4.1305464519443014, 6.0115315432737813, 1.442810887057671, 1.0991042969656719, 2.9747283442783647, 4.1470833362695361, 5.8537319437576532, 0.44199288625532318, 0.33308674229917523, 1.9819024648912311, 4.6085902660584726, 1.5002575452854443, 1.9173869300966202, 2.6896238600277331, 4.9896680969896812, -0.12983630770289367, -0.048758038118745017, 1.4626441105226251, 2.6024543829281805, 4.2709211598947086, 1.6834560896263997, 1.4242646735581961, 2.7252174861866632, 5.1830214014759983, 1.1166140078643974, 0.047913880455698443, 3.1630359803532042, 4.8052679195002845, 2.8030593815687137, 2.6585920825837914, 4.0114210888453341, 5.0847867731650052, 1.1344905326068608, 1.1128693438340205, 2.3174196204989088, 4.0040373237105245, 5.5605844249633574, 2.1885510905907926, 1.6542989759677615, 3.7572912057930816, 4.6922103646722038, -0.4874240414003011, -0.33725948694931168, 0.53742315231535043, 2.7579744773861354, 4.1100082425094397, 1.0592405077643012, 2.8001805605831098, 3.2957608675641219, 5.1548188972967477, -1.6390406328430702, -0.8069439514628568, -0.12990343033429508, 2.0744768244188321, 3.0510701804628049, 0.52176985453932601, 2.1576540484895448, 2.6182167794756128, 3.5912858843796478, 5.721864615277461, 1.9560583764353821, 1.1422295539814231, 3.6274583071828168, -1.086036068858756, -0.96859536438591876, 1.5097619888303786, 3.6864826762425018, 0.22843427532956317, 0.55405849914761496, 1.3647836507952387, 3.2524011627674043, 0.087914140826010706, -0.59053878170964658, 1.1517145417973578, -0.37254434572924289, -0.16132424844260512, 1.3883426678403659, 2.7107106753238739, 3.7555865791958047, 0.085580619676102698, 1.0403009250758579, 1.8485337092585066, 2.0241849797869111, 4.3295151603840267, 0.11910601442882431, 0.072146148859411008, 2.1258449619874469, 4.4620050424522857, 1.6886687174700101, 2.1519934653173252, 3.462763337220808, 5.8101260464192936, -0.082643123398976392, -0.47281005541337967, 1.0556146046621633, 3.2670759542788925, 3.7803104475453968, -1.5853201584162897, -2.2821408212139884, 0.10308994583314808, 1.9865307269309178, 0.86562407921713824, 1.317599957781356, 2.8790669737939814, 3.1813043394579186, 4.5148038075000567, 2.2915273743144633, 0.87313174463524834, 3.0530444428212147, 5.3903806698138066, 2.1261874893197716, 2.3786887686855942, 3.2938790119507275, 5.7995080446377054, 0.23961908633200313, 0.76015582296190909, 1.600551018040953, 4.6659279021158131, -1.3870438590865577, -1.0793774810651615, 0.6574297311083438, 1.1193956997845922, 0.46438272143475096, -0.23362961370690621, 0.66364402162110969, 3.3274620759236346, -1.0197643714864471, 1.2548160619990258, 1.7456334985688939, 3.0226561403325243, -0.12640228721215596, 1.6832441980846302, 1.8929974157761082, 3.5635847865804018, 4.8478065252141391, 0.89171875241667475, 2.2399048838372377, 2.5353187785935809, 1.5935525954690335, 2.1930141712773032, 3.5970582705936271, 5.9055770827039407, 5.9091633132954087, 1.8814116609614948, 1.7941506389544009, 3.3934877078851136, 5.2439509251546443, -0.1725799200672811, -0.018852046640865858, 1.5133143170773269, 4.1504018888039038, 0.40415111292339068, 0.23269084765218628, 1.3421658333094502, 4.5681091716251503, 0.61393705717202463, 0.21861621685105084, 1.8472000941029332, 3.2179639239535378, 0.00016517628887049995, 0.23093914727941028, 1.4169859074745219, 3.0918313409607312, -1.3871380871014782, 0.50083860435210115, 2.3974781433827959, 4.0340789543677751, -0.36650522111931783, 0.93211401205921218, 1.1794007940261579, 2.7634541523545066, 1.6623893626353305, 1.4230521927565924, 2.6502543647439119, 5.3872670075667521, 1.0688793412122963, 0.85085376483208797, 3.1897406623770808, 3.9622424912895924, 6.2722695549329508, 1.5161514747997646, 1.3510632170042558, 3.4678673562296871, 4.7910347354841978, 6.003696781609432, 0.22843223679756636, 1.1211012339137822, 2.5819760720657783, -0.99637579282207978, 0.64283548446327843, -0.27335489777694155, 2.1817900708105418, 4.884004398910867, 2.835622756735765, 3.06664208778006, 2.6959846323381984, 4.6045489035709313, 7.6956910446871252, -1.7876699200363322, -1.3762992604902859, -1.8512822859165861, 1.036349175683478, 3.7634080346083518, -0.80442845491586557, -0.47814713152316757, 0.70389198668308484, 1.8444193826614748, 3.6546175621548289, -0.69341665894867799, -0.96183188795567642, 0.14194947060404525, 2.5524684743751935, 4.673689759007658, -0.77682370710463511, 0.53324092914993215, 0.031358446918600784, 1.9810647606075118, 5.2233550301375908, 1.461772791384929, 1.7591687138416086, 1.786010606605549, 4.0505570373453912, 5.8635150738785544, 1.4220175259056886, 1.0384858175702982, 1.625371768839591, 3.7160935937248318, 6.7758598138064015, 2.0313338271360695, 2.5098767540144968, 1.2742103117306074, 4.1824373250933613, 7.1280704847008129, 0.26537538936571403, 0.74696385940167809, 1.0835883660661942, 4.0586325559482379, 5.97400885513016, -0.70794764518824249, 0.63396481187264542, 0.6129063676077896, 3.1305921236395302, 6.5722654549639499, -0.7061887088455745, 0.35245068887904701, 0.54753763228411001, 3.6467921668462493, 4.5244117408686986, -0.34471181392105721, -0.26473699792593414, -0.3298667332970675, 2.635466561806838, 4.5936882617814874, 1.1564390314282675, 0.72761169714337659, 0.76681368652634863, 3.4576201566671139, 5.9432482802667526, -1.818350296346299, -2.4231370866599207, -1.4584220760022568, 0.50437829693939951, 2.7189913871091811, 1.259576672040019, 1.7827292162857435, 2.2027895177631502, 4.5073284522197978, 6.251885995576103, 0.48748222644563549, 0.60392815774330744, 0.47269125788037214, 2.7500159164459088, 3.9743720791725856, 0.98688738028284995, 0.86658602423224818, 1.6115572938336786, 3.3020300893731425, 5.0455804705798242, -0.15255693832570116, 0.351688582565098, -0.15116218996617514, 2.8268994267573251, 5.4428563501065721, 0.99650016811745235, 3.0240380910607332, 1.6282526445066323, 4.0257637334750225, 6.8119015428120839, 0.73007463803261752, 0.56791385672904748, 0.50082578636258068, 3.2142948934490274, 5.664063798686275, 1.8750237444477307, 1.750425865224432, 0.91222721254557881, 3.9525846853780777, 6.4748006672939304, -0.040398464192992978, 0.13198339307983606, 0.60197752942308691, 2.8619466226070167, 5.4608635701070725, 0.33164335351274732, 0.87345620917205602, 1.1474130723493265, 3.9317817684664984, 5.6218022086559074, 1.9339138079724867, 2.3544083793357191, 2.9231920230163517, 3.9040425411420201, 6.790291580907672, -0.43138861665873562, -0.16903794307380682, -1.0134532004713825, 2.1721652287628728, 3.7815297233568885, -0.89807827482947089, -0.03436566911222011, -0.6495422872342691, 2.1041178747883027, 4.8773941656311948, 0.23278219582211029, 0.42200226389366274, 0.03374555143950031, 2.4804086657818045, 5.2747716245495821, 1.4242433619795642, 1.9600078997435861, 2.2040737707392424, 4.7005821490218782, 7.9917614803667112, -1.3691856759670546, -0.77106714848100943, -0.58958839366663318, 1.9655431382895829, 3.8786872726750743, 1.1903997633595373, 0.96814483813595142, 2.3412288502255909, 3.747909934680234, 5.5547755204466362, 0.50757001480997388, 0.55550709965366374, 0.39722837066622763, 3.6565458062007865, 5.9495311956274737, 0.93086271256688902, 1.9842736404020682, 1.8502264341001156, 3.4936544349027305, 6.3128365066997025, 0.94722521976988983, 1.6927845701325386, 0.96175222050971676, 3.9990492132047231, 5.8327426034062722, -0.43600359544296641, -0.4274396794835395, -0.038621302309943067, 2.5878955040473803, 3.8544900728697353, 0.51378734318285413, 0.57897741501716182, 0.19241079959198282, 3.3159026308221562, 6.770841856696153, -0.15674661138373816, 0.071714566495063109, 0.54807050804721968, 1.8796574211619079, 4.2817077367296648, -2.2073923298412224, -1.3309195313905435, -1.4813408991322328, 0.91702294693596675, 4.2912731600498688, 0.266764368136857, -0.16397068609252585, 0.80490832436905868, 2.8129761089923351, 5.0322572781967256, 0.65150230292457534, 1.8806552549566078, 1.9845503164382083, 3.4993719000008934, 5.5320772518693353, -0.82901856006296004, -0.11297763940612743, 0.046746840438465137, 2.1014977686185397, 4.9958316367429161, -0.43559506537342108, 0.10427744331409752, 0.73941300802485577, 2.5390933379210114, 4.7510360300824033, 1.4329305982059204, 1.0257352505241797, 1.5935279034373711, 4.5810552159511406, 7.1550617450884966, 2.6269736058644595, 2.7437520005790512, 2.0674048454419816, 4.6961712340539821, 6.7850739167864367, -0.4682367707468047, 0.91602260074597486, 1.7123385207277135, 3.3487812508215549, 6.6932600953503636, -0.13442430455718082, 0.018829713299711143, -0.20616305539378849, 2.5536822996586444, 5.1236246608462332, 0.55579992378068865, 0.96337756258654172, 1.6570991057453226, 4.2928376551181771, 4.83633862661876, -1.0163090777321853, 0.073211211981028812, 0.0090071284011860042, 1.6762891254488288, 4.8873576984482003, 1.0704239763764134, 1.8160583132547612, 1.1721826702027338, 4.3011994067117882, 7.21961591815072, 2.0997928748587626, 1.6172801842359532, 2.160668250690851, 4.3685783128190474, 7.0450988057457256, 0.5784309903100181, -0.030430293688701399, 0.42375632972589611, 0.65886960942747652, 0.8893734259839795, -1.1157585218293082, -0.87880302980824976, -0.72055395907553166, 0.67416893605903971, -0.1893542458971206, 0.53248985753523048, 0.83283078142528577, 1.3307256702580574, 0.90375180448678782, 0.47004921650423059, 2.6385633921826246, 3.2600990757066857, 2.8344508231909429, 2.8025020646351946, 3.3228739676018142, 1.2679401450924788, 0.23544090663601341, 1.0390879544854636, 1.0308601548244107, 1.2960062935875625, 0.44709874462179594, 0.87607402856258454, 0.66462674975752589, 0.88718833934330854, 1.8528782424071024, 3.0466230599825108, 3.2879726083928462, 3.6616908111466251, 4.446274982213497, 4.0665914090296136, -1.7303993358778949, -0.69318323103476054, 0.29563834934903921, -0.58968927518559167, -0.54622200739581539, -0.20502368340903276, 2.0149121049362178, 2.1874123426383241, 1.6302086391292996, 1.9203242148458965, -1.4895190597316799, -0.64276449361336807, -1.0362217320231937, -0.084837574424791007, -0.96717360756771453, 0.62967488628112989, 0.77255368265316604, 0.88800410522278339, 2.1734149971232064, 1.4241510946771299, 1.5765887140578412, 1.3061051589934753, 2.0255443135525137, 2.1232027886641243, 2.1545301747908652, -0.95740759725425839, -0.36314797116907077, -0.45816539204397044, 0.63633152653626257, 0.66000458081760693, 2.0701362587666368, 2.3356213447668845, 1.3013402804711647, 2.1641082724878329, 2.9371350350310776, 0.27934848540169499, 0.51715184044591145, 0.42880267342910505, 1.0995696736252536, 1.442802984760724, 0.74478184886742071, 1.8989077311876466, 2.3859259269249069, 2.8158643443875118, 2.4216388596626941, -0.26296196474947731, 1.1774526978215722, 1.7204775221450146, 1.4503635356350584, 2.0375941286045931, 0.42412023341682825, 0.24073872704882915, 1.1002979907747583, 1.0096425037136709, 1.7431783478928216, 2.5546884066287916, 2.7938119291853152, 2.9071079282694692, 3.8838473423840876, 2.788387822511512, 0.37634125813321201, 0.22539015292778233, -0.5316684095400136, 1.443534551190222, 1.0647836033981173, -0.52828850736696786, 0.20545682819202654, 0.37764643167107914, 0.20109396080442848, 0.90124671817036828, 0.7721340344921932, 0.018108281832134088, 1.3321179735130313, 0.86437044439858268, 1.2275835220153235, -1.5883803904018201, -0.94742284110544039, -1.4783222112653815, -1.0869235091016283, -1.0441449027130518, 2.9009985292135845, 1.5017968358434608, 2.9285027846151022, 2.9600827080795931, 3.4696578680494707, 1.135204596855937, 2.2198674240854492, 1.3192811878955197, 3.2302797522599089, 3.1620017260417486, 1.396483746416932, 2.498575336850116, 2.4157646148438201, 1.6286415113670787, 3.2495243016403004, 0.77013213466480224, 1.1395624562639377, 1.8781017880666491, 2.2726207167548065, 1.6631279627982334, -0.58134905565823281, -0.36234726159076502, 0.063539698299578895, -0.34469016391581353, 0.45966525043167905, 2.7856957329316359, 3.3827674516386761, 3.594636647599017, 3.4682477307213402, 4.0156240475953622, 0.8740998575549388, 0.60236478859082232, 1.496823939548598, 1.3689880666360159, 1.2486982213959972, -1.0357829923838793, 0.0055357914244873463, -1.2208008914841522, -1.4266406079631102, 0.083715548236624276, -0.3677245379978264, -0.4467837698630705, -0.12194140621618835, 0.29266939375118917, 0.14990600194909215, -0.46533837016281265, 0.48606292585677058, -0.22789612038208223, 0.20647613605429388, 0.39320593371498591, 1.1580077998293983, 1.0744233003464059, 1.165455785558581, 2.5456577705886905, 2.4539143634580132, -0.5532927784240762, -0.12610179915729869, 0.19485381978399158, 0.80966073430510588, 1.4106931206939581, 2.8707945099311272, 2.9965117633092837, 2.4813062904754988, 3.1007467267171971, 3.6662217209811958, 1.750225186057345, 1.7823045754057305, 2.7076951602480017, 3.6354661463148501, 2.4769931976774613, 1.8040887837264448, 1.4086761414766411, 1.4586609550216834, 1.7002017051817448, 1.4904202034235641, 0.26055441340331964, 0.81251164947314614, 0.83466592240617643, 1.7398312357621601, 2.4177963519701589, 0.35751109126973202, 1.0450407886775297, 1.7107473460578166, 1.6420551749463801, 1.0915287424186841, 0.74432986324683137, 1.1242691829953417, 1.7158765856653659, 1.5121531206938521, 2.1821148740032283, -0.1337106097212637, 0.36663918818079055, -0.47217012572845174, 0.40754612399440693, 0.54511237428527781, 0.85224220257822014, 0.049559497428789956, 1.4392449405706451, 2.4768611621879266, 1.704086639731581, -0.64055731160772778, 0.47626314564729416, 1.2785389203710089, 0.79272242629504575, 1.9011828187070026, -0.85891855019932317, -0.95093513388376427, -1.4155702377910213, 0.33686277882613813, 0.23995631218030444, -2.0440184200043636, -1.1658956769216404, -0.87552539278165631, -1.1747741095900053, -0.84529555175655213, -1.1702151860246077, -0.56032646435289779, 0.56775919868581193, 1.0559484100168333, 1.2317212243935685, -0.81615488349387277, -0.40710926553657267, 0.72818301314402911, 1.3487478471004335, 1.3459575248241566, 0.83485692062467942, 0.67087695938938352, 0.56729150854553478, 0.34193013230278879, 2.5241783736434042, -0.75779997786209008, -0.75575091294790164, 0.21010964444217517, 0.59123037776251963, 0.025719240724810799, -0.42585087468228955, 1.5336530592210205, 1.675810208820989, 1.3527909744842361, 1.546827151542014, 1.183329283919305, 1.9209442246703456, 3.2729106074180745, 2.7718979215878905, 3.5200738388206019, 0.33231652197131833, -0.44596115854543195, -0.67212080304465649, 0.32988867263004312, -0.096235870099458998, 1.2442323287384718, 0.62776869076215946, 0.6842756115894808, 1.2138073465060595, 1.7225487755422948, -0.11584054630254326, 0.060150271393291765, 1.3639070271884259, 1.7068407484828145, 1.1401377447161902, 0.10249518908572447, 0.38207344342622451, 1.2905398424319252, 1.4347917656604661, 1.5490291618491792, 0.20409137418126105, 0.25625134989467779, 1.3408659916980774, 1.4685122665587751, 1.6631785334794884, 0.51499177629782522, 0.79959428368139496, 1.1521507820381578, 1.6981458426005833, 2.6865101501344024, 0.41550001256966529, 0.32091840913839209, 1.2660131391180893, 1.0749214949785144, 2.3751306971901633, 1.0336820140914615, 0.73939986230755794, 1.0189946540759083, 1.5854674156735962, 1.0273894612809502, 1.039394894699406, 0.79357611910632142, 0.86424877546715617, 1.0244325239867389, 1.7798909381600909, -1.6364322984448783, -0.53083377030138246, -0.53468814051271774, 0.9482036929701827, 0.87346434049337474, 2.5769437761693883, 2.3313531976955186, 1.2669146232045523, 2.5306723183417619, 1.8522558895836436, 3.3423580060194475, 3.9106406949179862, 4.4308217621418269, 4.0995385486363212, 4.5969157263768494, -0.53429828809422575, -0.73128955538510254, 0.56023448020890065, 1.3283966933803351, 0.076424953923678873, 2.0565787387304191, 2.5565117316538459, 1.8700927473377036, 2.9817514814589532, 3.3237612606873501, 1.5410964638076847, 3.2409053576627573, 2.3246426589897666, 3.3836065913419611, 4.0428881838796178, 0.82966429185243751, 0.36678456701760298, 0.42419640538132197, 1.3489503179481632, 1.3472808825124682, 1.0085271609266995, 1.9256172300671341, 1.5594086990260645, 0.80084194652563467, 2.5442940572653625, 1.1037676203585134, 1.5349208842523132, 1.7201640983908177, 1.2253239798378019, 2.1481367291023896, 2.2326544069266339, 1.3306570818324985, 1.4146611140229777, 1.6153381065898311, 2.4172931725998561, 1.9188444953924106, 1.3458025967920306, 1.4666453794676004, 1.8259529405638315, 2.4722763149883358, -1.2684588834953505, -0.9879468906068446, -1.1694644947583535, -0.88154303738032747, -0.35581164024923551, 1.1037921681368781, 0.56903930442357142, 1.2034315685915786, 1.1527241064150313, 3.229839708247721, -0.25623776875060639, -0.036817710196296954, 0.57168908975378219, 0.22279451737867995, 0.56786043645088791, 0.75055083398521927, 1.4091767119661249, 1.6149789725534347, 1.4402259453659614, 2.1807900572219396, -0.71470430424725695, -0.35679788331352696, 0.075249597872698748, 0.96895395009178997, -0.043892339930590107, 1.4678284222771909, 0.2075201315787642, 0.70400445171754167, 2.0322732594433361, 0.8527219190745301, 0.66228008375589575, 0.79304739812457536, 1.4007469402865937, 1.9637402412148663, 1.2253937873679575, -1.2975012697223451, -0.72072976563130531, -0.43806589816419483, -0.37777254871054067, 1.0404390050020833, -1.7009597706908008, -2.2398602632581852, -0.3579503924865759, 0.10051929622583544, -0.079583011959706351, -1.3686780017247688, -1.5112517390290912, -2.3512156034404939, -0.68459240037545122, -0.38940859149806656, -1.8576745257973426, -1.3776846867411945, -0.62567439704572436, -0.80006329572784385, -0.12050614203439436, -1.8139510387325926, -1.634352057530303, -1.0336771335390977, -1.9413210910844205, -1.584071937985982, 1.2196412725914998, 1.6721575068842567, 1.8015624523980989, 2.5794240782876057, 2.5691111595429637, 2.1868592478506357, 1.9388056020244235, 3.0735639652103397, 4.1744713923815118, 2.4668620891951374, 0.72571907004966396, 1.3597275792563206, 1.6990217256891906, 1.6200073684925285, 1.6159745196827366, -1.1022358718356102, -1.0701242012235253, 0.020796193079299308, -0.75668625978184367, 0.046420266969644713, 0.29020314015664739, 1.2214956359265874, 0.9915450328729164, 2.5830566242284578, 2.8275742208354, -0.65177476694807934, -0.82069136705929391, -0.13831831352266794, -0.12873556839057543, 2.010991971351145, 0.17910229770288028, -0.2996948240163062, 1.1518484637709161, 0.87412673581870148, 1.2046773553403929, 0.66533404478195646, 0.36821675465618842, 1.4373197067295949, 1.3811198073443258, 2.471750284470112, 2.0330478547158752, 2.1206215570941338, 2.4268335785894779, 3.1255866460586197, 3.8641980158790661, -0.64723936510129443, 0.052865163590172792, -0.42666863945242384, 1.0392058266450652, 0.60982528574706685, 0.79211319025138738, 0.81771674189012855, 1.1519860850386783, 1.1147597732028831, 2.293370251799554, 0.50270994386915535, 0.52871463495855064, 1.6422610265512225, 1.7394164528700589, 1.7633574539057175, 1.3308351252288337, 1.1745070585882411, 1.346345075904946, 1.4854636405814376, 3.2304700161318207, -0.14219109698910093, 0.95127570285279894, 0.95562865816914799, 0.16237794983316123, 1.1860307286756397, 0.11488419162202392, 0.017614916492196153, 1.1411248711382069, 0.98398364110756464, 2.4897331496502209, -0.37905583274488808, -0.58269099464914753, 0.21668647192650545, 0.40683215503876424, 0.073572997835092746] + }, + "cells": { + "g": [3, 3, 3, 3, 4, 4, 4, 4], + "t": [2, 3, 4, 5, 2, 3, 4, 5], + "att": [0.031419477925689859, 1.0795162943988603, 2.1133518226847086, 3.0202671592981178, 0.17146546152483008, -0.26790091392588422, 2.0385657598869629, 4.1652407502122752], + "se": [0.13038507360189816, 0.11833070815931725, 0.18499992165315304, 0.14426780306320647, 0.10404928172836575, 0.10960941881338133, 0.10529614386033723, 0.13619665759579466] + }, + "event_study": { + "egt": [-2, -1, 0, 1, 2], + "att": [0.17146546152483008, -0.11824071800009718, 1.5590410271429116, 3.1392962864484919, 3.0202671592981178], + "se": [0.10404928172836575, 0.08235178064454754, 0.087721203057747971, 0.15306924462190835, 0.14426780306320641], + "overall_att": 2.5728681576298404, + "overall_se": 0.097608983575660169 + }, + "simple": { + "overall_att": 2.4833883572961848, + "overall_se": 0.094253486976800874 + }, + "group": { + "egt": [3, 4], + "att": [2.0710450921272288, 3.1019032550496188], + "se": [0.11305016015636252, 0.10548897353424015], + "overall_att": 2.5864741735884236, + "overall_se": 0.093219774066356093 + } +} diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index 3e353e95..fee3878f 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -384,6 +384,18 @@ class CallawaySantAnna( each period (stationarity). Uses cross-sectional DRDID (Sant'Anna & Zhao 2020, Section 4) with per-observation influence functions. + allow_unbalanced_panel : bool, default=False + When ``True`` and the input panel is unbalanced (some units are not + observed in every period), route the pooled observations through the + repeated-cross-section levels estimator (matching R + ``did::att_gt(allow_unbalanced_panel=TRUE)`` / ``DRDID::reg_did_rc``) + instead of within-cell panel differencing, and cluster the influence + function by unit for the standard error. **Inert on a balanced panel** + (results are byte-identical to the default). When ``False`` (default) + an unbalanced panel is handled by within-cell differencing and a + ``UserWarning`` is emitted. ATT matches R bit-for-bit; the SE matches + up to the documented CR1 ``sqrt(G/(G-1))`` finite-sample factor. + ``survey_design=`` combined with this flag raises ``NotImplementedError``. epv_threshold : float, default=10 Events Per Variable threshold for propensity score logit. When the ratio of minority-class observations to predictor @@ -498,6 +510,7 @@ def __init__( cband: bool = True, pscore_trim: float = 0.01, panel: bool = True, + allow_unbalanced_panel: bool = False, epv_threshold: float = 10, pscore_fallback: str = "error", vcov_type: str = "hc1", @@ -572,6 +585,13 @@ def __init__( self.cband = cband self.pscore_trim = pscore_trim self.panel = panel + # When True AND the input panel is unbalanced (some units unobserved in + # some periods), route through the repeated-cross-section (RC) levels + # estimator on the pooled observations — matching R + # `did::att_gt(allow_unbalanced_panel=TRUE)` (which sets panel=FALSE -> + # DRDID::reg_did_rc). Inert on a balanced panel (the default within-cell + # differencing path is byte-identical). See fit() for the routing. + self.allow_unbalanced_panel = allow_unbalanced_panel self.epv_threshold = epv_threshold self.pscore_fallback = pscore_fallback @@ -1842,6 +1862,130 @@ def fit( # second layer for the post-construction mutation path. self._validate_vcov_type(self.vcov_type) + # --- allow_unbalanced_panel routing (RC-on-panel = R's allow_unbalanced_panel) --- + # Detect an unbalanced panel (some units unobserved in some periods). + # Only meaningful for panel input; declared RC (panel=False) is handled + # by the branches below via `_use_rc`. + _is_unbalanced_panel = False + _finite_out = None + if self.panel and all(c in data.columns for c in (unit, time, outcome)): + # Complete-case balance (matches R `pre_process_did`, which + # complete-case-filters BEFORE deciding to flip to the RC path): a + # rectangular panel with a missing / non-finite OUTCOME in a cell is + # effectively unbalanced — the default panel path treats that cell as + # a missing panel cell (wide-pivot valid-mask) and silently does + # within-cell differencing. So detect unbalancedness on the + # finite-outcome rows, not just structural (unit, time) row presence. + _finite_out = np.isfinite( + pd.to_numeric(data[outcome], errors="coerce").to_numpy(dtype=np.float64) + ) + # Assess rectangularity on the complete-case frame for BOTH the cell + # count AND the units x periods denominator — the same frame the RC + # path estimates on (below). If a whole unit or a whole period is + # dropped by non-finite outcomes, the remaining sample may still be + # rectangular (balanced) and must stay inert, matching R's + # complete-case preprocessing. Using the unfiltered unit/period counts + # in the denominator would spuriously route such a panel to RC. + _cc = data.loc[_finite_out] + _is_unbalanced_panel = ( + int(_cc.drop_duplicates(subset=[unit, time]).shape[0]) + < _cc[unit].nunique() * _cc[time].nunique() + ) + # Route through the RC levels estimator ONLY when the flag is set AND the + # panel is actually unbalanced. Inert on a balanced panel (default + # within-cell differencing is byte-identical). This MATCHES R: R's + # pre_process_did recomputes `allow_unbalanced_panel` from the observed + # balance and only flips `panel <- FALSE` when the panel is genuinely + # unbalanced (verified: R keeps panel=TRUE on balanced data under the + # flag), so balanced-panel inertness is the R contract, not a deviation. + _route_as_rc = bool(self.allow_unbalanced_panel and _is_unbalanced_panel) + # `_use_rc` selects the RC precompute / per-cell loop / obs-level counting + # for BOTH declared RC (panel=False) and RC-routing (unbalanced + flag). + _use_rc = (not self.panel) or _route_as_rc + if _route_as_rc and survey_design is not None: + raise NotImplementedError( + "allow_unbalanced_panel=True is not yet supported together with " + "survey_design=. The RC-on-panel path carries per-observation " + "weights, whereas survey designs assume per-unit weights; the " + "combination needs a per-unit weight resolution that is not " + "implemented (deferred). Use a balanced panel, drop " + "survey_design=, or pass panel=False for genuine repeated " + "cross-sections." + ) + if _is_unbalanced_panel and not self.allow_unbalanced_panel: + warnings.warn( + "Unbalanced panel detected (some units are unobserved in some " + "periods). Each ATT(g,t) is estimated by within-cell panel " + "differencing on the units observed at BOTH the base period and " + "t — a valid but different estimand than R's repeated-cross-" + "section handling. Pass allow_unbalanced_panel=True to match R " + "`did::att_gt(allow_unbalanced_panel=TRUE)` (the DRDID " + "reg_did_rc levels estimator on the pooled observations).", + UserWarning, + stacklevel=2, + ) + if _route_as_rc: + warnings.warn( + "allow_unbalanced_panel=True: routing the unbalanced panel " + "through the repeated-cross-section (RC) levels estimator to " + "match R `did::att_gt(allow_unbalanced_panel=TRUE)`. The RC " + "estimator assumes the population distribution of (Y, X, G) is " + "stable across periods (not data-checkable); standard errors " + "cluster the per-observation influence function by unit.", + UserWarning, + stacklevel=2, + ) + # Complete-case analysis sample: DROP non-finite-outcome rows so that + # survey/PSU resolution, cohort masses, the RC precompute, AND + # estimation all operate on the SAME rows (matches R + # `pre_process_did`). The RC estimator slices the outcome arrays + # directly (no wide-pivot valid-mask), so an unfiltered NaN outcome + # would spuriously materialize its 2x2 cells as non-estimable and + # break the bit-exact ATT parity. Reassign `data` once, up front, so + # every downstream consumer sees the filtered sample. + if _finite_out is not None and not bool(_finite_out.all()): + data = data.loc[_finite_out].reset_index(drop=True) + # Validate panel structure BEFORE routing to the RC estimator — R's + # pre_process_did does the same. The RC precompute reads cohort and + # cluster PER OBSERVATION (no wide pivot), so a duplicate (unit, + # period) row, a unit whose treatment cohort changes over time, or a + # time-varying cluster would silently corrupt the reweighting / + # composition. The normal panel path catches duplicates via the wide + # pivot; the RC route must check explicitly. Fail closed. + # Column-presence guards: `unit`/`time` are guaranteed present (the + # unbalance detection above required them); for a missing `first_treat` + # / `cluster` column, fall through so the canonical required-column + # check below raises "Missing columns: ..." rather than a KeyError. + _dups = int(data.duplicated(subset=[unit, time]).sum()) + if _dups: + raise ValueError( + f"allow_unbalanced_panel=True requires at most one observation " + f"per ({unit}, {time}); found {_dups} duplicate row(s). Aggregate " + f"or de-duplicate before fitting." + ) + if first_treat in data.columns: + # Coerce to numeric first (like the normal fit path) so coercible + # labels such as "3" and "3.0" for the same unit are not falsely + # flagged as a changing cohort. + _ft = pd.to_numeric(data[first_treat], errors="coerce") + _ft_counts = _ft.groupby(data[unit].to_numpy()).nunique() + if (_ft_counts > 1).any(): + _bad = list(_ft_counts.index[_ft_counts > 1][:5]) + raise ValueError( + f"allow_unbalanced_panel=True requires a time-invariant treatment " + f"cohort ('{first_treat}') per unit; unit(s) {_bad} have a changing " + f"first_treat." + ) + if self.cluster is not None and self.cluster in data.columns: + _cl_counts = data.groupby(unit)[self.cluster].nunique() + if (_cl_counts > 1).any(): + _bad = list(_cl_counts.index[_cl_counts > 1][:5]) + raise ValueError( + f"allow_unbalanced_panel=True with cluster='{self.cluster}' " + f"requires a time-invariant cluster per unit; unit(s) {_bad} " + f"have a changing cluster value." + ) + if not self.panel: warnings.warn( "panel=False uses repeated cross-section DRDID estimators " @@ -1984,10 +2128,29 @@ def fit( _resolve_effective_cluster(resolved_survey, cluster_ids, self.cluster) + # allow_unbalanced_panel (RC-routing) with NO user cluster=: cluster the + # per-observation RC influence function by the ORIGINAL unit so the SE + # accounts for within-unit correlation — matching R's unit-level IF + # aggregation for allow_unbalanced_panel. Synthesize SurveyDesign(psu=unit) + # → the same PSU-meat aggregator + PSU multiplier bootstrap used for an + # explicit cluster=. When the user set cluster=X the block above already + # synthesized psu=X (their choice); do NOT double-cluster here. + if _route_as_rc and self.cluster is None and resolved_survey is None: + _unit_psu_design = SurveyDesign(psu=unit, weight_type="pweight") + ( + resolved_survey, + survey_weights, + survey_weight_type, + _, + ) = _resolve_survey_for_fit(_unit_psu_design, data, "analytical") + effective_survey_design = _unit_psu_design + # Validate within-unit constancy for panel survey designs (uses # effective_survey_design so synthesized designs are validated too). + # Skipped under RC-routing: the RC path is observation-level and the + # synthesized psu=unit design is trivially unit-constant. if resolved_survey is not None: - if self.panel: + if self.panel and not _route_as_rc: _validate_unit_constant_survey(data, unit, effective_survey_design) if resolved_survey.weight_type != "pweight": raise ValueError( @@ -2039,7 +2202,10 @@ def fit( treatment_groups = sorted([g for g in df[first_treat].unique() if g > 0]) if self.panel: - # Panel: count unique units + # Panel (incl. an unbalanced panel routed as RC): count unique + # UNITS — the units are real even when the aggregation runs the RC + # estimator. Only a declared repeated cross-section (panel=False) + # counts observations. unit_info = ( df.groupby(unit) .agg({first_treat: "first", "_never_treated": "first"}) @@ -2070,7 +2236,7 @@ def fit( # Per-cell SEs use IF-based variance; aggregated SEs use design-based # variance via compute_survey_if_variance() or PSU-level bootstrap. # Pre-compute data structures for efficient ATT(g,t) computation - if self.panel: + if not _use_rc: precomputed = self._precompute_structures( df, outcome, @@ -2125,7 +2291,7 @@ def fit( _skip_info = {"missing_period": [], "empty_cell": []} _n_skipped_other = 0 - if not self.panel: + if _use_rc: # --- Repeated cross-section path --- # No vectorized/Cholesky fast paths (panel-only optimizations). # Loop using _compute_att_gt_rc() for each (g,t). @@ -2653,6 +2819,11 @@ def fit( cluster_name_for_results: Optional[str] = survey_design.psu elif self.cluster is not None: cluster_name_for_results = self.cluster + elif _route_as_rc: + # RC-on-panel auto-synthesizes SurveyDesign(psu=unit); surface the + # effective unit-level clustering so introspection / summaries don't + # report cluster_name=None while n_clusters is populated. + cluster_name_for_results = unit else: cluster_name_for_results = None n_clusters_for_results: Optional[int] = ( @@ -2706,6 +2877,8 @@ def fit( event_study_vcov=event_study_vcov, event_study_vcov_index=event_study_vcov_index, panel=self.panel, + allow_unbalanced_panel=self.allow_unbalanced_panel, + used_rc_on_unbalanced_panel=_route_as_rc, epv_diagnostics=epv_diagnostics if epv_diagnostics else None, epv_threshold=self.epv_threshold, pscore_fallback=self.pscore_fallback, @@ -3544,11 +3717,42 @@ def _precompute_structures_rc( # For RCS, the resolved survey is already per-observation resolved_survey_rc = resolved_survey - # Fixed cohort masses: total observations per cohort across all periods. - # Used as aggregation weights so that n_treated is consistent with WIF. + # Fixed cohort masses used as aggregation weights (R's pg = n_g / N). + # Count UNIQUE UNITS per cohort, not observations: for a genuine repeated + # cross-section each obs is a distinct unit, so unique-unit-count == + # observation-count (a no-op); but for an unbalanced PANEL routed as RC + # (allow_unbalanced_panel), a cohort's obs-count exceeds its unit-count, + # and R `did::aggte` weights by the fixed UNIT cohort mass — so unique + # units is the R-correct weight on both paths. rcs_cohort_masses = {} + _units_per_cohort = df.groupby(first_treat)[unit].nunique() for g in treatment_groups: - rcs_cohort_masses[g] = int(np.sum(unit_cohorts == g)) + rcs_cohort_masses[g] = int(_units_per_cohort.get(g, 0)) + + # Per-unit aggregation cohort-mass basis for the R pg (= n_g / N over + # UNITS, incl. never-treated). Consumed by _get_agg_cache (fast-path pg) + # and _aggregate_event_study (survey_cohort_weights) so the point + # estimate AND the WIF weight by fixed unit cohort mass. No-op for a true + # RC (each obs is a distinct unit => per-unit == per-obs); the fix for an + # unbalanced panel routed as RC (allow_unbalanced_panel). + _first_pos = df.reset_index(drop=True).drop_duplicates(subset=[unit]).index.values + _unit_cohort_vals = unit_cohorts[_first_pos] + _unit_w = ( + survey_weights_arr[_first_pos] + if survey_weights_arr is not None + else np.ones(len(_first_pos), dtype=np.float64) + ) + agg_cohort_masses = { + float(cv): float(np.sum(_unit_w[_unit_cohort_vals == cv])) + for cv in np.unique(unit_cohorts) + } + agg_total_weight = float(np.sum(_unit_w)) + # Observations per unit, aligned to the per-observation IF index space. + # The WIF is a per-UNIT quantity; on a panel routed as RC each of a + # unit's observations carries its cohort's WIF, so the unit-clustered + # sum over-counts by this factor unless the per-obs WIF is divided by it + # (see _combined_if_fast). 1 for every obs on a true RC (no-op). + obs_per_unit = df.groupby(unit)[unit].transform("size").to_numpy(dtype=np.float64) return { "all_units": all_units, @@ -3574,6 +3778,9 @@ def _precompute_structures_rc( else None ), "rcs_cohort_masses": rcs_cohort_masses, + "agg_cohort_masses": agg_cohort_masses, + "agg_total_weight": agg_total_weight, + "obs_per_unit": obs_per_unit, } def _compute_att_gt_rc( @@ -4738,6 +4945,7 @@ def get_params(self) -> Dict[str, Any]: "cband": self.cband, "pscore_trim": self.pscore_trim, "panel": self.panel, + "allow_unbalanced_panel": self.allow_unbalanced_panel, "epv_threshold": self.epv_threshold, "pscore_fallback": self.pscore_fallback, } diff --git a/diff_diff/staggered_aggregation.py b/diff_diff/staggered_aggregation.py index 4d3ebbb7..bc36bf56 100644 --- a/diff_diff/staggered_aggregation.py +++ b/diff_diff/staggered_aggregation.py @@ -16,6 +16,35 @@ PrecomputedData = Dict[str, Any] +def fixed_cohort_agg_weights( + precomputed: Optional["PrecomputedData"], +) -> Optional[Dict[Any, float]]: + """Fixed per-cohort aggregation masses (R's ``pg = n_g / N`` numerator) for + the treated cohorts ``g > 0``, or ``None`` when the caller should fall back + to per-cell weights (``agg_weight`` / ``n_treated``). + + Priority: unit-level ``agg_cohort_masses`` (RC-on-panel and true RC, exposed + by ``_precompute_structures_rc``) → per-observation survey cohort mass + (survey designs) → ``None`` (panel non-survey). Preferring + ``agg_cohort_masses`` over the raw ``survey_weights`` sum is what makes an + unbalanced panel routed as RC (``allow_unbalanced_panel=True``, which + synthesizes ``SurveyDesign(psu=unit)``) weight every aggregation — simple, + event-study, group, AND the multiplier bootstrap — by fixed UNIT cohort + mass rather than observation count. Single source of truth so the analytical + and bootstrap paths cannot diverge. + """ + if precomputed is None: + return None + agg_masses = precomputed.get("agg_cohort_masses") + if agg_masses is not None: + return {g: m for g, m in agg_masses.items() if g > 0} + sw = precomputed.get("survey_weights") + if sw is not None: + unit_cohorts = precomputed["unit_cohorts"] + return {g: float(np.sum(sw[unit_cohorts == g])) for g in np.unique(unit_cohorts) if g > 0} + return None + + class CallawaySantAnnaAggregationMixin: """ Mixin class providing aggregation methods for CallawaySantAnna estimator. @@ -62,16 +91,10 @@ def _aggregate_simple( gt_pairs = [] groups_for_gt = [] - # For survey: compute fixed per-cohort weight sums from the full - # unit-level sample (matching R's did::aggte pg = n_g / N). - survey_cohort_weights = None - if precomputed is not None and precomputed.get("survey_weights") is not None: - sw = precomputed["survey_weights"] - unit_cohorts = precomputed["unit_cohorts"] - survey_cohort_weights = {} - for g in np.unique(unit_cohorts): - if g > 0: # exclude never-treated (0) - survey_cohort_weights[g] = float(np.sum(sw[unit_cohorts == g])) + # Fixed per-cohort aggregation weights (R's did::aggte pg = n_g / N), + # preferring the unit-level RC mass so allow_unbalanced_panel weights the + # overall ATT by fixed UNIT cohort mass, not observation count. + survey_cohort_weights = fixed_cohort_agg_weights(precomputed) for (g, t), data in group_time_effects.items(): # Only include post-treatment effects (t >= g - anticipation) @@ -237,7 +260,19 @@ def _get_agg_cache(precomputed: "PrecomputedData") -> Dict[str, Any]: return cache cohort_values, cohort_codes = np.unique(unit_cohorts, return_inverse=True) - if survey_w is not None: + agg_masses = precomputed.get("agg_cohort_masses") + if agg_masses is not None: + # RC path: pg basis is per-UNIT cohort mass (R's pg = n_g / N over + # units), exposed by _precompute_structures_rc. `cohort_codes` stays + # per-observation (the WIF scatter is per-obs, divided by + # obs_per_unit downstream). No-op for a true RC (per-unit == + # per-obs); the fix for an unbalanced panel routed as RC. + cohort_masses = np.array( + [float(agg_masses.get(float(cv), 0.0)) for cv in cohort_values], + dtype=np.float64, + ) + total_weight = float(precomputed.get("agg_total_weight", float(np.sum(cohort_masses)))) + elif survey_w is not None: # Survey-weighted cohort masses. np.bincount accumulation order # differs from the historical per-group mask-sums at the ~1 ULP # level (documented drift budget; REGISTRY CallawaySantAnna SE @@ -259,6 +294,9 @@ def _get_agg_cache(precomputed: "PrecomputedData") -> Dict[str, Any]: "cohort_codes": cohort_codes, "cohort_masses": cohort_masses, "total_weight": total_weight, + # Per-obs unit multiplicity for the WIF over-count correction (all + # 1.0 on panel / true RC → the division below is a no-op there). + "obs_per_unit": precomputed.get("obs_per_unit"), } precomputed["_agg_cache"] = cache return cache @@ -379,8 +417,16 @@ def _combined_if_fast( ) return np.full(n_units, np.nan), None - # Scale by 1/total_weight to match R's getSE formula - psi_wif = wif_contrib / total_weight + # Scale by 1/total_weight to match R's getSE formula. On a panel routed + # as RC, additionally divide by obs_per_unit: the WIF is a per-UNIT + # quantity but wif_contrib is per-observation, so the unit-clustered sum + # would otherwise over-count each unit's WIF by its observation count. + # obs_per_unit is 1.0 for panel / true RC (a no-op there). + obs_per_unit = cache.get("obs_per_unit") + if obs_per_unit is not None: + psi_wif = wif_contrib / (obs_per_unit * total_weight) + else: + psi_wif = wif_contrib / total_weight return psi_standard + psi_wif, None @@ -790,15 +836,10 @@ def _aggregate_event_study( # Organize effects by relative time, keeping track of (g,t) pairs effects_by_e: Dict[int, List[Tuple[Tuple[Any, Any], float, float]]] = {} - # Fixed per-cohort survey weights for aggregation - survey_cohort_weights = None - if precomputed is not None and precomputed.get("survey_weights") is not None: - sw = precomputed["survey_weights"] - unit_cohorts = precomputed["unit_cohorts"] - survey_cohort_weights = {} - for g in np.unique(unit_cohorts): - if g > 0: - survey_cohort_weights[g] = float(np.sum(sw[unit_cohorts == g])) + # Fixed per-cohort aggregation weights (shared with _aggregate_simple and + # the bootstrap): unit-level RC mass preferred so allow_unbalanced_panel + # weights each multi-cell horizon by fixed UNIT cohort mass, not obs count. + survey_cohort_weights = fixed_cohort_agg_weights(precomputed) for (g, t), data in group_time_effects.items(): e = t - g # Relative time diff --git a/diff_diff/staggered_bootstrap.py b/diff_diff/staggered_bootstrap.py index 2c3ce962..9dc66d56 100644 --- a/diff_diff/staggered_bootstrap.py +++ b/diff_diff/staggered_bootstrap.py @@ -238,26 +238,22 @@ def _run_multiplier_bootstrap( # analytical _aggregate_simple() path in staggered_aggregation.py. # Do NOT use per-cell survey_weight_sum (which varies by cell on # unbalanced panels). - survey_w = precomputed.get("survey_weights") if precomputed is not None else None - if survey_w is not None: - unit_cohorts = precomputed["unit_cohorts"] - # Precompute fixed cohort masses (same formula as _aggregate_simple) - _cohort_mass_cache: dict = {} - for gt in gt_pairs: - g = gt[0] - if g not in _cohort_mass_cache: - _cohort_mass_cache[g] = float(np.sum(survey_w[unit_cohorts == g])) - all_n_treated = np.array([_cohort_mass_cache[gt[0]] for gt in gt_pairs], dtype=float) - else: - # Use agg_weight if available (RCS: fixed cohort mass); - # fall back to n_treated for panel data - all_n_treated = np.array( - [ - group_time_effects[gt].get("agg_weight", group_time_effects[gt]["n_treated"]) - for gt in gt_pairs - ], - dtype=float, - ) + # Fixed per-cohort aggregation masses — the SAME single source of truth + # the analytical _aggregate_simple() / _aggregate_event_study() use, so + # the bootstrap weights an unbalanced-panel (allow_unbalanced_panel) RC + # aggregation by fixed UNIT cohort mass, not observation count. Returns + # None for panel non-survey (→ per-cell agg_weight/n_treated fallback). + from diff_diff.staggered_aggregation import fixed_cohort_agg_weights + + _fixed_masses = fixed_cohort_agg_weights(precomputed) + + def _agg_mass(gt): + g = gt[0] + if _fixed_masses is not None and g in _fixed_masses: + return _fixed_masses[g] + return group_time_effects[gt].get("agg_weight", group_time_effects[gt]["n_treated"]) + + all_n_treated = np.array([_agg_mass(gt) for gt in gt_pairs], dtype=float) post_n_treated = all_n_treated[post_treatment_mask] # Filter out NaN ATT(g,t) cells from overall aggregation (matches analytical path) @@ -720,17 +716,17 @@ def _prepare_event_study_aggregation( # Use fixed cohort survey masses (not per-cell survey_weight_sum) when # survey weights are present, matching the analytical # _aggregate_event_study() path. - survey_w = precomputed.get("survey_weights") if precomputed is not None else None - _cohort_mass: Optional[dict] = None - if survey_w is not None: - unit_cohorts = precomputed["unit_cohorts"] - _cohort_mass = {} + # Shared fixed-cohort masses (same source of truth as the analytical + # event-study path): unit-level RC mass preferred so the bootstrap + # event-study weights an unbalanced-panel (allow_unbalanced_panel) + # aggregation by fixed UNIT cohort mass, not observation count. + from diff_diff.staggered_aggregation import fixed_cohort_agg_weights + + _fixed_masses = fixed_cohort_agg_weights(precomputed) def _agg_weight(g: Any, t: Any) -> float: - if _cohort_mass is not None: - if g not in _cohort_mass: - _cohort_mass[g] = float(np.sum(survey_w[unit_cohorts == g])) - return _cohort_mass[g] + if _fixed_masses is not None and g in _fixed_masses: + return _fixed_masses[g] # Use agg_weight if available (RCS: fixed cohort mass) return group_time_effects[(g, t)].get( "agg_weight", group_time_effects[(g, t)]["n_treated"] diff --git a/diff_diff/staggered_results.py b/diff_diff/staggered_results.py index 4bb45f37..fde82418 100644 --- a/diff_diff/staggered_results.py +++ b/diff_diff/staggered_results.py @@ -155,6 +155,15 @@ class CallawaySantAnnaResults: # contract. anticipation: int = 0 panel: bool = True + # allow_unbalanced_panel routing (RC-on-panel). `allow_unbalanced_panel` + # records the fit-time flag; `used_rc_on_unbalanced_panel` is True ONLY when + # the panel was actually unbalanced and the RC (repeated-cross-section) + # levels estimator was used (the flag is inert on a balanced panel). When + # True the ATT(g,t) estimand is R's `allow_unbalanced_panel=TRUE` RC-on-panel + # estimand, not within-cell panel differencing — and the influence function + # is clustered by unit (see `cluster_name`). + allow_unbalanced_panel: bool = False + used_rc_on_unbalanced_panel: bool = False event_study_effects: Optional[Dict[int, Dict[str, Any]]] = field(default=None) group_effects: Optional[Dict[Any, Dict[str, Any]]] = field(default=None) influence_functions: Optional["np.ndarray"] = field(default=None, repr=False) diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 393cbe01..054e9b7e 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -562,19 +562,37 @@ Aggregations: adjustment, matching R's `did::aggte()`. The WIF accounts for uncertainty in estimating group-size aggregation weights. Group aggregation uses equal time weights (deterministic), so WIF is zero. - - **Deviation from R (unbalanced-panel event-study weighting):** within an event-study - (dynamic) horizon that pools MORE THAN ONE (g,t) cell, diff-diff weights each cell by its - per-cell aggregation weight (`agg_weight` / valid `n_treated`), whereas R `did::aggte` - weights by the fixed cohort probability `pg = n_g / N` (or survey mass) computed once from - the group column. On **balanced** panels these coincide exactly (a cell's valid count equals - the cohort mass), so event-study / group / simple aggregates match R to ~1e-5 -- verified, - including the universal zero-reference dilution case. On **unbalanced** panels a cell can - carry fewer valid units than the cohort mass, so a multi-cell horizon's relative weights (and - thus that horizon's ATT/SE) can differ from R -- e.g. on a dropped-unit panel a two-real-cell - horizon was diff-diff `-0.065` vs R `-0.136`. This is a **pre-existing, general** aggregation - convention (it predates and is independent of the universal reference cells -- it applies to - any multi-cell horizon, references or not) and is tracked in `TODO.md`. Single-cell horizons - (the majority) normalize to weight 1 and are unaffected. + - **Unbalanced panels — default within-cell differencing vs `allow_unbalanced_panel=True`:** + on an unbalanced panel (some units unobserved in some periods) the **default** path estimates + each ATT(g,t) by within-cell panel differencing on the units observed at BOTH the base period + and t, and weights a multi-cell event-study horizon by per-cell valid `n_treated`. This is a + valid but **different estimand** than R `did::att_gt(allow_unbalanced_panel=TRUE)`, which sets + `panel=FALSE` and runs the repeated-cross-section levels estimator (`DRDID::reg_did_rc`) on the + pooled observations, weighting by the fixed cohort probability `pg = n_g / N` over UNITS. Both + the cell estimator AND the weighting differ from R on unbalanced data (the estimator choice + dominates); on **balanced** panels they coincide exactly (each cell's valid count equals the + cohort mass), so the default path matches R. The default path emits a `UserWarning` on + unbalanced input (no-silent-failures) pointing to the flag. + - **`allow_unbalanced_panel=True` (RC-on-panel parity with R):** routes an unbalanced panel's + pooled observations through diff-diff's RC estimator (bit-exact vs `reg_did_rc`) and clusters + the per-observation influence function by the original unit. ATT matches R **bit-for-bit** — + cells AND dynamic aggregation, including the fixed unit-cohort-mass `pg` reweighting and the + per-unit WIF (the per-observation WIF is divided by each unit's observation count so the + unit-clustered sum is not over-counted). Inert on balanced panels (byte-identical to the + default) — matching R, whose `pre_process_did` likewise recomputes balance and keeps + `panel=TRUE` (differencing) on balanced input, so the flag engages RC only when the panel is + genuinely unbalanced. Panel structure is validated before routing (no duplicate `(unit, + period)` rows, time-invariant treatment cohort and cluster per unit), matching R's + preprocessing — fail-closed since the RC precompute reads cohort/cluster per observation. + `survey_design=` with the flag raises `NotImplementedError` (per-obs vs per-unit weight + resolution deferred). Verified vs R `did` 2.5.1 in + `tests/test_csdid_ported.py::TestAllowUnbalancedPanel` (golden + `benchmarks/data/cs_unbalanced_golden.json`). + - **Deviation from R:** the analytical SE equals R's up to the CR1 finite-sample factor + `sqrt(G/(G-1))` (G = number of units): diff-diff's cluster-robust variance applies the + `G/(G-1)` Bessel correction that R's `att_gt` `getSE` (`sqrt(mean(inf^2)/n)`) omits. Exact + factor, ~0.25% at G=200, vanishing as G → ∞ — the same convention class as the fixest + cluster-SE band (G2). - Bootstrap: Multiplier bootstrap with Rademacher, Mammen, or Webb weights. Bootstrap perturbs the combined influence function (standard IF + WIF) directly, not just fixed-weight re-aggregation. This correctly propagates weight estimation uncertainty. diff --git a/tests/test_csdid_ported.py b/tests/test_csdid_ported.py index 4bbe9e34..b829e1a4 100644 --- a/tests/test_csdid_ported.py +++ b/tests/test_csdid_ported.py @@ -1915,3 +1915,440 @@ def test_universal_reference_materialized_for_all_skipped_cohort(self): and np.isnan(ref["se"]) and ref["n_treated"] == 40 # cohort 7's fixed mass (40 units) ) + + +# --------------------------------------------------------------------------- +# allow_unbalanced_panel: RC-on-panel parity with R +# did::att_gt(allow_unbalanced_panel=TRUE) (panel=FALSE -> DRDID::reg_did_rc) +# --------------------------------------------------------------------------- + +_UNBAL_GOLDEN_PATH = Path(__file__).parents[1] / "benchmarks" / "data" / "cs_unbalanced_golden.json" + + +@pytest.fixture(scope="module") +def unbalanced_golden(): + """Load the unbalanced-panel R golden. Skip if absent.""" + if not _UNBAL_GOLDEN_PATH.exists(): + pytest.skip( + f"Golden {_UNBAL_GOLDEN_PATH} missing — regenerate via " + "`Rscript benchmarks/R/generate_cs_unbalanced_golden.R`." + ) + with open(_UNBAL_GOLDEN_PATH) as f: + return json.load(f) + + +def _unbal_df(golden): + d = golden["data"] + return pd.DataFrame( + { + "unit": d["unit"], + "period": d["period"], + "first_treat": d["first_treat"], + "outcome": d["outcome"], + } + ) + + +class TestAllowUnbalancedPanel: + """``CallawaySantAnna(allow_unbalanced_panel=True)`` matches R + ``did::att_gt(allow_unbalanced_panel=TRUE)`` (which sets ``panel=FALSE`` and + runs ``DRDID::reg_did_rc`` on the pooled observations). + + ATT is bit-exact (cells AND dynamic aggregation). The analytical SE equals + R's up to the CR1 ``G/(G-1)`` finite-sample factor diff-diff's cluster path + applies (R's ``att_gt`` getSE omits it) — a documented Deviation from R. The + per-observation RC influence function is clustered by the original unit so + the SE accounts for within-unit correlation. + """ + + def _fit(self, golden): + return CallawaySantAnna( + estimation_method="reg", + control_group="never_treated", + allow_unbalanced_panel=True, + ).fit( + _unbal_df(golden), + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + aggregate="event_study", + ) + + def test_cells_match_r(self, unbalanced_golden): + """Per-cell ATT bit-exact; per-cell SE == R * sqrt(G/(G-1)).""" + with pytest.warns(UserWarning, match="repeated-cross-section|allow_unbalanced"): + res = self._fit(unbalanced_golden) + c = unbalanced_golden["cells"] + n_units = unbalanced_golden["meta"]["n_units"] + fac = np.sqrt(n_units / (n_units - 1)) + for i in range(len(c["g"])): + key = (int(c["g"][i]), int(c["t"][i])) + v = res.group_time_effects.get(key) + assert v is not None and not v.get("is_reference"), f"missing cell {key}" + assert v["effect"] == pytest.approx(c["att"][i], abs=1e-9) + assert v["se"] == pytest.approx(c["se"][i] * fac, abs=1e-8) + + def test_event_study_matches_r(self, unbalanced_golden): + """Dynamic-aggregation ATT bit-exact; SE == R * sqrt(G/(G-1)). + + Exercises the unit-level pg reweighting + the unit-level WIF + (obs_per_unit) fix: on an unbalanced panel a multi-cell horizon would + otherwise weight cells by observation count, not fixed unit cohort mass. + """ + with pytest.warns(UserWarning): + res = self._fit(unbalanced_golden) + es = unbalanced_golden["event_study"] + n_units = unbalanced_golden["meta"]["n_units"] + fac = np.sqrt(n_units / (n_units - 1)) + for i, e in enumerate(es["egt"]): + v = res.event_study_effects.get(int(e)) + assert v is not None, f"missing event time {e}" + assert v["effect"] == pytest.approx(es["att"][i], abs=1e-9) + assert v["se"] == pytest.approx(es["se"][i] * fac, abs=1e-8) + + def test_flag_inert_on_balanced_panel(self): + """On a balanced panel the flag is a no-op: output byte-identical.""" + rng = np.random.default_rng(3) + rows = [] + for u in range(90): + g = [0, 3, 4][u % 3] + ufe = rng.normal() + for t in range(1, 6): + post = 1.0 if (g and t >= g) else 0.0 + rows.append( + { + "unit": u, + "period": t, + "first_treat": g, + "outcome": ufe + 0.3 * t + (1.5 * post if g else 0.0) + rng.normal(0, 0.5), + } + ) + bal = pd.DataFrame(rows) + kw = dict( + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + aggregate="event_study", + ) + r0 = CallawaySantAnna(estimation_method="reg").fit(bal, **kw) + r1 = CallawaySantAnna(estimation_method="reg", allow_unbalanced_panel=True).fit(bal, **kw) + for e in r0.event_study_effects: + assert r1.event_study_effects[e]["effect"] == pytest.approx( + r0.event_study_effects[e]["effect"], abs=1e-13 + ) + assert r1.event_study_effects[e]["se"] == pytest.approx( + r0.event_study_effects[e]["se"], abs=1e-13 + ) + + def test_unbalanced_default_path_warns(self, unbalanced_golden): + """Without the flag, an unbalanced panel warns (no-silent-failures).""" + with pytest.warns(UserWarning, match="Unbalanced panel detected"): + CallawaySantAnna(estimation_method="reg").fit( + _unbal_df(unbalanced_golden), + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + aggregate="event_study", + ) + + def test_user_cluster_honored(self, unbalanced_golden): + """A user-supplied cluster= is honored (not double-clustered by unit); + the ATT is unaffected (still bit-exact vs R).""" + df = _unbal_df(unbalanced_golden) + df["clus"] = (df["unit"] % 12).astype(int) + with pytest.warns(UserWarning): + res = CallawaySantAnna( + estimation_method="reg", + control_group="never_treated", + allow_unbalanced_panel=True, + cluster="clus", + ).fit( + df, + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + aggregate="event_study", + ) + es = unbalanced_golden["event_study"] + for i, e in enumerate(es["egt"]): + v = res.event_study_effects.get(int(e)) + if v is not None: + assert v["effect"] == pytest.approx(es["att"][i], abs=1e-9) + + def test_survey_design_with_flag_raises(self, unbalanced_golden): + """survey_design= x allow_unbalanced_panel= is fail-closed (deferred).""" + from diff_diff.survey import SurveyDesign + + df = _unbal_df(unbalanced_golden) + df["w"] = 1.0 + with pytest.raises(NotImplementedError, match="allow_unbalanced_panel"): + CallawaySantAnna( + estimation_method="reg", + control_group="never_treated", + allow_unbalanced_panel=True, + ).fit( + df, + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + survey_design=SurveyDesign(weights="w"), + ) + + def test_simple_overall_matches_r(self, unbalanced_golden): + """Simple `overall_att` matches R `aggte(type="simple")` — the aggregate + weighted by fixed UNIT cohort mass (not observation count). This is the + cross-surface twin of the event-study weighting: both go through the + shared `fixed_cohort_agg_weights`.""" + with pytest.warns(UserWarning): + res = CallawaySantAnna( + estimation_method="reg", + control_group="never_treated", + allow_unbalanced_panel=True, + ).fit( + _unbal_df(unbalanced_golden), + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + ) + s = unbalanced_golden["simple"] + n_units = unbalanced_golden["meta"]["n_units"] + fac = np.sqrt(n_units / (n_units - 1)) + assert res.overall_att == pytest.approx(s["overall_att"], abs=1e-9) + assert res.overall_se == pytest.approx(s["overall_se"] * fac, abs=1e-8) + + def test_group_effects_match_r(self, unbalanced_golden): + """Group (cohort) effects match R `aggte(type="group")` att.egt / se.egt.""" + with pytest.warns(UserWarning): + res = CallawaySantAnna( + estimation_method="reg", + control_group="never_treated", + allow_unbalanced_panel=True, + ).fit( + _unbal_df(unbalanced_golden), + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + aggregate="group", + ) + g = unbalanced_golden["group"] + n_units = unbalanced_golden["meta"]["n_units"] + fac = np.sqrt(n_units / (n_units - 1)) + assert res.group_effects is not None + for i, gval in enumerate(g["egt"]): + eff = res.group_effects.get(gval) or res.group_effects.get(float(gval)) + assert eff is not None, f"missing group {gval}" + assert eff["effect"] == pytest.approx(g["att"][i], abs=1e-9) + assert eff["se"] == pytest.approx(g["se"][i] * fac, abs=1e-8) + + def test_bootstrap_event_study_uses_unit_weights(self, unbalanced_golden): + """The multiplier bootstrap weights the event-study aggregation by fixed + UNIT cohort mass (via the same shared helper), so the bootstrap event- + study effects equal the (unit-weighted) analytical effects exactly — not + the observation-weighted values the pre-fix bootstrap path produced.""" + kw = dict( + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + aggregate="event_study", + ) + df = _unbal_df(unbalanced_golden) + with pytest.warns(UserWarning): + a = CallawaySantAnna( + estimation_method="reg", control_group="never_treated", allow_unbalanced_panel=True + ).fit(df, **kw) + with pytest.warns(UserWarning): + b = CallawaySantAnna( + estimation_method="reg", + control_group="never_treated", + allow_unbalanced_panel=True, + n_bootstrap=200, + seed=11, + ).fit(df, **kw) + for e in a.event_study_effects: + assert b.event_study_effects[e]["effect"] == pytest.approx( + a.event_study_effects[e]["effect"], abs=1e-12 + ) + + def test_result_metadata_reflects_rc_routing(self, unbalanced_golden): + """The result records the flag AND whether RC-on-panel actually engaged, + and labels the effective unit-level clustering — so downstream can tell + the RC-on-panel estimand from the default within-cell one.""" + df = _unbal_df(unbalanced_golden) + kw = dict(outcome="outcome", unit="unit", time="period", first_treat="first_treat") + # RC-on-panel actually used (unbalanced + flag) + with pytest.warns(UserWarning): + rc = CallawaySantAnna( + estimation_method="reg", control_group="never_treated", allow_unbalanced_panel=True + ).fit(df, **kw) + assert rc.allow_unbalanced_panel is True + assert rc.used_rc_on_unbalanced_panel is True + assert rc.cluster_name == "unit" # effective unit clustering surfaced + assert rc.n_clusters == unbalanced_golden["meta"]["n_units"] + # Default unbalanced (no flag): not routed, flag False + with pytest.warns(UserWarning, match="Unbalanced panel detected"): + d = CallawaySantAnna(estimation_method="reg").fit(df, **kw) + assert d.allow_unbalanced_panel is False + assert d.used_rc_on_unbalanced_panel is False + # Balanced + flag: flag set but inert (not routed) + rng = np.random.default_rng(4) + rows = [] + for u in range(60): + g = [0, 3, 4][u % 3] + ufe = rng.normal() + for t in range(1, 6): + post = 1.0 if (g and t >= g) else 0.0 + rows.append( + { + "unit": u, + "period": t, + "first_treat": g, + "outcome": ufe + 0.3 * t + (1.5 * post if g else 0.0) + rng.normal(0, 0.5), + } + ) + b = CallawaySantAnna(estimation_method="reg", allow_unbalanced_panel=True).fit( + pd.DataFrame(rows), **kw + ) + assert b.allow_unbalanced_panel is True + assert b.used_rc_on_unbalanced_panel is False # inert on balanced + + def test_rejects_duplicate_unit_period(self, unbalanced_golden): + """allow_unbalanced_panel=True fails closed on duplicate (unit, period) + rows — the RC route does not pivot, so it must validate explicitly.""" + df = _unbal_df(unbalanced_golden) + dup = pd.concat([df, df.iloc[[0]]], ignore_index=True) # duplicate one row + with pytest.raises(ValueError, match="at most one observation per"): + CallawaySantAnna( + estimation_method="reg", control_group="never_treated", allow_unbalanced_panel=True + ).fit(dup, outcome="outcome", unit="unit", time="period", first_treat="first_treat") + + def test_rejects_changing_first_treat(self, unbalanced_golden): + """allow_unbalanced_panel=True fails closed when a unit's treatment + cohort changes over time (would contaminate cohort composition).""" + df = _unbal_df(unbalanced_golden).copy() + # Flip one treated unit's first_treat in one period. + u = df[df["first_treat"] == 3]["unit"].iloc[0] + idx = df[(df["unit"] == u)].index[0] + df.loc[idx, "first_treat"] = 4 + with pytest.raises(ValueError, match="time-invariant treatment cohort"): + CallawaySantAnna( + estimation_method="reg", control_group="never_treated", allow_unbalanced_panel=True + ).fit(df, outcome="outcome", unit="unit", time="period", first_treat="first_treat") + + def test_rejects_time_varying_cluster(self, unbalanced_golden): + """allow_unbalanced_panel=True with a time-varying cluster= fails closed + (the RC route skips the panel unit-constant survey validator).""" + df = _unbal_df(unbalanced_golden).copy() + df["clus"] = 0 + u = df["unit"].iloc[0] + df.loc[df[df["unit"] == u].index[0], "clus"] = 1 # one unit's cluster changes + with pytest.raises(ValueError, match="time-invariant cluster"): + CallawaySantAnna( + estimation_method="reg", + control_group="never_treated", + allow_unbalanced_panel=True, + cluster="clus", + ).fit(df, outcome="outcome", unit="unit", time="period", first_treat="first_treat") + + def test_nan_outcome_triggers_rc_routing(self): + """A RECTANGULAR panel with a non-finite OUTCOME cell is effectively + unbalanced (complete-case, like R): the flag must engage RC routing, not + stay inert. Guards the structural-vs-complete-case detection gap.""" + rng = np.random.default_rng(6) + rows = [] + for u in range(90): + g = [0, 3, 4][u % 3] + ufe = rng.normal() + for t in range(1, 6): + post = 1.0 if (g and t >= g) else 0.0 + rows.append( + { + "unit": u, + "period": t, + "first_treat": g, + "outcome": ufe + 0.3 * t + (1.5 * post if g else 0.0) + rng.normal(0, 0.5), + } + ) + full = pd.DataFrame(rows) # fully rectangular / balanced + df_nan = full.copy() + df_nan.loc[full.index[10], "outcome"] = np.nan # one missing-outcome cell + df_drop = full.drop(index=full.index[10]).reset_index(drop=True) # same, row removed + kw = dict( + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + aggregate="event_study", + ) + + def mk(): + return CallawaySantAnna( + estimation_method="reg", control_group="never_treated", allow_unbalanced_panel=True + ) + + with pytest.warns(UserWarning, match="repeated-cross-section|allow_unbalanced"): + res = mk().fit(df_nan, **kw) + assert res.used_rc_on_unbalanced_panel is True + assert res.cluster_name == "unit" + # The RC path must ESTIMATE on the complete-case sample (drop the NaN row), + # not materialize the NaN cell as non-estimable: the NaN-outcome fit must + # equal the structurally-row-removed fit bit-for-bit (att AND se). + ref = mk().fit(df_drop, **kw) + assert res.n_obs == ref.n_obs # NaN row dropped + for k, v in res.group_time_effects.items(): + if v.get("is_reference"): + continue + assert v["effect"] == pytest.approx(ref.group_time_effects[k]["effect"], abs=1e-12) + for e, v in res.event_study_effects.items(): + assert v["effect"] == pytest.approx(ref.event_study_effects[e]["effect"], abs=1e-12) + assert res.overall_att == pytest.approx(ref.overall_att, abs=1e-12) + + def test_complete_case_rectangular_stays_inert(self): + """When a WHOLE unit and a WHOLE period are dropped by non-finite + outcomes but the remaining complete-case sample is still rectangular, the + flag must stay inert (balance is assessed on the complete-case frame for + BOTH the cell count and the units x periods denominator, matching R).""" + rng = np.random.default_rng(3) + rows = [] + for u in range(90): + g = [0, 3, 4][u % 3] + ufe = rng.normal() + for t in range(1, 6): + post = 1.0 if (g and t >= g) else 0.0 + rows.append( + { + "unit": u, + "period": t, + "first_treat": g, + "outcome": ufe + 0.3 * t + (1.5 * post if g else 0.0) + rng.normal(0, 0.5), + } + ) + df = pd.DataFrame(rows) + df.loc[df["unit"] == 0, "outcome"] = np.nan # a whole (never-treated) unit + df.loc[df["period"] == 5, "outcome"] = np.nan # a whole period + # complete-case sample = 89 units x 4 periods, still rectangular -> inert. + res = CallawaySantAnna(estimation_method="reg", allow_unbalanced_panel=True).fit( + df, outcome="outcome", unit="unit", time="period", first_treat="first_treat" + ) + assert res.used_rc_on_unbalanced_panel is False + + def test_coercible_first_treat_not_rejected(self, unbalanced_golden): + """A unit with coercible-equal first_treat labels ('3' vs '3.0') is NOT + rejected as a changing cohort (validation coerces before comparing).""" + df = _unbal_df(unbalanced_golden).copy() + df["first_treat"] = df["first_treat"].astype(object) + # Represent cohort-3 rows as a mix of "3" and 3.0 for the same units. + mask3 = df["first_treat"].isin([3, 3.0, "3", "3.0"]) + df.loc[mask3, "first_treat"] = ["3" if i % 2 else 3.0 for i in range(int(mask3.sum()))] + with pytest.warns(UserWarning): + res = CallawaySantAnna( + estimation_method="reg", control_group="never_treated", allow_unbalanced_panel=True + ).fit(df, outcome="outcome", unit="unit", time="period", first_treat="first_treat") + assert res.used_rc_on_unbalanced_panel is True # fit succeeds, not rejected