perf(imputation,two_stage): route demeaning through shared MAP engine + fix zero-weight replicate bug (#615)

Author: igerber

CHANGELOG.md

@@ -20,6 +20,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   `survey_design=SurveyDesign(psu=<cluster_col>)`. No behavior change for unclustered fits.
 
 ### Fixed
+- **ImputationDiD/TwoStageDiD covariate fits with zero-weight replicate designs (JK1/plain
+  BRR) now produce finite SEs.** Replicate weights that zero out whole PSUs reach Step 1
+  unmasked; the previous per-estimator pandas demeaning loops divided 0/0 on zero-total-weight
+  groups, NaN-poisoning the demeaned design so EVERY replicate refit failed inside
+  `solve_ols` ("All replicate refits failed. Returning NaN variance." after a
+  non-convergence warning storm — measured: 198 warnings and SE=NaN on a 350k-row panel
+  with 16 JK1 replicates, now 0 warnings and a finite SE). A main fit combining zero-weight
+  rows with covariates previously raised an opaque `ValueError`; it now fits, with the
+  zero-weight unit's FE surfacing as NaN (spillover convention — keys retained for the
+  rank-condition check, never a silent finite 0.0) and its unidentified cells NaN across
+  all inference fields. TwoStageDiD's per-replicate "non-finite imputed outcomes" warning
+  is suppressed inside replicate-refit closures only (`warn_nan=False`); the main-fit
+  warning is unchanged.
 - **Dube, Girardi, Jordà & Taylor (2025) citation corrected to *J. Applied Econometrics*
   40(**7**):741-758** (was cited as issue 5) across `docs/references.rst`,
   `docs/methodology/REGISTRY.md`, `docs/methodology/papers/dube-2025-review.md`,
@@ -29,6 +42,25 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   original attribution as a record of what was claimed at release time; this entry
   supersedes it.
 
+### Changed
+- **ImputationDiD/TwoStageDiD demeaning modernized onto the shared MAP engine.** The
+  private per-estimator pandas `_iterative_demean` loops (rebuilt a
+  `pd.Series.groupby().transform()` hash table every alternating-projection iteration)
+  are deleted; the covariate and pre-trend-lead within-transformations now route through
+  `demean_by_groups` (factorize-once + `np.bincount`, optional Rust kernel, one dispatch
+  for all columns), and both `_iterative_fe` FE solvers route through a new shared
+  bincount Gauss-Seidel helper (`diff_diff.utils._iterative_fe_solve`, modeled on
+  SpilloverDiD's). `max_iter` modernized 100 → 10,000 (the R `fixest`/`pyfixest`
+  convention already used by the shared engine). Estimates are preserved: measured ATT
+  deltas ≤ 2e-15 and SE deltas ≤ 2e-12 relative across a 5-scenario before/after grid
+  (2.35M-row panels, R-parity suites unchanged at 1e-6/1e-7). Measured speedups (median
+  of 3): ImputationDiD covariate fit 4.18s → 1.72s (2.4x) and no-covariate 1.40x at
+  2.35M rows; replicate-weight survey variance 20.3s → 3.6s (5.7x, 32 replicates,
+  350k rows) and 30.3s → 1.8s on zeroed-PSU designs (17x — the old loop burned
+  `max_iter` futile iterations per replicate). TwoStageDiD fit time unchanged
+  (GMM-variance-dominated). Accumulation-order numerics documented in
+  `docs/methodology/REGISTRY.md` (both estimator sections + "Absorbed Fixed Effects").
+
 ## [3.6.2] - 2026-07-03
 
 ### Added

TODO.md

@@ -51,7 +51,8 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m
 | `SpilloverDiD` sparse cKDTree path for the staggered nearest-treated-distance helper (mirrors the static helper's sparse branch). `_compute_nearest_treated_distance_staggered` always builds dense `(n_units, n_treated_by_onset)` matrices per cohort; add a sparse branch gated on `n > _CONLEY_SPARSE_N_THRESHOLD`. | `spillover.py` | Wave B | Mid | Low |
 | `HeterogeneousAdoptionDiD` Phase 3 Stute: Appendix-D vectorized form replaces the per-iteration OLS refit with a single precomputed `M = I - X(X'X)^{-1}X'` applied to `eps*eta` (~2× faster, functionally identical). Shipped the literal-refit form to match paper text. | `had_pretests.py::stute_test` | Phase 3 | Mid | Low |
 | Multiplier-bootstrap weight chunking (CallawaySantAnna, EfficientDiD, and HAD — all wired through `diff_diff/bootstrap_chunking.py`) covers the **unstratified** survey-PSU generation (the default unit-level bootstrap — `cluster=None`, equivalently `cluster="unit"` — the large-`n_units` OOM case). Remaining gap: the **stratified** survey-PSU generator (`generate_survey_multiplier_weights_batch`, per-stratum + lonely-PSU pooling + FPC) still materializes the full `(n_bootstrap × n_psu)` matrix (consumed via sliced blocks). Stratified designs have few PSUs so this rarely OOMs; tile per-stratum generation over draws (each stratum's draws are independent → contiguous draw-blocks reproduce the stream bit-identically) if a large-PSU stratified design hits memory. | `diff_diff/bootstrap_chunking.py::iter_survey_multiplier_weight_blocks` | follow-up | Mid | Low |
-| `ImputationDiD._iterative_demean` (`imputation.py:1064`) and its near-identical twin `TwoStageDiD._iterative_demean` (`two_stage.py:2201`) rebuild a `pd.Series`+`groupby().transform()` every alternating-projection iteration. Precompute the unit/time group codes once and use `np.bincount` for the per-iteration group means. **Not bit-identical** on pandas 3.0 (`np.bincount` is naive accumulation; pandas' `group_mean` is Kahan-compensated → ~5.8e-11, the same order as the demean's `tol=1e-10`), so this needs a tolerance/REGISTRY-Note treatment + golden re-validation. Modest peak effect (the per-iteration Series are transient), analytical-path-only (the bootstrap reuses #562's cached projection). Optimize both twins together (cross-surface-twins). | `imputation.py`, `two_stage.py` | PR-C deferral | Mid | Low |
+| Migrate `spillover._iterative_fe_subset` onto the shared `diff_diff.utils._iterative_fe_solve` (same Gauss-Seidel-on-codes recursion, specialized to a masked Butts subsample; ImputationDiD/TwoStageDiD already route through the shared helper — this takes the FE-solver copy count from 2 to 1). Preserve the SpilloverDiD positive-weight/NaN-FE REGISTRY contract and `_FE_ITER_MAX=100` budget (or align it to 10k with a REGISTRY note). | `spillover.py` | demean-modernization | Mid | Low |
+| Adopt `snap_absorbed_regressors` (FE-spanned regressor two-stage snap + LSMR confirmation) on the ImputationDiD lead-indicator path — lead columns are the most plausible FE-spanned regressors, and the truncated-MAP-iterate exposure documented at `utils.py` applies; today only `solve_ols` rank detection guards it. Behavior change beyond the demean-modernization refactor, so deferred from that PR. | `imputation.py::_compute_lead_coefficients` | demean-modernization | Mid | Low |
 
 ### Testing / docs
 

diff_diff/imputation.py

@@ -37,7 +37,7 @@
     ImputationDiDResults,
 )
 from diff_diff.linalg import solve_ols
-from diff_diff.utils import safe_inference, warn_if_not_converged
+from diff_diff.utils import _iterative_fe_solve, demean_by_groups, safe_inference
 
 
 class _UntreatedProjection(NamedTuple):
@@ -983,22 +983,28 @@ def _iterative_fe(
         unit_vals: np.ndarray,
         time_vals: np.ndarray,
         idx: pd.Index,
-        max_iter: int = 100,
+        max_iter: int = 10_000,
         tol: float = 1e-10,
         weights: Optional[np.ndarray] = None,
     ) -> Tuple[Dict[Any, float], Dict[Any, float]]:
         """
         Estimate unit and time FE via iterative alternating projection (Gauss-Seidel).
 
-        Converges to the exact OLS solution for both balanced and unbalanced panels.
-        For balanced panels, converges in 1-2 iterations (identical to one-pass).
-        For unbalanced panels, typically 5-20 iterations.
+        Thin wrapper over the shared bincount solver
+        (``diff_diff.utils._iterative_fe_solve``): factorize unit/time once,
+        solve on integer codes, map the level arrays back to dicts.
+        Converges to the exact (W)LS solution for balanced and unbalanced
+        panels; balanced panels converge in 1-2 iterations.
 
         Parameters
         ----------
+        idx : pd.Index
+            Unused; retained for call-site stability.
         weights : np.ndarray, optional
-            Survey weights. When provided, uses weighted group means
-            (sum(w*x)/sum(w)) instead of unweighted means.
+            Survey weights (weighted group means ``sum(w*x)/sum(w)``). A
+            unit/period whose observations ALL carry zero weight has no
+            identifying contribution and gets ``NaN`` FE (its key is kept so
+            the rank-condition membership check still sees the group).
 
         Returns
         -------
@@ -1007,118 +1013,28 @@ def _iterative_fe(
         time_fe : dict
             Mapping from time -> time fixed effect.
         """
-        n = len(y)
-        alpha = np.zeros(n)  # unit FE broadcast to obs level
-        beta = np.zeros(n)  # time FE broadcast to obs level
-
-        # Precompute per-group weight sums (invariant across iterations)
-        if weights is not None:
-            w_series = pd.Series(weights, index=idx)
-            wsum_t = w_series.groupby(time_vals).transform("sum").values
-            wsum_u = w_series.groupby(unit_vals).transform("sum").values
-
-        converged = False
-        with np.errstate(invalid="ignore", divide="ignore"):
-            for iteration in range(max_iter):
-                resid_after_alpha = y - alpha
-                if weights is not None:
-                    wr_t = pd.Series(resid_after_alpha * weights, index=idx)
-                    beta_new = wr_t.groupby(time_vals).transform("sum").values / wsum_t
-                else:
-                    beta_new = (
-                        pd.Series(resid_after_alpha, index=idx)
-                        .groupby(time_vals)
-                        .transform("mean")
-                        .values
-                    )
-
-                resid_after_beta = y - beta_new
-                if weights is not None:
-                    wr_u = pd.Series(resid_after_beta * weights, index=idx)
-                    alpha_new = wr_u.groupby(unit_vals).transform("sum").values / wsum_u
-                else:
-                    alpha_new = (
-                        pd.Series(resid_after_beta, index=idx)
-                        .groupby(unit_vals)
-                        .transform("mean")
-                        .values
-                    )
-
-                # Check convergence on FE changes
-                max_change = max(
-                    np.max(np.abs(alpha_new - alpha)),
-                    np.max(np.abs(beta_new - beta)),
-                )
-                alpha = alpha_new
-                beta = beta_new
-                if max_change < tol:
-                    converged = True
-                    break
-        warn_if_not_converged(converged, "ImputationDiD iterative FE solver", max_iter, tol)
-
-        unit_fe = pd.Series(alpha, index=idx).groupby(unit_vals).first().to_dict()
-        time_fe = pd.Series(beta, index=idx).groupby(time_vals).first().to_dict()
+        unit_codes, unit_uniques = pd.factorize(unit_vals, sort=False)
+        time_codes, time_uniques = pd.factorize(time_vals, sort=False)
+        if (unit_codes < 0).any() or (time_codes < 0).any():
+            raise ValueError(
+                "ImputationDiD: unit or time column contains NaN. Drop or "
+                "impute missing group keys before fitting."
+            )
+        unit_fe_arr, time_fe_arr = _iterative_fe_solve(
+            np.asarray(y, dtype=np.float64),
+            unit_codes.astype(np.intp, copy=False),
+            time_codes.astype(np.intp, copy=False),
+            len(unit_uniques),
+            len(time_uniques),
+            weights=weights,
+            max_iter=max_iter,
+            tol=tol,
+            method_name="ImputationDiD iterative FE solver",
+        )
+        unit_fe = dict(zip(unit_uniques, unit_fe_arr))
+        time_fe = dict(zip(time_uniques, time_fe_arr))
         return unit_fe, time_fe
 
-    @staticmethod
-    def _iterative_demean(
-        vals: np.ndarray,
-        unit_vals: np.ndarray,
-        time_vals: np.ndarray,
-        idx: pd.Index,
-        max_iter: int = 100,
-        tol: float = 1e-10,
-        weights: Optional[np.ndarray] = None,
-    ) -> np.ndarray:
-        """Demean a vector by iterative alternating projection (unit + time FE removal).
-
-        Converges to the exact within-transformation for both balanced and
-        unbalanced panels. For balanced panels, converges in 1-2 iterations.
-
-        Parameters
-        ----------
-        weights : np.ndarray, optional
-            Survey weights. When provided, uses weighted group means
-            (sum(w*x)/sum(w)) instead of unweighted means.
-        """
-        result = vals.copy()
-
-        # Precompute per-group weight sums (invariant across iterations)
-        if weights is not None:
-            w_series = pd.Series(weights, index=idx)
-            wsum_t = w_series.groupby(time_vals).transform("sum").values
-            wsum_u = w_series.groupby(unit_vals).transform("sum").values
-
-        converged = False
-        with np.errstate(invalid="ignore", divide="ignore"):
-            for _ in range(max_iter):
-                if weights is not None:
-                    wr_t = pd.Series(result * weights, index=idx)
-                    time_means = wr_t.groupby(time_vals).transform("sum").values / wsum_t
-                else:
-                    time_means = (
-                        pd.Series(result, index=idx).groupby(time_vals).transform("mean").values
-                    )
-                result_after_time = result - time_means
-                if weights is not None:
-                    wr_u = pd.Series(result_after_time * weights, index=idx)
-                    unit_means = wr_u.groupby(unit_vals).transform("sum").values / wsum_u
-                else:
-                    unit_means = (
-                        pd.Series(result_after_time, index=idx)
-                        .groupby(unit_vals)
-                        .transform("mean")
-                        .values
-                    )
-                result_new = result_after_time - unit_means
-                if np.max(np.abs(result_new - result)) < tol:
-                    result = result_new
-                    converged = True
-                    break
-                result = result_new
-        warn_if_not_converged(converged, "ImputationDiD iterative demean", max_iter, tol)
-        return result
-
     @staticmethod
     def _compute_balanced_cohort_mask(
         df_treated: pd.DataFrame,
@@ -1240,16 +1156,24 @@ def _fit_untreated_model(
             X_raw = df_0[covariates].values.copy()
             units = df_0[unit].values
             times = df_0[time].values
-            n_cov = len(covariates)
-
-            # Step A: Iteratively demean Y and all X columns to remove unit+time FE
-            y_dm = self._iterative_demean(y, units, times, df_0.index, weights=w_0)
-            X_dm = np.column_stack(
-                [
-                    self._iterative_demean(X_raw[:, j], units, times, df_0.index, weights=w_0)
-                    for j in range(n_cov)
-                ]
+
+            # Step A: within-transform Y and all X columns through the shared
+            # MAP engine (factorize-once + bincount + optional Rust kernel),
+            # one dispatch for every column. within_transform pins
+            # [unit, time]; [time, unit] here preserves the historical
+            # time-then-unit sweep order of the per-estimator loops.
+            narrow = df_0[[outcome, *covariates, time, unit]].copy()
+            demeaned, _ = demean_by_groups(
+                narrow,
+                [outcome, *covariates],
+                [time, unit],
+                inplace=True,
+                weights=w_0,
+                max_iter=10_000,
+                tol=1e-10,
             )
+            y_dm = demeaned[outcome].to_numpy(dtype=np.float64)
+            X_dm = demeaned[covariates].to_numpy(dtype=np.float64)
 
             # Step B: OLS for covariate coefficients on demeaned data
             result = solve_ols(
@@ -2283,31 +2207,27 @@ def _compute_lead_coefficients(
             df_0[col_name] = indicator
             lead_cols.append(col_name)
 
-        # Within-transform via iterative demeaning (survey-weighted when present)
-        y_dm = self._iterative_demean(
-            df_0[outcome].values,
-            df_0[unit].values,
-            df_0[time].values,
-            df_0.index,
-            weights=survey_weights_0,
-        )
-
         all_x_cols = lead_cols[:]
         if covariates:
             all_x_cols.extend(covariates)
 
-        X_dm = np.column_stack(
-            [
-                self._iterative_demean(
-                    df_0[col].values,
-                    df_0[unit].values,
-                    df_0[time].values,
-                    df_0.index,
-                    weights=survey_weights_0,
-                )
-                for col in all_x_cols
-            ]
+        # Within-transform through the shared MAP engine (survey-weighted when
+        # present), one dispatch for outcome + leads + covariates. Demean into
+        # a narrow copy: df_0's raw lead indicators must survive for the
+        # per-horizon n_obs counts below. within_transform pins [unit, time];
+        # [time, unit] here preserves the historical time-then-unit sweep order.
+        narrow = df_0[[outcome, *all_x_cols, time, unit]].copy()
+        demeaned, _ = demean_by_groups(
+            narrow,
+            [outcome, *all_x_cols],
+            [time, unit],
+            inplace=True,
+            weights=survey_weights_0,
+            max_iter=10_000,
+            tol=1e-10,
         )
+        y_dm = demeaned[outcome].to_numpy(dtype=np.float64)
+        X_dm = demeaned[all_x_cols].to_numpy(dtype=np.float64)
 
         # OLS for point estimates + VCV. When survey VCV will replace the
         # cluster-robust VCV, skip cluster_ids to avoid errors on domains

diff_diff/spillover.py

@@ -1287,7 +1287,8 @@ def _convert_treatment_to_first_treat(
 # =============================================================================
 
 # Convergence tolerance for the iterative alternating-projection FE solver
-# (Gauss-Seidel style; mirrors `TwoStageDiD._iterative_fe`).
+# (Gauss-Seidel style; same recursion as the shared
+# `diff_diff.utils._iterative_fe_solve` used by ImputationDiD/TwoStageDiD).
 _FE_ITER_MAX = 100
 _FE_ITER_TOL = 1e-10
 
@@ -1408,8 +1409,9 @@ def _iterative_fe_subset(
     ``NaN`` at positions whose unit / time is not represented in the
     subsample (rank-deficient cells).
 
-    Mirrors ``TwoStageDiD._iterative_fe`` structurally but operates on
-    integer-coded factors via ``np.bincount`` for speed.
+    Same Gauss-Seidel-on-integer-codes recursion as the shared
+    ``diff_diff.utils._iterative_fe_solve`` (which ImputationDiD/TwoStageDiD
+    now route through), specialized here to a masked Butts subsample.
 
     **Wave E.1 weighted path** — when ``weights`` is supplied, the solver
     minimizes ``sum_i w_i * (y_i - mu_i - lambda_t)^2`` (WLS-FE under