Add cell-period IF allocator for dCDH survey variance

[igerber]

Summary

• Extends dCDH analytical survey variance with a cell-period IF allocator that decomposes the per-group influence function U[g] into per-(g, t)-cell attributions and expands to observation level as psi_i = U_centered_per_period[g_i, t_i] * (w_i / W_{g_i, t_i}).
• Relaxes the within-group constancy validator to within-cell constancy (trivially satisfied in one-obs-per-cell panels). Designs with strata or PSU that vary across cells of a group are now accepted by the analytical TSL path.
• Adds two NotImplementedError gates for out-of-scope combinations: heterogeneity= with within-group-varying PSU/strata (follow-up PR will extend the WLS psi_obs path), and n_bootstrap > 0 with within-group-varying PSU (follow-up PR will extend the PSU-level wild bootstrap map).
• Adds a Monte Carlo coverage simulation (tests/test_dcdh_cell_period_coverage.py, @pytest.mark.slow) on a DGP with PSU varying across the cells of each group. Empirical 95% coverage hits 0.956 (±2.5pp band).
• Adds a deterministic row-sum invariant test on a hand-computed 4-group × 3-period panel verifying: closed-form joiner/stable_0 attribution, U_per_period.sum(axis=1) == U, post-period attribution placement, and cohort-centering preservation.
• Inverts TestSurveyWithinGroupValidation rejection tests into acceptance tests under the cell allocator; adds positive-rule tests for within-cell strata / PSU variation rejection and gate tests for the two NotImplementedError combos.
• Documents the new contract in docs/methodology/REGISTRY.md: post-period attribution is framed as a library convention (adopted because it preserves the group/PSU-sum identities of the prior group-level expansion and produces approximately nominal MC coverage on the test DGP) rather than a theorem derived from observation-level survey linearization. A TODO.md Methodology/Correctness entry tracks the open validation question.
• Under within-group-constant PSU (the pre-allocator accepted input), per-cell sums telescope to U_centered[g], so PSU-level Binder aggregation is byte-identical to the prior group-level expansion up to single-ULP floating-point noise. All existing survey, core, methodology, parity, and HonestDiD tests pass without modification.

Methodology references (required if estimator / math changes)

• Method name(s): ChaisemartinDHaultfoeuille (dCDH, reversible-treatment DID)
• Paper / source link(s):
  • de Chaisemartin & D'Haultfœuille (2020), AER 110(9), https://doi.org/10.1257/aer.20181169
  • de Chaisemartin & D'Haultfœuille (2022, rev. 2024), NBER WP 29873 — Web Appendix Section 3.7.3 (cohort-recentered plug-in variance)
• Any intentional deviations from the source (and why): The paper's plug-in variance assumes iid sampling; the library's survey IF expansion (proportional per-group / per-cell weight allocation) is not in the papers. The cell-period allocator extends the pre-existing library convention. Post-period attribution is documented as a convention rather than a derived result — see docs/methodology/REGISTRY.md Note "Survey IF expansion — library convention" and the TODO.md entry tracking the open derivation question.

Validation

• Tests added/updated:
  • tests/test_survey_dcdh.py — 7 new/inverted tests in TestSurveyWithinGroupValidation (acceptance + within-cell rejection + two NotImplementedError gates + row-sum invariant on a hand-computed panel)
  • tests/test_dcdh_cell_period_coverage.py — new file, 1 test, @pytest.mark.slow, 500 MC replications at 95.6% empirical coverage
  • tests/test_methodology_chaisemartin_dhaultfoeuille.py — unpack the new tuple return in test_cohort_recentering_not_grand_mean
• Backtest / simulation / notebook evidence (if applicable): MC coverage simulation is the primary empirical validation (tests/test_dcdh_cell_period_coverage.py); all 333 non-slow tests and the slow coverage test pass.

Security / privacy

• Confirm no secrets/PII in this PR: Yes

---

Generated with Claude Code

[github-actions]

Overall Assessment

⚠️ Needs changes

Executive Summary

• The PR affects ChaisemartinDHaultfoeuille survey inference for DID_M, DID_+, DID_-, event-study horizons, and placebo horizons by introducing a cell-period IF allocator.
• The core methodology deviation from the papers is documented in the registry and tracked in TODO.md, so I am treating that part as informational, not a blocker.
• There is one unmitigated P1: the newly advertised support for within-group-varying strata is incomplete when survey_design.psu is omitted, because the auto-injected PSU path still fails before the cell allocator runs.
• There is also a non-blocking performance regression: dense per-period IF tensors are now built on all paths, including non-survey fits that never consume them.
• Tests add useful invariant and coverage checks, but they do not cover the omitted-psu auto-inject path where the P1 issue appears.

Methodology

• Severity P1. Impact: the new public contract says that, with no explicit psu, fit() auto-injects psu=<group_col> as a safe default while also supporting strata that vary across cells of a group, but the auto-injected path preserves nest=False, so a group that spans multiple strata still trips the survey resolver’s cross-stratum PSU check before the new allocator can run. See diff_diff/chaisemartin_dhaultfoeuille.py:648, diff_diff/chaisemartin_dhaultfoeuille.py:753, and diff_diff/survey.py:313. Concrete fix: either synthesize a stratum-nested injected PSU when auto-injecting under within-group-varying strata, or narrow the documented contract and raise a targeted pre-resolution error telling users that this case requires explicit psu plus nest=True.
• Severity P3. Impact: the cell-period/post-period allocation itself is a documented library convention, not a paper-derived survey linearization, and the open derivation question is already tracked. I am not counting that against approval. See docs/methodology/REGISTRY.md:652 and TODO.md:60. Concrete fix: none required for approval.

Code Quality

• No material findings beyond the contract issue above.

Performance

• Severity P2. Impact: the PR now allocates dense U_per_period tensors unconditionally, even on non-survey/plugin paths that never use them. That adds O(G*T) memory to the main ATT path and O(G*T*L_max) to event-study/placebo paths for all fits, not just the new survey case. See diff_diff/chaisemartin_dhaultfoeuille.py:1643, diff_diff/chaisemartin_dhaultfoeuille.py:4374, and diff_diff/chaisemartin_dhaultfoeuille.py:4942. Concrete fix: only build/store the 2D IF tensors when a survey design actually needs the cell allocator, or compute them lazily on demand.

Maintainability

• Severity P3. Impact: new internal comments describe the allocator as “methodologically correct” / “valid” even though the public methodology note and TODO.md explicitly say the derivation is still open. That creates an avoidable maintenance trap. See diff_diff/chaisemartin_dhaultfoeuille.py:5044, diff_diff/chaisemartin_dhaultfoeuille.py:5233, docs/methodology/REGISTRY.md:652, and TODO.md:60. Concrete fix: soften those comments to match the registry wording (“library convention”, “empirically validated”, “formal derivation deferred”).

Tech Debt

• Severity P3. Impact: the main methodological limitation introduced here is already tracked under the repo’s deferred-work policy, so it is not a blocker. See TODO.md:60. Concrete fix: none required in this PR.

Security

• No findings.

Documentation/Tests

• Severity P2. Impact: the new varying-strata acceptance coverage only exercises the explicit psu="group", nest=True path; it does not cover the omitted-psu auto-inject path that the docs now describe as a safe default, which is why the P1 issue is currently unguarded. See tests/test_survey_dcdh.py:1120 and diff_diff/chaisemartin_dhaultfoeuille.py:648. Concrete fix: add a regression test with SurveyDesign(weights="pw", strata="stratum") on a panel where a group spans multiple strata, and assert the intended supported or documented behavior.
• I could not execute the tests locally in this review environment because the available python is missing numpy.

Path to Approval

1. Resolve the omitted-psu / varying-strata interaction: either make auto-injection synthesize a stratum-nested PSU when groups span multiple strata, or explicitly narrow the public contract and raise a targeted error before survey resolution.
2. Add a regression test for SurveyDesign(weights="pw", strata="stratum") with within-group-varying strata and no explicit psu, so the intended behavior is enforced.

[igerber]

/ai-review

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: dedc2a6e38ca50aefeaf2411fdddf7d91c7c8938

---

Overall Assessment

⚠️ Needs changes

Executive Summary

• The prior review’s main issue is only partially resolved: the PR now turns the default no-psu + varying-strata failure into a targeted error, but the fix overcorrects.
• Severity P1 [Newly identified]: the new pre-auto-inject guard rejects SurveyDesign(..., strata=..., nest=True) with no explicit psu, even though the auto-inject path already preserves nest=True and SurveyDesign.resolve() can make the synthesized psu=group unique within strata.
• The cell-period allocator itself remains a documented library convention rather than a paper-derived survey linearization; that is properly noted in the Methodology Registry and tracked in TODO.md, so it is informational only.
• The earlier non-blocking performance concern looks resolved: dense per-period IF tensors are now only built when survey inference is active.
• I could not run the tests locally in this environment because numpy is not installed, so this review is diff-based.

Methodology

• Severity P1 [Newly identified]. Impact: the new auto-inject guard hard-errors on a valid nested-strata design. The comment says the failure mode is specifically the nest=False resolver path, but the implementation raises whenever strata vary within group, without checking nest. That blocks SurveyDesign(weights="pw", strata="stratum", nest=True) even though the auto-injected design preserves nest=True and the resolver explicitly supports making (stratum, psu) unique within strata. Locations: diff_diff/chaisemartin_dhaultfoeuille.py:L729-L759, diff_diff/chaisemartin_dhaultfoeuille.py:L785-L795, diff_diff/survey.py:L299-L303, docs/methodology/REGISTRY.md:L652-L652. Concrete fix: only raise on the no-psu varying-strata path when the synthesized psu=<group> would be resolved with nest=False; allow the nest=True case to proceed and update the docs accordingly.

Code Quality

No new findings.

Performance

No new findings. The previous dense-allocation concern appears addressed by the new compute_per_period=_obs_survey_info is not None gating.

Maintainability

No new findings. The new comments are now aligned with the Registry/TODO wording and no longer overclaim a formal derivation.

Tech Debt

• Severity P3. Impact: the open derivation question for the cell-period/post-period survey IF allocator is explicitly documented and tracked, so it is not a blocker under the project’s deferred-work policy. Locations: docs/methodology/REGISTRY.md:L652-L653, TODO.md:L59-L60. Concrete fix: none required for approval.

Security

No findings.

Documentation/Tests

• Severity P2. Impact: the new tests only cover the unsupported default nest=False auto-inject path, so the over-broad rejection of the still-valid nest=True path is unguarded. Location: tests/test_survey_dcdh.py:L1175-L1192. Concrete fix: add a regression test using SurveyDesign(weights="pw", strata="stratum", nest=True) with no explicit psu, and assert it succeeds and matches the explicit psu="group", nest=True path on SE / df_survey.

Local execution note: I could not run the targeted pytest files here because the environment is missing numpy.

Path to Approval

1. Narrow the new pre-auto-inject ValueError so it only fires for the actually unsupported case: varying strata with no explicit psu and nest=False.
2. Add a regression test for no-psu, varying-strata, nest=True, and compare it against explicit SurveyDesign(weights="pw", strata="stratum", psu="group", nest=True).
3. Update the fit() docstring and Methodology Registry note so they distinguish the unsupported nest=False auto-inject path from the valid nest=True auto-inject path.

[igerber]

/ai-review

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: ca7b6bfa5d0e321f3a2fb60fb3340148a44b3d7f

---

Overall Assessment

✅ Looks good

Executive Summary

• The prior re-review blocker on omitted psu + varying strata + nest=True appears resolved in code and tests: the pre-auto-inject guard now only rejects the nest=False path, and the new regression covers the valid nest=True path at diff_diff/chaisemartin_dhaultfoeuille.py:L750-L810 and tests/test_survey_dcdh.py:L1175-L1213.
• I did not find new unmitigated P0/P1 issues in the changed estimator/variance paths. The new per-period IF tensor is threaded through overall, joiner/leaver, event-study, and placebo SE computation, and inference still routes through safe_inference().
• The cell-period/post-period survey IF allocator remains a library convention rather than a paper-derived observation-level linearization, but that deviation is explicitly documented in the Methodology Registry and tracked in TODO.md, so it is informational only under the project policy at docs/methodology/REGISTRY.md:L652-L653 and TODO.md:L60.
• Local execution note: this is a diff-based re-review. I could not run pytest here because the environment is missing numpy and pytest.

Methodology

• No unmitigated P0/P1 findings. The changed behavior matches the documented contract in the Methodology Registry for the nest=True auto-inject case, within-cell constancy checks, and the documented survey-allocator scope at diff_diff/chaisemartin_dhaultfoeuille.py:L750-L865 and docs/methodology/REGISTRY.md:L652-L653.

Code Quality

• No findings.

Performance

• No findings. The dense per-period attribution tensors are now only requested when survey inference is active, which resolves the prior always-allocate concern visible in the changed call sites.

Maintainability

• No findings. The fit docstring, in-code comments, tests, and Registry note now describe the same allocator contract and follow-up PR boundaries.

Tech Debt

• Severity P3. Impact: the survey cell-period allocator’s post-period attribution is still a library convention, not a paper-derived observation-level linearization, but that limitation is explicitly tracked and therefore non-blocking under the project’s deferred-work policy at docs/methodology/REGISTRY.md:L652-L653 and TODO.md:L60. Concrete fix: none required for approval.

Security

• No findings.

Documentation/Tests

• No new findings in the diff. The updated tests cover the resolved nest=True auto-inject path, within-cell rejection rules, the row-sum invariant, and the new slow coverage simulation at tests/test_survey_dcdh.py:L1175-L1543 and tests/test_dcdh_cell_period_coverage.py:L1-L134. Local execution note: I could not run them in this environment.