Skip to content

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

Merged
igerber merged 1 commit into
mainfrom
fix/cs-unbalanced-pg-event-study
Jul 7, 2026
Merged

feat(callaway): allow_unbalanced_panel — RC-on-panel parity with R did#631
igerber merged 1 commit into
mainfrom
fix/cs-unbalanced-pg-event-study

Conversation

@igerber

@igerber igerber commented Jul 7, 2026

Copy link
Copy Markdown
Owner

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=FALSEDRDID::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

github-actions Bot commented Jul 7, 2026

Copy link
Copy Markdown

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.

@igerber igerber force-pushed the fix/cs-unbalanced-pg-event-study branch from 70b017b to f8c25c9 Compare July 7, 2026 10:04
@github-actions

github-actions Bot commented Jul 7, 2026

Copy link
Copy Markdown

🔁 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.

@igerber igerber force-pushed the fix/cs-unbalanced-pg-event-study branch from f8c25c9 to 0ece407 Compare July 7, 2026 10:10
@github-actions

github-actions Bot commented Jul 7, 2026

Copy link
Copy Markdown

🔁 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.

@igerber igerber force-pushed the fix/cs-unbalanced-pg-event-study branch from 0ece407 to b34c750 Compare July 7, 2026 10:19
@github-actions

github-actions Bot commented Jul 7, 2026

Copy link
Copy Markdown

🔁 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.

@igerber igerber force-pushed the fix/cs-unbalanced-pg-event-study branch 2 times, most recently from 93fa908 to b0c0e7d Compare July 7, 2026 10:29
@github-actions

github-actions Bot commented Jul 7, 2026

Copy link
Copy Markdown

🔁 AI review rerun (requested by @igerber)

Head SHA: b0c0e7da6c46c7b6490a9edb80f67bfdc4fe831a


Overall Assessment

⚠️ Needs changes

Executive Summary

  • The previous NaN-outcome contamination issue is mostly addressed: RC routing now drops non-finite outcome rows before precompute/estimation.
  • The Registry documents RC-on-panel routing, fixed unit-cohort pg, per-unit WIF correction, and CR1 SE deviation, so those are not defects.
  • P1 remains in the complete-case balance detector: it uses finite-outcome rows for the numerator but unfiltered units/periods for the denominator.
  • This can route to RC when R-style complete-case data would be balanced and should keep the panel path inert.
  • Duplicate (unit, period), changing first_treat, and time-varying cluster= checks are covered on the RC route.
  • I could not run tests locally: numpy and pytest are not installed. git diff --check passes.

Methodology

P1 — Complete-case balance detection can incorrectly route balanced complete-case panels to RC
Impact: The implementation says balance should match R’s complete-case preprocessing, but _is_unbalanced_panel computes _cc_cells on finite-outcome rows while the denominator still uses data[unit].nunique() * data[time].nunique() from the unfiltered data. If all rows for one unit, or all rows for one period, have non-finite outcomes, the complete-case sample can be balanced after dropping those rows, but this code still sets _route_as_rc=True. That violates the documented balanced-inert/R-preprocessing contract and can silently switch estimator/variance path, especially with covariates. See diff_diff/staggered.py:L1879-L1883 and diff_diff/staggered.py:L1936-L1937.
Concrete fix: Build a complete-case frame first, e.g. cc = data.loc[_finite_out], and compute both numerator and denominator from cc: cc.drop_duplicates([unit, time]).shape[0] < cc[unit].nunique() * cc[time].nunique(). Add regression cases where an entire unit and an entire period are removed by non-finite outcomes and the flag remains inert if the resulting complete-case panel is rectangular.

Code Quality

No findings beyond the methodology issue above.

Performance

No findings.

Maintainability

No findings.

Tech Debt

P3 — Remaining survey_design= and covariate/IPW/DR parity work is tracked
Impact: TODO.md records the deferred survey_design= × allow_unbalanced_panel work and covariate/IPW/DR R-parity verification, and the code fails closed for survey_design= on the routed path. 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 — Complete-case routing tests miss the “dropped whole unit/period leaves balanced panel” case
Impact: The new NaN-outcome test compares a single missing cell to a row-dropped fit, which covers the prior contamination bug, but it does not catch the remaining denominator issue above. See tests/test_csdid_ported.py:L2260-L2311.
Concrete fix: Add tests where all outcomes for one unit, and separately all outcomes for one period, are NaN; assert allow_unbalanced_panel=True does not set used_rc_on_unbalanced_panel when the complete-case sample is balanced.

Path to Approval

  1. Change _is_unbalanced_panel to compute both observed complete-case cells and the unit×period denominator from the finite-outcome complete-case frame.
  2. Add regression tests for all-NaN unit and all-NaN period cases that become rectangular after complete-case filtering, asserting the flag is inert and results match the panel path.

@igerber igerber force-pushed the fix/cs-unbalanced-pg-event-study branch from b0c0e7d to 896da63 Compare July 7, 2026 10:37
@github-actions

github-actions Bot commented Jul 7, 2026

Copy link
Copy Markdown

🔁 AI review rerun (requested by @igerber)

Head SHA: 896da632e1c5f2df50aecd8cca392279392218ac


Overall Assessment

✅ Looks good

Executive Summary

  • The prior P1 is resolved: complete-case balance now uses the finite-outcome frame for both observed cells and the unit×period denominator.
  • allow_unbalanced_panel=True routing, unit clustering, fixed cohort pg weights, and WIF correction match the documented registry contract.
  • The CR1 SE difference from R is documented as a Registry deviation, so it is not a defect.
  • New parameter propagation is present in constructor, get_params(), aggregations, bootstrap paths, and results metadata.
  • survey_design= and covariate/IPW/DR parity limitations are tracked in TODO.md, so they are informational.

Methodology

No unmitigated findings.

The previous complete-case routing issue is fixed at diff_diff/staggered.py:L1879-L1893, with regression coverage at tests/test_csdid_ported.py:L2313-L2340. The changed RC-on-panel behavior is documented in docs/methodology/REGISTRY.md:L565-L594, and the implementation follows it via RC routing and complete-case filtering at diff_diff/staggered.py:L1901-L1947, panel-structure validation at diff_diff/staggered.py:L1959-L1987, unit clustering at diff_diff/staggered.py:L2138-L2145, and fixed cohort aggregation/WIF correction at diff_diff/staggered_aggregation.py:L19-L45 and diff_diff/staggered_aggregation.py:L420-L429.

Code Quality

No findings.

Performance

No findings.

Maintainability

No findings.

Tech Debt

Severity: P3 informational
Impact: survey_design= × allow_unbalanced_panel and covariate/IPW/DR R-parity verification remain deferred, but they are explicitly tracked in TODO.md:L132 and the survey path fails closed on routed unbalanced panels at diff_diff/staggered.py:L1905-L1914.
Concrete fix: No action required for this PR.

Security

No findings. I did not see secrets or PII in the changed files; the new golden data is synthetic benchmark data.

Documentation/Tests

No blocking findings.

Coverage includes R golden parity for cells/simple/event-study/group, bootstrap aggregation weighting, metadata, panel validation, NaN complete-case routing, and the prior complete-case rectangular regression case in tests/test_csdid_ported.py:L1978-L2354.

Verification: git diff --check passes. I could not run the targeted pytest class because this environment does not have pytest installed (/usr/bin/python: No module named pytest).

@igerber igerber added the ready-for-ci Triggers CI test workflows label Jul 7, 2026
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
@igerber igerber force-pushed the fix/cs-unbalanced-pg-event-study branch from 896da63 to cfea355 Compare July 7, 2026 11:32
@github-actions

github-actions Bot commented Jul 7, 2026

Copy link
Copy Markdown

🔁 AI review rerun (requested by @igerber)

Head SHA: cfea355a07a5e026ae78b1cfedbce6215c17fbd7


Overall Assessment

✅ Looks good

Executive Summary

  • The prior complete-case routing issue is resolved: balance is assessed on finite-outcome rows, and routed RC fits drop those rows consistently.
  • CallawaySantAnna(allow_unbalanced_panel=True) now follows the Registry contract for unbalanced panels: RC levels estimator, unit clustering, fixed unit-cohort pg weights, and WIF over-count correction.
  • The analytical SE difference from R is documented as a Registry deviation, so it is not a defect.
  • New parameter propagation is present in constructor, get_params(), result metadata, aggregations, and bootstrap paths.
  • No unmitigated P0/P1 findings.

Methodology

No unmitigated findings.

The affected method is CallawaySantAnna unbalanced-panel handling. The routing and complete-case behavior match the documented Registry contract: detection/filtering at diff_diff/staggered.py:L1865-L1947, panel-structure validation at diff_diff/staggered.py:L1959-L1987, RC precompute cohort/unit masses at diff_diff/staggered.py:L3720-L3755, fixed aggregation weights at diff_diff/staggered_aggregation.py:L19-L45, WIF correction at diff_diff/staggered_aggregation.py:L262-L431, and bootstrap reuse at diff_diff/staggered_bootstrap.py:L241-L256 and diff_diff/staggered_bootstrap.py:L723-L789.

The CR1 finite-sample SE factor is documented as a deviation at docs/methodology/REGISTRY.md:L591-L595, so it is informational, not a methodology defect.

Code Quality

No findings. The changed inference paths use safe_inference() / safe_inference_batch() rather than adding inline t-stat / p-value / CI calculations.

Performance

No findings. The shared fixed_cohort_agg_weights() helper avoids duplicating the cohort-mass logic across analytical and bootstrap aggregation paths.

Maintainability

No findings. The new parameter is propagated through __init__, get_params(), fit metadata, results, analytical aggregations, and bootstrap aggregation.

Tech Debt

Severity: P3 informational
Impact: survey_design= with routed unbalanced panels and covariate/IPW/DR R-parity verification remain deferred, but this is explicitly tracked in TODO.md:L131, and the survey route fails closed at diff_diff/staggered.py:L1905-L1914.
Concrete fix: No action required for this PR.

Security

No findings. I did not see secrets or PII in the changed files; the golden data appears synthetic.

Documentation/Tests

Severity: P3 informational
Impact: The docs/docstring wording says survey_design= with the flag raises, while the implementation raises only when the flag actually routes an unbalanced panel. That behavior is consistent with balanced-panel inertness, but the wording could be clearer.
Concrete fix: Clarify the wording to “survey_design= with a routed unbalanced panel raises NotImplementedError; balanced panels remain inert.”

Coverage added in tests/test_csdid_ported.py:L1978-L2354 covers R golden parity, simple/event-study/group aggregation, bootstrap weighting, metadata, validation, default warnings, and complete-case NaN routing.

Verification: git diff --check passed. I could not run the targeted pytest class because this environment does not have pytest installed (No module named pytest).

@igerber igerber merged commit a564bd7 into main Jul 7, 2026
27 checks passed
@igerber igerber deleted the fix/cs-unbalanced-pg-event-study branch July 7, 2026 13:21
@igerber igerber mentioned this pull request Jul 9, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ready-for-ci Triggers CI test workflows

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant