Add survey-aware bootstrap for all estimators (Phase 6)

[igerber]

Summary

• Add survey-aware bootstrap inference for all 8 bootstrap-using estimators
• Two strategies: PSU-level multiplier bootstrap (CS, ImputationDiD, TwoStageDiD, ContinuousDiD, EfficientDiD) and Rao-Wu rescaled bootstrap (SunAbraham, SyntheticDiD, TROP)
• Expand CallawaySantAnna analytical survey support to full strata/PSU/FPC via compute_survey_if_variance()
• Add shared infrastructure: generate_survey_multiplier_weights_batch, generate_rao_wu_weights, compute_survey_if_variance, aggregate_to_psu
• Thread survey weights through bootstrap aggregation/IF/GMM score computation for all estimators
• Add edge-case guards: lonely_psu="adjust" rejection, FPC validation, single-PSU handling, SyntheticDiD placebo+full-design guard
• Update REGISTRY.md with Phase 6 survey bootstrap methodology section
• Iteratively refined through 5 rounds of AI review (gpt-5.4-pro)

Methodology references (required if estimator / math changes)

• Method name(s): Rao-Wu rescaled bootstrap, PSU-level multiplier bootstrap, Taylor Series Linearization
• Paper / source link(s):
  • Rao & Wu (1988) "Resampling Inference with Complex Survey Data", JASA 83(401)
  • Rao, Wu & Yue (1992) "Some Recent Work on Resampling Methods for Complex Surveys", Survey Methodology 18(2)
  • Kolenikov (2010) "Resampling Variance Estimation for Complex Survey Data"
  • Shao (2003) "Impact of the Bootstrap on Sample Surveys", Statistical Science 18(2)
• Any intentional deviations from the source (and why):
  • FPC enters Rao-Wu via adjusted resample size m_h = round((1-f_h)*(n_h-1)) per Rao-Wu-Yue (1992) Section 3
  • TROP uses cross-classified pseudo-strata (survey_stratum × treatment_group) for Rao-Wu
  • lonely_psu="adjust" rejected for bootstrap paths (analytical path supports it)
  • Rust TROP bootstrap remains pweight-only; Python fallback for full design

Validation

• Tests added/updated: tests/test_survey.py, tests/test_survey_phase3.py, tests/test_survey_phase4.py, tests/test_survey_phase5.py
• 285 survey tests passing across all phases
• All deferral tests converted to positive tests
• Smoke tests + scale invariance + uniform weight equivalence tests

Security / privacy

• Confirm no secrets/PII in this PR: Yes

---

Generated with Claude Code

[github-actions]

Overall Assessment

⛔ Blocker

The PR introduces one unmitigated P0 in the shared Rao-Wu bootstrap path and three P1 methodology/documentation issues in the new survey-bootstrap support. The P0 is a silent statistical correctness bug in shared code, so this should not merge as-is.

Executive Summary

• The shared Rao-Wu rescaled bootstrap cannot produce the zero-variance full-census FPC case, so estimators that use it can report nonzero bootstrap uncertainty where the survey design implies zero variance.
• ContinuousDiD’s new survey bootstrap targets the wrong ACRT^{glob} estimand under non-uniform survey weights: the point estimator and analytical IF are weighted, but the bootstrap perturbation is unweighted.
• Both new survey-bootstrap generators reject lonely_psu="adjust", but the new Phase 6 registry text does not disclose that restriction.
• TROP’s Rao-Wu path changes the no-strata design by treating treatment status as strata, which changes the bootstrap/FPC construction in the strata=None case without a clear Phase 6 note.
• The new tests are mostly smoke tests; they do not cover Rao-Wu full-census FPC or ContinuousDiD’s survey-bootstrap overall_acrt_* outputs, which is why the methodology regressions above slip through.

Methodology

• Severity: P0. Method: Rao-Wu rescaled bootstrap (SunAbraham, SyntheticDiD, TROP). Impact: Within this library’s own survey methodology, full-census FPC (f_h = 1) is a zero-variance case in the analytical survey path at diff_diff/survey.py:658 and diff_diff/survey.py:713, with regression coverage at tests/test_survey.py:2504 and tests/test_survey.py:2533. The new shared Rao-Wu generator clamps m_h to at least 1 at diff_diff/bootstrap_utils.py:608, so its “full census” guard at diff_diff/bootstrap_utils.py:612 is unreachable. Callers at diff_diff/sun_abraham.py:1372, diff_diff/synthetic_did.py:760, diff_diff/trop_local.py:1163, and diff_diff/trop_global.py:1149 can therefore emit non-degenerate replicate weights and nonzero SEs for full-census FPC designs. That is silent wrong statistical output. Concrete fix: Special-case f_h >= 1 before drawing PSUs and keep original weights / zero perturbation for that stratum; also update the Rao-Wu Phase 6 note at docs/methodology/REGISTRY.md:1960 to document the census case.

• Severity: P1. Method: ContinuousDiD PSU-level multiplier bootstrap for ACRT^{glob}. Impact: The registry defines ACRT^{glob} as the treated average marginal effect at docs/methodology/REGISTRY.md:468. The survey point estimator and analytical IF both use survey-weighted treated averages at diff_diff/continuous_did.py:955, diff_diff/continuous_did.py:1005, and diff_diff/continuous_did.py:1176. But the new survey bootstrap perturbation uses an unweighted mean at diff_diff/continuous_did.py:1441. With non-uniform treated survey weights, overall_acrt_se, CI, and p-value no longer target the same estimand as the point estimate. Concrete fix: Use the stored weighted dpsi_bar (or equivalent weighted average when w_treated is present) in _bootstrap_gt_cell, and add a regression test covering overall_acrt_* under non-uniform treated survey weights.

• Severity: P1. Method: survey-aware bootstrap support in generate_survey_multiplier_weights_batch() and generate_rao_wu_weights(). Impact: The general survey registry still presents lonely_psu modes including "adjust" at docs/methodology/REGISTRY.md:1887, and the new Phase 6 section advertises survey-aware bootstrap support at docs/methodology/REGISTRY.md:1945 and docs/methodology/REGISTRY.md:1955. But both new generators hard-reject "adjust" at diff_diff/bootstrap_utils.py:454 and diff_diff/bootstrap_utils.py:570. That is an undocumented behavioral/methodology restriction in a newly advertised feature. Concrete fix: Either implement "adjust" for bootstrap paths or add explicit **Note:** / **Deviation from R:** entries in REGISTRY.md stating that bootstrap currently supports only "remove" and "certainty", plus tests for the present rejection.

• Severity: P1. Method: TROP Rao-Wu resampling design. Impact: In the TROP full-design bootstrap builders, the survey_design.strata is None case still creates strata from treatment group at diff_diff/trop_local.py:1073 and diff_diff/trop_global.py:1091. The shared survey layer otherwise supports genuinely unstratified FPC designs at diff_diff/survey.py:212. The Phase 6 Rao-Wu note describes resampling within actual strata at docs/methodology/REGISTRY.md:1960, while the TROP pseudo-strata note only describes the case where survey strata and treatment groups both exist at docs/methodology/REGISTRY.md:1977. In the strata=None case, the implementation therefore changes the bootstrap/FPC design without a clear matching Phase 6 note. Concrete fix: Either keep TROP unstratified when survey_design.strata is None, or document and justify the treatment-group pseudo-strata rule for that case, including its FPC interpretation, and add no-strata+FPC tests for both local and global TROP.

Code Quality

• No additional findings beyond the methodology issues above.

Performance

• No findings.

Maintainability

• No additional findings beyond the methodology/doc alignment issues above.

Tech Debt

• No TODO.md mitigation applies to the P0/P1 items above. Under the stated rubric, these are correctness or undocumented-methodology issues, not deferrable tech debt.

Security

• No findings.

Documentation/Tests

• Severity: P2. Impact: The new tests cover smoke/success paths but miss the key edge cases above. The ContinuousDiD survey-bootstrap test only checks overall_att / overall_att_se at tests/test_survey_phase3.py:673, not overall_acrt_*. The Phase 5 “full design” fixture omits fpc at tests/test_survey_phase5.py:109, so the SyntheticDiD and TROP full-design smoke tests at tests/test_survey_phase5.py:179, tests/test_survey_phase5.py:532, and tests/test_survey_phase5.py:547 never exercise the broken full-census Rao-Wu branch. Concrete fix: Add targeted regression tests for full-census FPC Rao-Wu behavior, lonely_psu="adjust" rejection/support, and ContinuousDiD overall_acrt_* under non-uniform survey weights.

Path to Approval

1. Fix diff_diff/bootstrap_utils.py so Rao-Wu handles f_h >= 1 as a census case with zero perturbation, then add stratified and unstratified full-census regression tests in at least one Rao-Wu estimator and a generator-level unit test.
2. Fix diff_diff/continuous_did.py so survey bootstrap ACRT^{glob} uses the same weighted treated derivative average as the point estimator / analytical IF, then add a regression test that asserts overall_acrt_se, CI, and p-value under non-uniform survey weights.
3. Resolve the lonely_psu="adjust" mismatch by either implementing it in both survey-bootstrap generators or documenting the unsupported mode with explicit **Note:** / **Deviation from R:** labels in docs/methodology/REGISTRY.md, plus tests for the current behavior.
4. Resolve the TROP strata=None design mismatch by either keeping the Rao-Wu path unstratified in that case or explicitly documenting the treatment-group pseudo-strata rule, including FPC semantics, and add SurveyDesign(weights=..., fpc=...) no-strata tests for both TROP methods.