perf(rust): compiled MAP demeaning kernel - rayon across columns, 1.4-2.8x on FE-absorption workloads

[igerber]

Summary

• Adds the optional Rust demean_map kernel for FE-absorption demeaning (rust/src/demean.rs): the method-of-alternating-projections sweeps run compiled and rayon-parallel across the demeaned variables, with the GIL released (Python::detach) around the compute. No BLAS dependency - works on the plain Windows build.
• Exact mirror of the canonical numpy engine (python-canonical policy): the numpy MAP loops are extracted into diff_diff.utils._demean_map_numpy (the directly-testable reference), and the kernel reproduces its per-column independent convergence loops, dimension sweep order, row-order scatter-add accumulation (np.bincount parity), per-group-sum division, zero-total-weight inert guard, and max|x - x_old| < tol stopping rule including NaN-poisoning semantics. Equivalence tests assert iteration-count EQUALITY then assert_allclose(atol=1e-12); observed outputs are bit-identical on all fixtures (asserted at allclose per policy - never a bit-identity claim).
• Resilient wiring: demean_map is imported independently in _backend.py, so a stale/mixed-version extension missing only the new symbol degrades to the numpy engine without disabling older Rust accelerations; the wrapper honors a documented None-fallback contract (absent symbol, degenerate shapes, or kernel-side validation errors) and re-raises anything else. DIFF_DIFF_BACKEND=python disables as usual. Everything downstream (non-convergence warning contract, FE-spanned-regressor snap + LSMR confirmation, one-way delegation, n_effects) is unchanged.
• Measured on frozen code (fresh same-session numpy-backend denominator, strictly sequential, CVs <= 4.1%; committed baselines/fe_absorption_rust.json): firm_churn (2.4M-row SunAbraham) 48.98s -> 25.84s = 1.90x (CV-adjusted lower bound 1.78x - clears the pre-agreed >= 1.5x submit gate), correlated-FE stress 31.74s -> 11.37s (2.79x), 5M-row geo experiment 0.976s -> 0.484s (2.02x, now ahead of the pyfixest yardstick's 0.623s), scanner TWFE 1.59x, survey BRR 1.44x, county event study 1.36x, small-panel guard unregressed. Estimate identity vs the committed baselines: 1e-13-1e-16 on every scenario including the stress case.
• Disclosed trade-off: the kernel holds one owned input copy plus the result, so firm_churn peak RSS rises 13.4 GB -> 21.0 GB; chunking variables through the kernel is the recorded follow-up if a real workload hits it. The now-solver-bound county-class profile is promoted to an Actionable TODO row (supersedes a stale parked note with fresh attribution).

Methodology references (required if estimator / math changes)

• Method name(s): Method of alternating projections (MAP) for N-way fixed-effects absorption (FWL residualization) - implementation acceleration only, no estimand/inference change
• Paper / source link(s): fixest/reghdfe/lfe demeaning convention, documented in docs/methodology/REGISTRY.md "Absorbed Fixed Effects with Survey Weights" (new Note covers the Rust kernel contract)
• Any intentional deviations from the source (and why): None - numpy remains the canonical reference; the kernel is documented as required to mirror it, with equivalence enforced by tests (iteration-count equality + allclose at 1e-12)

Validation

• Tests added/updated: tests/test_rust_backend.py::TestDemeanMapKernel (16 cases: weighted/unweighted x 2-way/3-way x balanced/unbalanced/contiguous >100-iteration regimes, zero-total-weight inertness, NaN-in-variable non-convergence parity, k=1/k=64, non-convergence flag parity, forced-fallback and stale-symbol fallback both proven to run the numpy engine, warning parity, estimator-level ATT/SE + FE-spanned snap-decision parity); cargo unit tests in rust/src/demean.rs. Full targeted sweeps: 960 passed with the Rust build, 875 passed under DIFF_DIFF_BACKEND=python.
• Backtest / simulation / notebook evidence (if applicable): committed baselines/fe_absorption_rust.json (multi-run protocol, CVs recorded); findings table in docs/performance-plan.md gains Rust columns via gen_findings_tables.py; gate evaluation documented in the performance-plan narrative.

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 method: MAP fixed-effect absorption in demean_by_groups / within_transform, with downstream impact on TWFE, SunAbraham, Bacon, Wooldridge, DiD/MPD absorb=, and survey replicate refits.
• Methodology is documented in docs/methodology/REGISTRY.md:L4301-L4353; the Rust path is explicitly a numpy-canonical implementation acceleration, not an estimand or inference change.
• Rust and numpy paths preserve the sensitive contracts I checked: sweep order, per-column convergence, zero-weight group guard, NaN non-convergence, fallback behavior, and warning propagation.
• No new inline inference, SE, p-value, CI, weighting, or control-group logic issue was introduced.
• One P3 performance trade-off remains: higher peak RSS from marshalling/copying, already disclosed in docs/performance-plan.md:L69-L81.

Methodology

• Finding: No methodology defect.
  Severity: N/A
  Impact: The PR affects absorbed fixed-effects MAP demeaning, but the registry documents the Rust kernel as an exact mirror of the canonical numpy engine, with numpy remaining the reference implementation. The implementation matches that contract in diff_diff/utils.py:L2645-L2747 and rust/src/demean.rs:L24-L75.
  Concrete fix: None.

Code Quality

• Finding: No blocking code-quality issue found.
  Severity: N/A
  Impact: Backend import/fallback is scoped so a stale extension missing only demean_map falls back without disabling older Rust functions (diff_diff/_backend.py:L68-L85). Kernel validation errors route back to numpy, and public validation-sensitive cases still fail through the canonical path if inputs are invalid.
  Concrete fix: None.

Performance

• Finding: Rust path has a documented peak-memory trade-off.
  Severity: P3 informational
  Impact: _demean_map_rust builds a contiguous Python matrix before calling Rust (diff_diff/utils.py:L2729-L2737), and Rust owns working/output buffers (rust/src/demean.rs:L183-L219). This can raise peak RSS on large FE workloads; the performance plan documents firm_churn RSS increasing from 13.4 GB to 21.0 GB and names variable chunking as the follow-up (docs/performance-plan.md:L79-L81).
  Concrete fix: No action required for approval. If this becomes user-visible, chunk variables through the Rust kernel.

Maintainability

• Finding: No maintainability blocker.
  Severity: N/A
  Impact: Extracting _demean_map_numpy gives the Rust implementation a directly testable reference surface, which reduces drift risk for a methodology-sensitive routine.
  Concrete fix: None.

Tech Debt

• Finding: Solver-bound follow-up is tracked.
  Severity: P3 informational
  Impact: The PR adds an actionable backlog item for post-demeaning solver bottlenecks in SunAbraham-family fits (TODO.md:L48-L50). This is not a correctness blocker.
  Concrete fix: None for this PR.

Security

• Finding: No obvious secret or security issue found.
  Severity: N/A
  Impact: Changed files do not appear to introduce credentials, external command execution, or unsafe user-input handling.
  Concrete fix: None.

Documentation/Tests

• Finding: No blocking documentation/test gap.
  Severity: N/A
  Impact: Registry, changelog, performance docs, and doc dependency metadata were updated. Tests cover weighted/unweighted two-way and three-way parity, slow convergence, zero-total-weight groups, NaN non-convergence, stale-symbol fallback, warning parity, and estimator-level ATT/SE parity (tests/test_rust_backend.py:L2918-L3159).
  Concrete fix: None.

Review note: I did not run the full test suite in this read-only review environment; static inspection included git diff --check, which passed.