feat(estimators): iterative alternating-projection demeaning for N-way absorbed FE

[igerber]

Summary

• Fixes an unbalanced-panel correctness bug: N>1 absorbed fixed effects (absorb=[a, b, ...] in DifferenceInDifferences / MultiPeriodDiD) and the shared unweighted two-way within_transform (used by TwoWayFixedEffects, SunAbraham, BaconDecomposition) used single-pass sequential demeaning, which is the exact (weighted) Frisch-Waugh-Lovell residualization only when the FE subspaces are orthogonal (balanced fully-crossed panels). On unbalanced panels it was a biased approximation (coefficients off by ~1e-2 in tested cases).
• New N-way method-of-alternating-projections engine diff_diff.utils.demean_by_groups(): each variable is demeaned by each FE dimension in turn until convergence — the exact (weighted) FWL residual onto the combined column space of all absorbed dummies, matching R fixest / reghdfe / lfe. The two-way within_transform() now delegates to it.
• The four DiD/MPD absorb= call sites (analytical + replicate-refit, ×2 estimators) route through the engine; the weighted-multi-absorb ValueError rejection is lifted (now supported via weighted MAP).
• Byte-stability preserved: single-absorb delegates to the unchanged one-way demean_by_group; the weighted within_transform output is bit-identical (Wooldridge golden at atol=1e-14 stays green); balanced multi-way matches the prior closed-form demean to machine precision. The unweighted two-way path now also emits the non-convergence UserWarning (previously only the weighted path could).
• DOF accounting (n_absorbed_effects = sum_d(nunique_d - 1)) is preserved unchanged; multi-way-FE DOF correctness is out of scope.

Methodology references (required if estimator / math changes)

• Method name(s): Frisch-Waugh-Lovell within transformation for multi-way fixed effects via the method of alternating projections (von Neumann / Halperin; Guimarães-Portugal; Gaure lfe; Correia reghdfe).
• Paper / source link(s): documented in docs/methodology/REGISTRY.md (TwoWayFixedEffects within-transform Note; "Absorbed Fixed Effects with Survey Weights" Note).
• Any intentional deviations from the source (and why): None. This brings diff-diff into agreement with fixest/reghdfe/lfe for N>1 absorbed FE on unbalanced panels (previously it deviated via single-pass).

Validation

• Tests added/updated: tests/test_utils.py (new TestDemeanByGroups: engine vs full-dummy OLS for unbalanced 2-way and 3-way, len==1 byte-identity, weighted byte-identity guard vs a frozen copy of the old loop, orthogonality, non-convergence warning); tests/test_within_transform.py (relaxed 5 balanced closed-form assertions to assert_allclose ~1e-12, added unbalanced-correctness + convergence-warning tests); tests/test_methodology_did.py (new TestMultiAbsorbIterativeDemean: unbalanced unweighted and non-uniform-weighted DiD/MPD absorb= vs fixed_effects= parity at atol≈1e-8, FE-collinear regressor reports NaN coef); tests/test_survey.py (flipped the two multi-absorb+survey rejection tests to now-supported + full-dummy parity).
• R-parity goldens unchanged: test_methodology_twfe.py / test_methodology_sun_abraham.py / test_methodology_bacon.py / test_methodology_wooldridge.py all green (balanced panels: MAP == closed-form to ~1 ULP; weighted byte-identical).
• Backtest / simulation / notebook evidence (if applicable): N/A.

Security / privacy

• Confirm no secrets/PII in this PR: Yes

[github-actions]

PR Review: ⚠️ Needs changes

Overall Assessment

⚠️ Needs changes — one unmitigated P1 edge-case/methodology-adjacent issue remains in the new weighted multi-absorb path.

Executive Summary

• MAP demeaning is documented in REGISTRY.md and is the right methodological direction for unbalanced multi-way FE.
• The main coefficient path now matches the registry intent for nonzero-weight panels.
• P1: weighted MAP can divide by zero when an absorbed group has zero total survey weight, despite zero-weight survey/domain rows being supported by the registry.
• Replicate refits correctly drop zero-weight rows, but the full-sample weighted absorb path does not.
• Minor docs issue: within_transform() still describes the old additive two-way formula as if it were the actual computation.

Methodology

Finding 1 — P1: Zero-total-weight absorbed groups break weighted MAP

Location: diff_diff/utils.py:L2726-L2733, called from diff_diff/estimators.py:L460-L466 and diff_diff/estimators.py:L1612-L1618

Impact:
demean_by_groups() computes weighted group means as wx_sum / w_sum without guarding w_sum == 0. Survey weights explicitly allow zero-weight observations, and the registry documents subpopulation/domain support with non-negative weights and zero-weight padding: docs/methodology/REGISTRY.md:L4529-L4535.

This PR newly supports DifferenceInDifferences / MultiPeriodDiD with weighted multi-absorb. If a survey subpopulation or padded domain leaves an absorbed category with only zero-weight rows, the new weighted MAP path produces NaN/Inf during residualization before WLS can treat those rows as inert. That can turn a valid survey/domain fit into a production error or unstable downstream inference.

The replicate refit closures do drop w_r > 0 before demeaning, consistent with the registry note at docs/methodology/REGISTRY.md:L4549-L4555, but the full-sample absorb calls pass the full survey_weights directly.

Concrete fix:
Make zero-weight rows/groups inert in the weighted MAP path. For example, compute weighted means with guarded division:

means = np.divide(wx_sum, w_sum, out=np.zeros_like(wx_sum, dtype=float), where=w_sum > 0)
return x - means

Then add tests for both DifferenceInDifferences and MultiPeriodDiD where survey_design has zero-weight rows and at least one absorbed category has zero total weight. Assert absorb=[a, b] matches fixed_effects=[a, b] on the positive-weight effective sample and does not emit NaN/Inf from demeaning.

Code Quality

No additional code-quality blockers found. The new helper is reasonably scoped and centralizes the MAP implementation.

Performance

No blocking performance finding. The per-variable pandas groupby loop is heavier than the prior closed-form balanced transform, but that is a consequence of moving to exact MAP on unbalanced panels.

Maintainability

No maintainability blocker found. Delegating within_transform() to demean_by_groups() reduces duplicated demeaning logic.

Tech Debt

No untracked blocking tech debt found. The removed TODO entry for weighted multi-absorb is appropriate once the zero-weight edge case is handled.

Security

No security findings.

Documentation/Tests

Finding 2 — P2: within_transform() docstring still describes the old closed-form transform

Location: diff_diff/utils.py:L2799-L2805

Impact:
The docstring says within_transform() computes y_it - y_i. - y_.t + y_.., which is exactly the old balanced-panel formula. The implementation now uses MAP, and the registry correctly says the additive formula is only exact on balanced fully-crossed panels.

Concrete fix:
Update the docstring to say the function applies two-way MAP demeaning, and mention the additive formula only as the balanced-panel special case.

Path to Approval

1. Guard weighted MAP division by zero in demean_by_groups() so zero-total-weight absorbed groups cannot produce NaN/Inf.
2. Add DiD and MPD regression tests with zero-weight survey rows where one absorbed group has zero total weight, comparing absorb=[a, b] to fixed_effects=[a, b].
3. Update the within_transform() docstring to describe MAP rather than the old additive formula as the general computation.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 7af6c6eb4450ca2c37843632dbcf30e8eb9f60b3

---

PR Review: ✅ Looks good

Overall Assessment

✅ Looks good — no unmitigated P0 or P1 findings. The previous P1 zero-total-weight absorbed-group issue and previous P2 within_transform() docstring issue are addressed.

Executive Summary

• The affected methods are DifferenceInDifferences(absorb=...), MultiPeriodDiD(absorb=...), and shared two-way within_transform() consumers.
• The PR now aligns the implementation with the Methodology Registry’s MAP demeaning notes for unbalanced multi-way absorbed FE.
• The prior zero-weight survey/domain edge case is guarded with np.divide(..., where=w_sum > 0) in both one-way and N-way weighted demeaning.
• Replicate absorb refits still drop zero-weight rows before re-demeaning, consistent with the registry.
• Only one P3 informational cleanup remains: TODO.md still contains a stale cross-reference saying N>1 absorbed-FE + weights is gated.

Methodology

No methodology blockers found.

The new MAP path is documented in docs/methodology/REGISTRY.md:L4270-L4292 and implemented in the DiD/MPD absorb paths at diff_diff/estimators.py:L449-L468 and diff_diff/estimators.py:L1604-L1620. This matches the documented FWL residualization requirement for multiple absorbed dimensions.

The previous zero-weight-group concern is resolved by guarded weighted means in diff_diff/utils.py:L2621-L2632 and diff_diff/utils.py:L2736-L2745, with DiD/MPD zero-total-weight absorbed-category tests at tests/test_methodology_did.py:L1685-L1704 and tests/test_methodology_did.py:L1706-L1738.

Code Quality

No findings.

demean_by_groups() centralizes the MAP implementation and preserves one-way delegation to demean_by_group() at diff_diff/utils.py:L2709-L2719.

Performance

No findings.

The iterative groupby loop is more expensive than the old closed-form unweighted two-way transform, but that cost follows from switching to exact MAP on unbalanced panels.

Maintainability

No findings.

The within_transform() docstring now describes MAP and identifies the additive formula as the balanced-panel special case at diff_diff/utils.py:L2810-L2856.

Tech Debt

Finding 1 — P3: Stale TODO cross-reference still says weighted N>1 absorb is gated

Location: TODO.md:L220-L226

Impact:
TODO.md still says “N>1 absorbed-FE + weights composition remains gated by the open multi-absorb row,” but this PR removes that gate and the open row. This is documentation/tracker drift only; runtime behavior and methodology are unaffected.

Concrete fix:
Update or remove that parenthetical so TODO.md no longer points to a removed Actionable item.

Security

No findings.

Documentation/Tests

No blocking findings.

The registry and docstring updates cover the methodology change. Focused test execution could not be run in this environment because pytest is not installed (python -m pytest: No module named pytest).

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 91c6381d0d55edafba40251179a83f02bacaff4b

---

PR Review: ✅ Looks good

Overall Assessment

✅ Looks good — no unmitigated P0 or P1 findings. The prior weighted multi-absorb zero-weight concern is addressed, and the stale TODO cross-reference from the previous review is corrected.

Executive Summary

• Affected methods: DifferenceInDifferences(absorb=...), MultiPeriodDiD(absorb=...), and shared within_transform() users.
• The new MAP demeaning path matches the Methodology Registry’s documented FWL residualization requirement for unbalanced multi-way FE.
• Weighted multi-absorb no longer rejects and now routes through demean_by_groups().
• Zero-total-weight absorbed groups are guarded in both one-way and N-way weighted demeaning.
• No new inline inference or partial NaN-guard anti-patterns were introduced.
• I could not run tests in this sandbox because pandas is not installed.

Methodology

No findings.

The methodology change is documented in docs/methodology/REGISTRY.md:L4270-L4293 and docs/methodology/REGISTRY.md:L350-L356. The implementation routes DiD and MPD absorbed regressors/outcomes through demean_by_groups() at diff_diff/estimators.py:L449-L468 and diff_diff/estimators.py:L1594-L1619, including replicate refits at diff_diff/estimators.py:L633-L641 and diff_diff/estimators.py:L1827-L1843.

Code Quality

No findings.

The new helper centralizes MAP behavior at diff_diff/utils.py:L2644-L2796, delegates one-way demeaning back to demean_by_group() at diff_diff/utils.py:L2709-L2719, and keeps within_transform() as a thin wrapper at diff_diff/utils.py:L2883-L2892.

Performance

No findings.

The iterative MAP loop is more expensive than the old closed-form unweighted two-way transform, but that follows from switching to exact residualization on unbalanced panels.

Maintainability

No findings.

The prior stale TODO concern is resolved: the old actionable multi-absorb row is gone, and the SE consistency note now says N>1 absorbed FE with weights is supported via iterative demeaning at TODO.md:L220-L225.

Tech Debt

No findings.

Security

No findings.

Documentation/Tests

No blocking findings.

Coverage includes unbalanced full-dummy parity and 3-way MAP checks in tests/test_utils.py:L1502-L1534, DiD/MPD absorb-vs-full-dummy parity in tests/test_methodology_did.py:L1589-L1669, zero-total-weight group guards in tests/test_methodology_did.py:L1671-L1745, and the lifted survey guard in tests/test_survey.py:L2647-L2728. The within_transform() docstring now reflects the MAP behavior at diff_diff/utils.py:L2813-L2856.