Extend dCDH heterogeneity SE to cell-period allocator

[igerber]

Summary

• Lift the NotImplementedError gate for heterogeneity= combined with within-group-varying PSU/strata under survey_design.
• Replace the group-level IF expansion ψ_i = ψ_g * (w_i / W_g) with the cell-period post-period single-cell expansion ψ_i = ψ_g * (w_i / W_{g, out_idx}), mirroring the DID_l convention shipped in v3.1.x.
• Under PSU=group the per-obs distribution differs but the PSU-level aggregate telescopes to ψ_g, so Binder TSL and Rao-Wu replicate variance are byte-identical; under within-group-varying PSU, mass lands in the post-period PSU.
• n_bootstrap > 0 + within-group-varying PSU remains gated (follow-up PR).

Methodology references (required if estimator / math changes)

• Method name(s): ChaisemartinDHaultfoeuille._compute_heterogeneity_test (Web Appendix Section 1.5, Lemma 7); cell-period survey IF allocator.
• Paper / source link(s): de Chaisemartin & D'Haultfœuille (2020, AER); (2022, NBER WP 29873, Web Appendix §§1.5, 3.7.3). Binder (1983) stratified-PSU variance. Survey IF expansion is a library extension not derived in the source papers (see REGISTRY.md Survey IF expansion Note).
• Any intentional deviations from the source (and why): Post-period single-cell attribution for ψ_g is a library convention (documented in REGISTRY.md heterogeneity Note + Survey IF expansion Note). Formal observation-level derivation remains open (tracked in TODO.md). Empty post-period cells with zero weight drop the group's contribution (matches ATT cell allocator's empty-cell convention).

Validation

• Tests added/updated:
  • tests/test_survey_dcdh.py: flipped gating test (test_heterogeneity_with_varying_psu_succeeds — asserts all 5 inference fields finite); added TestHeterogeneityCellPeriod::test_psu_level_byte_identity_under_psu_equals_group (math-aligned PSU-sum unit test); added test_heterogeneity_auto_inject_with_varying_strata_nest_true_succeeds (newly-unblocked nest=True regime); added test_heterogeneity_multi_horizon_varying_psu_succeeds.
  • tests/test_dcdh_heterogeneity_cell_period_coverage.py: slow-tier 500-rep MC null-coverage test on a DGP with within-group-varying PSU; empirical 95% coverage inside [0.925, 0.975].
• Existing regression guard: tests/test_survey_dcdh.py::TestSurveyHeterogeneity (PSU=group numerics) passes byte-for-byte.
• Backtest / simulation / notebook evidence (if applicable): MC coverage test validates the allocator empirically on the newly-unblocked regime. No tutorial updates.

Security / privacy

• Confirm no secrets/PII in this PR: Yes

[github-actions]

Overall Assessment

⚠️ Needs changes — one unmitigated P1 in the replicate-weight heterogeneity SE path.

Executive Summary

• The affected method is ChaisemartinDHaultfoeuille._compute_heterogeneity_test (Web Appendix §1.5 / Lemma 7), specifically its survey-design variance path.
• The TSL/Binder cell-period allocator itself is documented in REGISTRY.md, and the lack of a formal observation-level derivation is already tracked in TODO.md; I am not counting that documented library extension as a defect.
• P1: the PR also applies the new post-period psi_obs allocator to replicate-weight variance, but compute_replicate_if_variance() is observation-level, not PSU-level. That breaks the PR’s claimed “byte-identical under PSU=group” property and makes the replicate branch an undocumented methodology change.
• The changed tests validate analytical TSL and Binder byte-identity, but they do not exercise the newly-unblocked replicate_weights= + within-group-varying PSU heterogeneity path.
• Minor internal docstrings/comments still describe heterogeneity + within-group-varying PSU as gated.

Methodology

• P1 Affected method: ChaisemartinDHaultfoeuille._compute_heterogeneity_test survey heterogeneity SE path in diff_diff/chaisemartin_dhaultfoeuille.py:L3931-L3972, with the replicate helper in diff_diff/survey.py:L1707-L1742. Impact: the new code assumes that the legacy group allocator and new post-period cell allocator are identical for Rao-Wu replicate variance under psu=group, and the same claim appears in docs/methodology/REGISTRY.md:L618, docs/methodology/REGISTRY.md:L652, and CHANGELOG.md:L12. That is not true for the shipped helper: compute_replicate_if_variance() forms each replicate as theta_r = sum_i ratio_ir * psi_i, so moving psi_g from all rows in a group to only the post-period cell changes theta_r whenever replicate ratios vary within group. A minimal counterexample from the implemented formula is old psi=[0.5,0.5], new psi=[0,1], replicate ratios [0.5,1.5] / [1.5,0.5], which gives replicate totals [1,1] vs [1.5,0.5]. This silently changes heterogeneity SE/p-value/CI for previously supported replicate-weight fits and is an undocumented methodology deviation, not a mere implementation choice. Concrete fix: either keep the new allocator out of the replicate branch until a replicate-specific derivation/validation exists, or explicitly document the behavioral change and change the implementation/test suite to match that contract; if byte-identity is the intended contract, the replicate path must aggregate at the group/PSU level before applying Rao-Wu reweighting.

Code Quality

• No findings.

Performance

• No findings.

Maintainability

• P3 diff_diff/chaisemartin_dhaultfoeuille.py:L5431-L5438 still says “heterogeneity and n_bootstrap > 0 currently require within-group constancy,” and tests/test_survey_dcdh.py:L1074-L1080 still says heterogeneity + within-group-varying PSU raises NotImplementedError. Impact: the internal contract description is now stale and will mislead the next maintainer who touches the gating logic. Concrete fix: update both docstrings/comments so they say only n_bootstrap > 0 + within-group-varying PSU remains gated.

Tech Debt

• P3 Informational only: the lack of a formal observation-level derivation for the post-period cell allocator is already tracked in TODO.md:L60 and documented in docs/methodology/REGISTRY.md:L652. No action required for this PR unless you decide to widen the TODO to cover the replicate-weight branch explicitly.

Security

• No findings.

Documentation/Tests

• P2 The changed tests in tests/test_survey_dcdh.py:L1141-L1278 and tests/test_dcdh_heterogeneity_cell_period_coverage.py:L91-L150 cover only the analytical TSL path and Binder byte-identity. They do not cover the newly-unblocked replicate-weight heterogeneity path where the P1 lives. Impact: the current suite would not catch the replicate non-invariance described above. Concrete fix: add a targeted regression with replicate_weights= that vary within group and heterogeneity= under psu=group or within-group-varying PSU, and assert the intended documented behavior.
• Validation note: this review was source-based; I could not run pytest here because the sandbox Python lacks the project’s NumPy/Pandas dependencies.

Path to Approval

1. Decide the contract for replicate-weight heterogeneity under the new allocator. Either re-gate/retain the legacy group-level expansion for replicate designs for now, or explicitly register the new replicate behavior as a methodology deviation instead of claiming byte-identity under psu=group.
2. Add a regression test for heterogeneity= + replicate_weights= where replicate ratios vary within group. A helper-level test against compute_replicate_if_variance() is acceptable; an end-to-end fit() test is better.
3. Update the stale internal docstrings/comments so the codebase has one consistent statement of what is and is not still gated.

[igerber]

/ai-review

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 6205b15c7aa941c143a26e0cf57d61a59f5e8b31

---

Overall Assessment

Looks good — no unmitigated P0 or P1 findings in the reviewed diff.

Executive Summary

• Affected method: ChaisemartinDHaultfoeuille._compute_heterogeneity_test survey heterogeneity SE path, plus the fit-time gate for within-group-varying survey structure in diff_diff/chaisemartin_dhaultfoeuille.py:3683 and diff_diff/chaisemartin_dhaultfoeuille.py:826.
• Cross-checking against Web Appendix §1.5 / Lemma 7, the paper-level heterogeneity estimand is still the OLS regression of S_g(Y_{g,F_g-1+\ell} - Y_{g,F_g-1}) on X_g plus F_g × D_{g,1} × S_g indicators; the PR changes only the library’s survey-design variance plumbing around that regression, not the core paper estimand. (nber.org)
• The prior replicate-variance P1 is resolved: Binder TSL now uses the documented post-period cell allocator, while the Rao-Wu replicate branch intentionally keeps the legacy group-level allocator in diff_diff/chaisemartin_dhaultfoeuille.py:3947, diff_diff/chaisemartin_dhaultfoeuille.py:3977, and diff_diff/chaisemartin_dhaultfoeuille.py:4004.
• That Binder-vs-replicate split is now explicitly documented in docs/methodology/REGISTRY.md:618, docs/methodology/REGISTRY.md:652, and CHANGELOG.md:12, so I do not see an undocumented methodology deviation in the changed path.
• The changed heterogeneity inference still goes through safe_inference() after converting nonpositive/invalid variance to NaN, so I do not see a new partial-NaN or inline-inference anti-pattern in the modified code path at diff_diff/chaisemartin_dhaultfoeuille.py:4024.
• Test coverage now includes end-to-end varying-PSU heterogeneity fits, a helper regression for replicate non-invariance, and a slow MC coverage check in tests/test_survey_dcdh.py:1142, tests/test_survey_dcdh.py:1705, and tests/test_dcdh_heterogeneity_cell_period_coverage.py:93. I could not execute them here because pytest is not installed in this environment.

Methodology

• Severity: P3 informational. Impact: Web Appendix §1.5 defines the heterogeneity target and estimator as the coefficient on X_g in an OLS regression of S_g(Y_{g,F_g-1+\ell} - Y_{g,F_g-1}) on X_g and saturated F_g × D_{g,1} × S_g indicators, with Lemma 7 proving unbiasedness under Assumptions 1, 3, and 15. The PR preserves that paper-level object and limits its change to the survey-design IF allocation/variance layer described in diff_diff/chaisemartin_dhaultfoeuille.py:3683 and docs/methodology/REGISTRY.md:618. Concrete fix: None. (nber.org)
• Severity: P3 informational. Impact: The previous undocumented replicate-variance mismatch is addressed: the code now explicitly keeps the legacy group allocator on the replicate branch, matching the updated registry note and eliminating the prior silent Rao-Wu behavior change. Concrete fix: None.

Code Quality

• No findings.

Performance

• No findings.

Maintainability

• No findings.

Tech Debt

• Severity: P3 informational. Impact: The open question of a formal observation-level derivation for the survey cell-period allocator remains tracked in TODO.md:60 and documented in docs/methodology/REGISTRY.md:652. Per the review policy, that tracked limitation is mitigated and not a blocker. Concrete fix: None.

Security

• No findings.

Documentation/Tests

• Severity: P3 informational. Impact: The PR adds targeted tests for the newly-unblocked Binder path and for the replicate-path counterexample that motivated the prior re-review, including tests/test_survey_dcdh.py:1228, tests/test_survey_dcdh.py:1614, and tests/test_dcdh_heterogeneity_cell_period_coverage.py:93. I could not run them here because pytest is unavailable. Concrete fix: None.