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

[igerber]

Summary

• Adds CallawaySantAnna(allow_unbalanced_panel=True) — parity with R did::att_gt(allow_unbalanced_panel=TRUE) on unbalanced panels. When the flag is set and the input panel is unbalanced (some units unobserved in some periods), 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 — per-cell AND every aggregation (simple / event-study / group): the aggregation weights use fixed UNIT cohort mass pg = n_g/N (not observation count), via a single shared fixed_cohort_agg_weights() helper used by the analytical simple/event-study paths and the multiplier bootstrap; 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 — same convention class as the fixest cluster band).
• Inert on balanced panels (byte-identical to the default) — this matches R, whose pre_process_did recomputes balance and keeps panel=TRUE on balanced input under the flag (verified). Independently, the default path now emits a UserWarning on unbalanced input (previously silent — a no-silent-failures fix) pointing to the flag.
• Complete-case, like R: unbalancedness is detected on finite-outcome rows and the RC analysis sample drops non-finite-outcome rows, so a rectangular panel with a NaN outcome routes to RC and estimates on the row-dropped sample (bit-identical to structural removal), not materializing the cell as non-estimable.
• Panel-structure validation before routing (matching R's preprocessing): no duplicate (unit, period) rows, time-invariant treatment cohort and cluster per unit — fail-closed, since the RC precompute reads cohort/cluster per observation.
• Result metadata records allow_unbalanced_panel + used_rc_on_unbalanced_panel and labels the effective cluster_name="unit", so downstream can distinguish the RC-on-panel estimand from the default.
• survey_design= combined with the flag raises NotImplementedError (per-obs vs per-unit weight resolution deferred).

Methodology references (required if estimator / math changes)

• Method name(s): CallawaySantAnna unbalanced-panel handling; R did::att_gt(allow_unbalanced_panel=TRUE) → panel=FALSE → DRDID::reg_did_rc (repeated-cross-section levels estimator); dynamic/simple/group aggregation pg = n_g/N weighting + weight-influence-function correction.
• Paper / source link(s): Callaway & Sant'Anna (2021, JoE); Sant'Anna & Zhao (2020) RC estimator; R did 2.5.1. See docs/methodology/REGISTRY.md (CallawaySantAnna section).
• Any intentional deviations from the source (and why): (1) the analytical SE carries the CR1 sqrt(G/(G-1)) finite-sample factor that R's att_gt getSE omits (documented Deviation from R; vanishes as G→∞) — the ATT and the routing/complete-case/validation behavior otherwise match R; (2) survey_design= × flag and covariate/ipw/dr × flag R-parity are deferred (fail-closed / TODO).

Validation

• Tests added/updated: tests/test_csdid_ported.py::TestAllowUnbalancedPanel (15 tests: per-cell + simple/event-study/group parity vs R, bootstrap-by-unit, user-cluster honored, survey_design= fail-closed, balanced-inert byte-identity, default-path warning, result metadata, panel-structure validation [duplicate rows / changing cohort / time-varying cluster], NaN-outcome complete-case routing [== structural removal], coercible first_treat). New R golden benchmarks/data/cs_unbalanced_golden.json (generator benchmarks/R/generate_cs_unbalanced_golden.R), verified vs R did 2.5.1. 587 CS tests pass (RC / survey / panel / triple-diff paths unchanged).
• Backtest / simulation / notebook evidence (if applicable): N/A.

Security / privacy

• Confirm no secrets/PII in this PR: Yes

---

🤖 Generated with Claude Code

https://claude.ai/code/session_0178r8ZK6VRiXbaBiknrWjQB

[github-actions]

Overall Assessment

⚠️ Needs changes

Executive Summary

• The RC-on-unbalanced-panel methodology is broadly aligned with R did: unbalanced panel input is routed through the repeated-cross-section estimator while keeping unit-level IF aggregation. R’s preprocessing also validates panel structure before that route. (raw.githubusercontent.com)
• One P1 blocker: the new route bypasses panel-structure validation, so invalid panel data can be accepted silently under allow_unbalanced_panel=True.
• Documented deviations in REGISTRY.md are correctly treated as non-blocking: default within-cell differencing and the CR1 SE factor are documented.
• I could not run tests locally because the environment is missing pandas; this review is static plus source cross-check.

Methodology

P1 — Missing panel-structure validation before RC routing
Impact: R validates (id, period) uniqueness, time-invariant treatment cohort, and time-invariant cluster variables before flipping an unbalanced panel into the RC path. This PR detects unbalancedness before normal validation and then routes to _precompute_structures_rc, where cohort membership is row-level. An unbalanced panel with duplicated (unit, time) rows or a unit whose first_treat changes over time can therefore be treated as valid RC-on-panel data, silently reweighting duplicated observations or contaminating cohort/control composition. The same skip also lets time-varying cluster= through because _validate_unit_constant_survey is skipped under _route_as_rc. See diff_diff/staggered.py:L1869-L1881, diff_diff/staggered.py:L2075-L2081, diff_diff/staggered.py:L2166-L2189, and diff_diff/staggered.py:L3613-L3628.
Concrete fix: after required-column checks and numeric coercion, but before RC routing, add self.panel validations for duplicate [unit, time], within-unit first_treat constancy, and within-unit cluster constancy when cluster is set. Add regression tests where allow_unbalanced_panel=True raises for each invalid shape.

P3 — Documented deviations, no action required
Impact: The Registry explicitly documents that the default unbalanced-panel path remains within-cell panel differencing and warns, and that RC-on-panel analytical SE differs from R by the CR1 sqrt(G/(G-1)) factor. These are documented deviations, not defects. See docs/methodology/REGISTRY.md:L565-L590.
Concrete fix: None required.

Code Quality

P3 — Stale/unclear inline comment
Impact: diff_diff/staggered.py:L1875-L1877 says this behavior is narrower than R because R “forces RC even on balanced data.” Current R preprocessing recomputes allow_unbalanced_panel from actual balance before flipping panel <- FALSE, so this comment is at least misleading against the cited source. (raw.githubusercontent.com)
Concrete fix: Replace the comment with the actual contract: route only when the observed panel is unbalanced; balanced-panel inertness is intentional and tested.

Performance

No findings.

Maintainability

P3 — Existing event-study/bootstrap bucket duplication remains tracked
Impact: The shared fixed_cohort_agg_weights() helper reduces one important divergence risk, but event-study bucket construction duplication remains tracked in TODO.md. This is documented backlog, not a blocker. See TODO.md:L132-L133.
Concrete fix: None for this PR.

Tech Debt

P3 — Deferred covariate/IPW/DR parity is tracked
Impact: TODO.md tracks covariate / IPW / DR R-parity verification for allow_unbalanced_panel=True. Because it is tracked and the no-covariate reg golden covers the landed path, this is informational. See TODO.md:L132.
Concrete fix: None required for approval.

Security

No findings.

Documentation/Tests

P1 — Add tests for invalid panel shapes on the new route
Impact: Current added tests cover R parity, simple/event-study/group rows, bootstrap weighting, metadata, warning, user cluster, and survey fail-closed behavior, but not invalid panel structure that the new route can now bypass. See tests/test_csdid_ported.py:L1978-L2220.
Concrete fix: Add tests asserting ValueError under allow_unbalanced_panel=True for: duplicated (unit, period) rows on an otherwise unbalanced panel, a unit with changing first_treat, and a time-varying user cluster= column.

Path To Approval

1. Add panel-shape validation before _route_as_rc is used: duplicate (unit, time) rows, within-unit first_treat constancy, and time-invariant cluster= for panel input.
2. Add the three regression tests listed above under TestAllowUnbalancedPanel.
3. Update the misleading balanced-panel/R comment at diff_diff/staggered.py:L1875-L1877.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: f8c25c90c5fceb907786c2e896844ac05cf25d14

---

Overall Assessment

✅ Looks good

Executive Summary

• Previous P1 findings are resolved: RC routing now validates duplicate (unit, period), time-varying first_treat, and time-varying cluster= before using the observation-level RC precompute.
• Methodology is documented in REGISTRY.md: default unbalanced-panel differencing, RC-on-panel routing, fixed unit-cohort pg, WIF over-count correction, and the CR1 SE deviation.
• Analytical and bootstrap aggregation now share fixed_cohort_agg_weights(), reducing the prior weighting divergence risk.
• One P2 code-quality issue remains: the new routing validation runs before required-column checks and numeric normalization, so some invalid or coercible inputs can raise poor/over-strict errors.
• I could not run the added tests locally because pytest is not installed.

Methodology

P3 — Documented deviations, no action required
Impact: The affected method is CallawaySantAnna, specifically unbalanced-panel handling and simple/event-study/group aggregation. The registry now documents the default within-cell differencing behavior and the allow_unbalanced_panel=True RC-on-panel route, including fixed unit-cohort-mass pg, per-unit WIF correction, and the analytical SE CR1 factor. These are documented deviations or documented parity behavior, not defects. See docs/methodology/REGISTRY.md:L565-L595.
Concrete fix: None required.

P3 — Prior panel-validation methodology issue appears resolved
Impact: R preprocessing validates panel structure before unbalanced-panel RC routing and recomputes the unbalanced flag before switching to panel=FALSE; the new code mirrors that by checking duplicates, time-invariant treatment cohort, and time-invariant user cluster before _precompute_structures_rc. See diff_diff/staggered.py:L1869-L1947, plus regression tests at tests/test_csdid_ported.py:L2222-L2258. citeturn1view0
Concrete fix: None required.

Code Quality

P2 — New routing validation precedes required-column checks and normalization
Impact: With allow_unbalanced_panel=True on an unbalanced input, fit() can access first_treat before the normal missing-column check, producing a raw KeyError instead of the documented ValueError. It also checks time-invariance before pd.to_numeric() / np.inf -> 0 normalization, so mixed but coercible cohort values may be rejected unnecessarily. See diff_diff/staggered.py:L1869-L1938 versus the normal required-column and coercion block at diff_diff/staggered.py:L2126-L2149.
Concrete fix: Move required-column validation, df = data.copy(), numeric conversion, and np.inf normalization before the _route_as_rc validation block, then run duplicate/cohort/cluster validation on the normalized df.

Performance

No findings.

Maintainability

P3 — Remaining aggregation bucket duplication is tracked
Impact: Event-study bucket construction remains duplicated between analytical and bootstrap paths, but this is already tracked as deferred refactor work in TODO.md. See TODO.md:L133.
Concrete fix: None required for this PR.

Tech Debt

P3 — Deferred survey and covariate/IPW/DR parity work is tracked
Impact: survey_design= with allow_unbalanced_panel=True fails closed, and covariate/IPW/DR parity verification is explicitly tracked in TODO.md; these are deferrable and not silent correctness issues in the landed no-covariate reg golden path. See TODO.md:L132 and diff_diff/staggered.py:L1884-L1893.
Concrete fix: None required for approval.

Security

No findings.

Documentation/Tests

No blocking findings. The added tests cover R golden parity, fixed unit-cohort weighting across simple/event-study/group, bootstrap weighting, metadata, default warning, survey fail-closed behavior, and the three prior invalid-panel-shape cases. See tests/test_csdid_ported.py:L1978-L2258.

Verification note: attempted PYTHONDONTWRITEBYTECODE=1 pytest -q tests/test_csdid_ported.py::TestAllowUnbalancedPanel -p no:cacheprovider, but the environment returned pytest: command not found.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 0ece40735393046810367906f50f4d5447b41bed

---

Overall Assessment

⚠️ Needs changes

Executive Summary

• Prior P1 issues around duplicate (unit, period), changing first_treat, and time-varying cluster= are addressed.
• The RC-on-panel methodology and CR1 SE deviation are documented in REGISTRY.md.
• Analytical and bootstrap aggregation now share fixed cohort-mass weighting.
• P1 remains: allow_unbalanced_panel=True only detects structurally absent rows, not missing/non-finite outcome rows that the estimator later treats as missing panel cells.
• I could not run tests locally: pytest and pandas are not installed in this environment.

Methodology

P1 — RC routing misses complete-case unbalanced panels
Impact: R did complete-finite-filters rows before deciding whether allow_unbalanced_panel should flip panel data to the repeated-cross-section path; then it tests nrow(data) != n_units * n_periods and sets .rowid to the original unit id for unbalanced-panel RC inference. (raw.githubusercontent.com) This implementation computes _is_unbalanced_panel from only distinct (unit, time) row presence before validation/normalization, so a rectangular panel with outcome=np.nan in one cell is treated as balanced and the flag stays inert. The later panel path explicitly treats NaN outcomes as missing cells via the pivot/valid-mask logic, so the fit silently uses within-cell panel differencing instead of the documented RC-on-panel estimand. See diff_diff/staggered.py:L1869-L1883, diff_diff/staggered.py:L789-L815, and diff_diff/staggered.py:L991-L996.
Concrete fix: Compute the routing balance check after required-column validation and numeric normalization on the effective analysis sample, treating missing/non-finite outcome rows as absent, consistent with R complete-case preprocessing. Add a regression test with a full Cartesian panel plus one NaN outcome asserting allow_unbalanced_panel=True sets used_rc_on_unbalanced_panel=True, emits the RC warning, and clusters by unit.

Code Quality

P2 — Routing validation still precedes numeric normalization
Impact: The missing-column part of the prior P2 appears fixed, but the time-invariant first_treat check still runs on raw values before pd.to_numeric(). Coercible labels such as "3" and "3.0" for the same unit can be rejected as a changing cohort even though the normal fit path would normalize both to 3.0. See diff_diff/staggered.py:L1935-L1943 versus normalization at diff_diff/staggered.py:L2140-L2149.
Concrete fix: Move required-column validation, working-copy creation, pd.to_numeric() coercion, and np.inf -> 0 normalization ahead of the RC-routing validation, then validate duplicates/cohort/cluster on the normalized df.

Performance

No findings.

Maintainability

P3 — Remaining event-study bucket duplication is tracked
Impact: Analytical and bootstrap event-study bucket construction remains duplicated, but this is already tracked in TODO.md. See TODO.md:L133.
Concrete fix: None required for this PR.

Tech Debt

P3 — Deferred survey and covariate/IPW/DR parity work is tracked
Impact: survey_design= with allow_unbalanced_panel=True fails closed, and covariate/IPW/DR parity verification is documented as deferred work. See TODO.md:L132 and diff_diff/staggered.py:L1884-L1893.
Concrete fix: None required for approval.

Security

No findings.

Documentation/Tests

P2 — Missing test for NaN-outcome-triggered unbalance
Impact: The new tests cover structurally missing rows but not a rectangular panel with missing/non-finite outcomes, which is exactly where the routing bug above escapes. See tests/test_csdid_ported.py:L1978-L2258.
Concrete fix: Add the regression test described in Methodology.

Path to Approval

1. Change _is_unbalanced_panel to evaluate balance on the post-validation effective analysis sample, counting missing/non-finite outcome rows as absent.
2. Add a regression test for allow_unbalanced_panel=True on a rectangular panel with one NaN outcome that verifies RC routing, unit clustering metadata, and the warning path.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: b34c750548956ab2db2bd0b7b75e653d8786f4b5

---

Overall Assessment

⚠️ Needs changes

Executive Summary

• Prior duplicate (unit, period), changing first_treat, and time-varying cluster= checks appear addressed.
• The Registry documents RC-on-panel routing, fixed unit-cohort aggregation weights, per-unit WIF correction, and the CR1 SE deviation, so those documented differences are not defects.
• P1 remains: missing/non-finite outcome rows are used for routing detection but are not removed from the RC estimation sample.
• The new NaN-outcome test verifies routing metadata only, not that ATT cells/aggregates match the complete-case sample.
• I could not run tests locally because this environment is missing numpy.

Methodology

P1 — RC-on-panel routes on complete-case balance but estimates on unfiltered rows
Impact: The code correctly detects unbalance using finite-outcome rows, matching the stated complete-case contract, but then creates df = data.copy() and passes the original rows into _precompute_structures_rc(). The RC precompute stores df[outcome].values, and _compute_att_gt_rc() slices those arrays directly. A rectangular panel with one NaN outcome will now route to RC, but any affected 2x2 cell receives NaN in y_gt / y_gs / y_ct / y_cs, then gets materialized as a non-estimable cell instead of being estimated on the row-dropped complete-case sample. This breaks the Registry’s “ATT matches R bit-for-bit” contract for allow_unbalanced_panel=True and can change simple/event-study/group aggregates through finite-cell dropping. See diff_diff/staggered.py:L1878-L1882, diff_diff/staggered.py:L2154-L2163, diff_diff/staggered.py:L3665-L3668, diff_diff/staggered.py:L3849-L3853, and diff_diff/staggered.py:L2308-L2320.
Concrete fix: Build one effective analysis sample after required-column validation and numeric normalization, dropping at least non-finite outcome rows for the RC-on-panel path, and use that same filtered df for routing, survey/PSU resolution, validation, _precompute_structures_rc(), cohort masses, and estimation.

Code Quality

No findings beyond the methodology issue above.

Performance

No findings.

Maintainability

No findings.

Tech Debt

P3 — Deferred survey and covariate/IPW/DR parity work is tracked
Impact: TODO.md tracks remaining survey_design= and covariate/IPW/DR parity work for allow_unbalanced_panel=True; this is informational and not approval-blocking by itself. See TODO.md:L132.
Concrete fix: None required for this PR.

Security

No findings.

Documentation/Tests

P2 — NaN-outcome regression test does not assert complete-case estimation
Impact: test_nan_outcome_triggers_rc_routing only checks used_rc_on_unbalanced_panel and cluster_name, so it would pass even though the NaN row still contaminates RC cell estimation. See tests/test_csdid_ported.py:L2260-L2286.
Concrete fix: Extend the test to compare the NaN-outcome fit against the same data with that row dropped, asserting finite matching ATT cells and simple/event-study/group aggregates.

Path to Approval

1. Apply the complete-case filter to the actual RC-on-panel estimation sample, not only to _is_unbalanced_panel.
2. Add a regression test where a rectangular panel has one NaN outcome and allow_unbalanced_panel=True matches the row-dropped complete-case result for cells and aggregations.