perf(callaway-santanna): chunk multiplier-bootstrap weight generation to bound peak memory

[igerber]

Summary

• Chunk the CallawaySantAnna multiplier-bootstrap weight generation and consumption over draw-blocks so the dense (n_bootstrap × n_units) weight matrix is never materialized. New diff_diff/bootstrap_chunking.py helper (iter_weight_blocks, iter_survey_multiplier_weight_blocks, compute_block_size); staggered_bootstrap.py refactored to tile the bootstrap over draws into preallocated small output arrays.
• Memory: peak RSS at the production default of 999 reps drops ~79% at 500k units (11.6 GB → 2.4 GB) and ~68% at 1M units (10.8 GB → 3.4 GB). The millions-of-units × 999-rep regime — previously ~40 GB+ just for the weight matrix (hard OOM) — now stays near the fit's memory floor.
• The generated weight stream is bit-identical on both backends (Rust absolute per-row seeding; NumPy in-order stream); end-to-end bootstrap SEs match the un-chunked path to within BLAS floating-point reassociation (~1 ULP, far below bootstrap Monte-Carlo error).

Methodology references

• Method: CallawaySantAnna multiplier bootstrap. This is an implementation/performance change only — the bootstrap still perturbs the same ATT(g,t) influence functions and combined IF+WIF aggregations, and survey/cluster weights are still generated at PSU level. No change to the estimator or variance formula (see docs/methodology/REGISTRY.md CallawaySantAnna).
• Intentional deviations: None. Numerically equivalent to the prior path within floating-point reassociation of the BLAS reductions.

Validation

• Tests: tests/test_bootstrap_chunking.py (new, 33 tests both backends): block sizing, weight-stream bit-identity (chunked vs single-block and vs the legacy generator), survey-helper parity vs generate_survey_multiplier_weights_batch (unstratified tiled, psu=None, FPC scaling, stratified fallback), and end-to-end CS chunk-invariance for cluster="unit", cluster=None (default), and non-identity cluster fan-out — covering ATT(g,t), overall, event-study, group, and cband. tests/test_methodology_callaway.py import repointed to the canonical generator. All pass under DIFF_DIFF_BACKEND=rust and =python.
• Deferred (tracked in TODO.md Actionable Backlog → Performance): stratified survey-PSU generation and the EfficientDiD/HAD bootstraps still build full weight matrices; wiring those through the new helper is follow-up work.

Security / privacy

• Confirm no secrets/PII in this PR: Yes

[github-actions]

Overall Assessment

✅ Looks good — no unmitigated P0/P1 findings.

Executive Summary

• Affected method: CallawaySantAnna multiplier bootstrap, including PSU-level survey/cluster bootstrap.
• The change tiles weight generation over bootstrap draws and preserves the documented multiplier-bootstrap target: perturbing ATT(g,t) IFs and combined IF+WIF aggregations.
• I found no undocumented methodology deviation, missing assumption check, or incorrect SE/variance change.
• NaN propagation remains guarded through the existing bootstrap stats helpers and the single-PSU cluster variance guard.
• Deferred non-CS / stratified chunking gaps are tracked in TODO.md, so they are P3 informational only.
• I could not run tests locally because this environment lacks pytest and numpy.

Methodology

Finding: None

• Severity: None
• Impact: The changed path matches the registry contract: CS bootstrap perturbs combined influence functions, including WIF, and survey bootstrap uses PSU-level multipliers with FPC scaling.
• Concrete fix: None.

References checked: docs/methodology/REGISTRY.md:L493-L496, docs/methodology/REGISTRY.md:L610-L614, docs/methodology/REGISTRY.md:L4237-L4245; implementation in diff_diff/staggered_bootstrap.py:L348-L499 and diff_diff/bootstrap_chunking.py:L103-L185.

Code Quality

Finding: None

• Severity: None
• Impact: The iterator abstraction is scoped and keeps RNG/block semantics localized. The old full-matrix consumers are rewritten to fill preallocated output arrays block-by-block without changing aggregation formulas.
• Concrete fix: None.

Performance

Finding: Deferred chunking gaps are tracked

• Severity: P3 informational
• Impact: Stratified survey PSU generation still materializes the full PSU-weight matrix, and EfficientDiD/HAD still use full weight matrices. This is performance debt, not a correctness defect, and the PR adds it to TODO.md.
• Concrete fix: No action required for this PR. Follow the tracked work in TODO.md:L69 if those paths become memory bottlenecks.

Maintainability

Finding: None

• Severity: None
• Impact: The new helper separates block-size calculation, standard weight blocks, and survey PSU blocks cleanly enough for future consumers.
• Concrete fix: None.

Tech Debt

Finding: Tracked deferred work only

• Severity: P3 informational
• Impact: The unchunked stratified/EfficientDiD/HAD surfaces are documented as follow-up work, which satisfies the project’s deferred-work policy.
• Concrete fix: None for approval.

Security

Finding: None

• Severity: None
• Impact: No secrets, external calls, unsafe deserialization, or security-sensitive changes found in the diff.
• Concrete fix: None.

Documentation / Tests

Finding: Local tests not executed

• Severity: P3 informational
• Impact: Static review only; this environment has neither pytest nor numpy, so I could not execute tests/test_bootstrap_chunking.py.
• Concrete fix: Ensure CI runs the added chunking tests under both DIFF_DIFF_BACKEND=rust and DIFF_DIFF_BACKEND=python, as the PR intends.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 4c6f20a7d3b7bd4d8ce4dac363634b63a8518a81

---

Overall Assessment

✅ Looks good — no unmitigated P0/P1 findings.

Executive Summary

• Affected method: CallawaySantAnna multiplier bootstrap.
• The chunked draw-block implementation preserves the documented CS bootstrap target: ATT(g,t) IFs plus combined IF/WIF aggregations.
• No NaN/inference anti-patterns found in the changed bootstrap code.
• Deferred stratified/EfficientDiD/HAD chunking gaps are tracked in TODO.md.
• I found P2 documentation/test regressions around survey methodology tracking and references, but no runtime correctness blocker.

Methodology

Finding: CS chunking is methodology-equivalent

• Severity: None
• Impact: The registry requires multiplier bootstrap perturbation of combined IFs, including WIF, and the changed code still does that for overall/event-study surfaces while group aggregation reuses the per-draw ATT(g,t) rows. See docs/methodology/REGISTRY.md:L493-L496, diff_diff/staggered_bootstrap.py:L423-L497.
• Concrete fix: None.

Finding: Survey df citation regressed

• Severity: P2
• Impact: docs/methodology/REGISTRY.md:L4176-L4177 now cites Korn & Graubard (1990) as JASA 85(409), while the methodology materials identify the survey-df/Bonferroni-t source as The American Statistician 44(4), 270-276. This weakens source-material traceability but does not change code output.
• Concrete fix: Restore the corrected venue/title details in REGISTRY.md and restore the matching reference entry in docs/references.rst.

Code Quality

Finding: None

• Severity: None
• Impact: diff_diff/bootstrap_chunking.py:L44-L185 keeps block sizing, standard weights, and survey PSU blocks localized. The staggered bootstrap consumer fills preallocated output arrays block-by-block without duplicating final inference logic.
• Concrete fix: None.

Performance

Finding: Deferred non-CS/stratified chunking gaps are tracked

• Severity: P3 informational
• Impact: EfficientDiD/HAD and stratified survey PSU generation can still materialize full weight matrices, but this is documented in TODO.md:L69.
• Concrete fix: No action required for this PR.

Maintainability

Finding: None

• Severity: None
• Impact: The new helper is narrow and reusable, and docs/doc-deps.yaml:L910-L916 records its methodology documentation dependency.
• Concrete fix: None.

Tech Debt

Finding: Tracked only

• Severity: P3 informational
• Impact: Remaining performance debt is explicitly tracked in TODO.md:L69, so it is not an approval blocker under the project policy.
• Concrete fix: No action required for this PR.

Security

Finding: None

• Severity: None
• Impact: No secrets, external calls, unsafe deserialization, or security-sensitive changes found in the diff.
• Concrete fix: None.

Documentation / Tests

Finding: Survey methodology validation was removed/downgraded

• Severity: P2
• Impact: tests/test_methodology_survey.py is no longer present, and METHODOLOGY_REVIEW.md:L92 / METHODOLOGY_REVIEW.md:L1429-L1443 downgrade Survey Data Support back to In Progress. This is a methodology-validation regression, though not a runtime correctness bug.
• Concrete fix: Restore the deleted methodology survey tests and the Complete review block, or add a tracked TODO with a PR reference if this deferral is intentional.

Finding: References page no longer lists cited survey sources

• Severity: P2
• Impact: REGISTRY.md still cites Lumley (2004) and Solon-Haider-Wooldridge (2015) at docs/methodology/REGISTRY.md:L4099-L4101, L4196-L4218, and L4399-L4423, but docs/references.rst:L86-L97 now jumps from Binder/Gerber to Placebo references.
• Concrete fix: Restore the Lumley (2004), Korn-Graubard (1990), and Solon-Haider-Wooldridge (2015) entries in docs/references.rst.

Finding: Tests not executed in this review

• Severity: P3 informational
• Impact: Static review only; I did not run the test suite in this read-only environment.
• Concrete fix: Ensure CI runs tests/test_bootstrap_chunking.py under both DIFF_DIFF_BACKEND=python and DIFF_DIFF_BACKEND=rust.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 173289bd82cd987080ebf8f49e44adfa2b3c4392

---

Overall Assessment

✅ Looks good — no unmitigated P0/P1 findings.

Executive Summary

• Affected method: CallawaySantAnna multiplier bootstrap.
• The PR tiles multiplier weight generation over bootstrap draws without changing the documented estimand, weight distributions, WIF handling, or survey PSU bootstrap contract.
• Overall/event-study bootstrap still perturbs combined IFs; group aggregation still uses per-draw ATT(g,t) rows with deterministic equal-time weights.
• Prior survey documentation/test findings appear resolved in the current tree.
• I did not execute tests because the local review environment lacks numpy.

Methodology

Finding: CS bootstrap chunking is methodology-equivalent

• Severity: None
• Impact: REGISTRY.md requires multiplier bootstrap with Rademacher/Mammen/Webb weights, block structure preserving within-unit correlation, and combined IF perturbation for WIF-aware aggregations. The new loop still computes ATT(g,t) perturbations, overall combined IF bootstrap, event-study combined IF bootstrap, and group aggregation from the per-draw ATT(g,t) rows. See docs/methodology/REGISTRY.md:L493-L499, diff_diff/staggered_bootstrap.py:L415-L497.
• Concrete fix: None.

Finding: Survey PSU bootstrap contract preserved

• Severity: None
• Impact: The registry’s survey bootstrap contract is PSU-level multiplier generation with within-stratum independence and FPC scaling. The new helper tiles only the unstratified generation path and falls back to the existing full generator for stratified/lonely-PSU paths, preserving the documented survey behavior. See docs/methodology/REGISTRY.md:L4239-L4248, diff_diff/bootstrap_chunking.py:L143-L183.
• Concrete fix: None.

Code Quality

Finding: No new inference anti-patterns found

• Severity: None
• Impact: I found no new inline t_stat = effect / se or partial NaN guard pattern in the changed bootstrap code. Bootstrap inference remains routed through the existing bootstrap stats helpers. See diff_diff/staggered_bootstrap.py:L499-L600.
• Concrete fix: None.

Performance

Finding: Remaining non-CS/stratified chunking work is tracked

• Severity: P3 informational
• Impact: EfficientDiD/HAD and stratified survey PSU generation can still materialize full bootstrap weight matrices, but this is explicitly tracked in TODO.md:L69, so it is not a blocker under the deferred-work policy.
• Concrete fix: No action required for this PR.

Maintainability

Finding: Chunking helper is scoped and documented

• Severity: None
• Impact: The new helper centralizes block sizing, standard multiplier blocks, and survey PSU blocks, and docs/doc-deps.yaml:L910-L916 links it to methodology documentation.
• Concrete fix: None.

Tech Debt

Finding: Tracked only

• Severity: P3 informational
• Impact: The remaining chunking adoption gaps are tracked in TODO.md:L69.
• Concrete fix: No action required.

Security

Finding: None

• Severity: None
• Impact: No secrets, external calls, unsafe deserialization, or security-sensitive behavior found in the changed files.
• Concrete fix: None.

Documentation / Tests

Finding: Prior survey documentation/test regressions appear resolved

• Severity: None
• Impact: The current tree has corrected survey references and restored survey methodology coverage: docs/references.rst:L97-L107, METHODOLOGY_REVIEW.md:L92, and tests/test_methodology_survey.py:L1.
• Concrete fix: None.

Finding: Tests not executed in this review

• Severity: P3 informational
• Impact: Static review only; local Python import failed with ModuleNotFoundError: No module named 'numpy'.
• Concrete fix: Ensure CI runs tests/test_bootstrap_chunking.py under both DIFF_DIFF_BACKEND=python and DIFF_DIFF_BACKEND=rust.