perf(imputation,two_stage): shared MAP demeaning engine + zero-weight replicate fix

[igerber]

Summary

• Delete the executably-identical private pandas _iterative_demean twins in ImputationDiD/TwoStageDiD; the covariate and pre-trend-lead within-transformations now route through the shared MAP engine demean_by_groups (factorize-once + np.bincount, optional Rust kernel, one dispatch for all columns; group order [time, unit] preserves the historical time-then-unit sweep).
• Replace both _iterative_fe bodies with a shared bincount Gauss-Seidel helper (diff_diff.utils._iterative_fe_solve, modeled on spillover._iterative_fe_subset); estimator methods stay as thin wrappers with unchanged signatures. max_iter modernized 100 → 10,000 (the shared-engine / R fixest convention).
• Bug fix: covariate fits with zero-weight replicate designs (JK1/plain BRR zero whole PSUs; keep_mask only drops always-treated) previously divided 0/0 in the pandas loops → NaN-poisoned demeaned design → every replicate refit failed in solve_ols(check_finite=True) → NaN SEs plus a non-convergence warning storm (measured: 198 warnings, SE=NaN). Now: finite SE, zero warnings. Zero-total-weight groups surface as NaN FE with keys retained (SpilloverDiD REGISTRY convention — never a silent finite 0.0). Main fits with zero-weight rows + covariates no longer raise an opaque ValueError.
• TwoStageDiD's per-replicate "non-finite imputed outcomes" warning is suppressed inside replicate-refit closures only (warn_nan=False threaded through the three stage-2 helpers); the main-fit warning is unchanged.
• Measured (median of 3, frozen code): ImputationDiD covariate fit 2.4x, no-covariate 1.4x at 2.35M rows; replicate-weight survey variance 5.7x (32 positive replicates) and 17x (zeroed-PSU JK1). TwoStageDiD unchanged (GMM-variance-dominated). ATT deltas ≤ 2e-15, SE deltas ≤ 1.5e-12 relative across all five scenarios.

Methodology references (required if estimator / math changes)

• Method name(s): Borusyak-Jaravel-Spiess (2024) imputation estimator Step 1; Gardner (2022) two-stage DiD Stage 1. No estimand or equation changes — same Gauss-Seidel/MAP WLS projections on integer codes instead of pandas groupby loops.
• Paper / source link(s): BJS 2024 REStud 91(6); Gardner 2022 arXiv:2207.05943 (see docs/methodology/REGISTRY.md sections, Notes updated in this PR).
• Any intentional deviations from the source (and why): None methodological. Accumulation-order numerics (~1e-10 order, bincount vs pandas Kahan-compensated means) and the zero-weight NaN-FE convention are documented with **Note:** labels in both REGISTRY sections, cross-linked to the "Absorbed Fixed Effects with Survey Weights" engine notes.

Validation

• Tests added/updated: tests/test_utils.py (new TestIterativeFeSolve: frozen-pandas-loop drift guard at atol=1e-9, zero-weight NaN-FE contract, warn label); tests/test_imputation.py + tests/test_two_stage.py (new TestZeroWeightGroups estimator-level locks: finite SE on zeroed-PSU replicate designs, no warning storm, main-fit nan-ytilde warning unchanged, NaN cohort cell via assert_nan_inference; covariate-delta-vs-full-dummy-OLS lock; _iterative_fe zero-weight NaN tests). Direct tests of the deleted _iterative_demean retired with an intentional-coverage-narrowing note (engine-level non-convergence coverage lives in test_utils.py::TestDemeanByGroups).
• Backtest / simulation / notebook evidence (if applicable): R-parity suites (test_methodology_imputation.py vs didimputation, test_methodology_two_stage.py vs did2s) pass unchanged at ATT abs=1e-6 / SE abs=1e-7 under both DIFF_DIFF_BACKEND=python and =rust; full targeted sweeps (439 tests) green under both backends; 5-scenario before/after benchmark grid on 2.35M-row panels.

Security / privacy

• Confirm no secrets/PII in this PR: Yes

---

Generated with Claude Code

[github-actions]

Overall Assessment

✅ Looks good — no unmitigated P0/P1 findings.

Executive Summary

• Affected methods: BJS ImputationDiD Step 1 / pretrend within-transform, and Gardner TwoStageDiD Stage 1 / covariate within-transform.
• The shared MAP/bincount migration, max_iter=10_000, accumulation-order drift, and zero-weight NaN-FE convention are documented in docs/methodology/REGISTRY.md.
• I found no undocumented methodology mismatch, incorrect variance/SE change, control-group issue, or new inline inference anti-pattern.
• Zero-weight replicate behavior is covered in both estimator tests; warn_nan=False is only threaded inside TwoStageDiD replicate refits.
• I could not run tests here because pytest is not installed.

Methodology

P3 — Documented Implementation Deviation / No Action Required

Impact: ImputationDiD and TwoStageDiD now use shared demean_by_groups() / _iterative_fe_solve() instead of private pandas loops. This changes numerical accumulation order and zero-total-weight group behavior, but the deviations are explicitly documented with **Note:** entries and preserve the cited estimands and variance contracts.

Locations: diff_diff/imputation.py:L980-L1036, diff_diff/imputation.py:L1166-L1176, diff_diff/imputation.py:L2220-L2230, diff_diff/two_stage.py:L2167-L2179, diff_diff/two_stage.py:L2229-L2239, diff_diff/utils.py:L2980-L3089, docs/methodology/REGISTRY.md:L1413-L1414, docs/methodology/REGISTRY.md:L1498-L1499

Concrete fix: None required.

Code Quality

No findings.

Performance

No findings. The refactor removes duplicated pandas groupby loops and uses the shared factorize-once MAP path.

Maintainability

No findings.

Tech Debt

P3 — Remaining Solver Consolidation / Tracked in TODO.md

Impact: spillover._iterative_fe_subset remains a sibling implementation rather than using _iterative_fe_solve, and ImputationDiD lead-indicator snapping is deferred. Both are explicitly tracked, so they are not blockers.

Locations: TODO.md:L54-L55

Concrete fix: No action required for this PR; address the tracked TODO items in follow-up work.

Security

No findings.

Documentation/Tests

No findings. The PR updates methodology notes and adds focused tests for shared FE drift, zero-weight NaN FE, replicate finite SEs, and TwoStageDiD warning suppression.

Verification not run: pytest is unavailable in this environment (pytest: command not found).