feat(callaway): allow_unbalanced_panel — RC-on-panel parity with R did

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) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0178r8ZK6VRiXbaBiknrWjQB

Author: igerber

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

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 |

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")