igerber · igerber · Jul 7, 2026
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -454,6 +454,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   memory-traffic cleanup, not a headline speedup). The `zero_mask` abs scan is retained
   (correctness); the remaining conditional-path lever is the sieve/nuisance stage
   (see TODO).
+### Changed
+- **`SpilloverDiD` stage-1 FE solver routed through the shared Gauss-Seidel engine.**
+  `spillover._iterative_fe_subset` is now a thin Butts-subsample wrapper over
+  `diff_diff.utils._iterative_fe_solve` (the engine ImputationDiD / TwoStageDiD already
+  use), taking the FE-solver copy count in the library from 2 to 1. The wrapper keeps the
+  SpilloverDiD front door (empty-Omega_0 / empty positive-weight-Omega_0 `ValueError`
+  gates); the shared engine owns the iteration, the zero-weight NaN-FE convention, and
+  the `warn_if_not_converged` non-convergence warning (now labelled "SpilloverDiD stage-1
+  FE (Butts Omega_0 subsample)", replacing the caller-side message). Per sweep the shared
+  engine computes the identical group means and convergence metric, so converged fits are
+  bit-identical; `max_iter` is aligned from the historical local cap of 100 to the shared
+  10,000 convention (fits that previously exhausted 100 iterations and warned may now
+  converge instead — strictly more accurate FE; `tol=1e-10` unchanged). REGISTRY
+  SpilloverDiD section documents the routing.
 
 ## [3.6.2] - 2026-07-03
 

diff --git a/TODO.md b/TODO.md
@@ -52,7 +52,6 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m
 | Wild cluster bootstrap CI inversion calls `_t_star(r)` ~O(100) times, each materializing a fresh `(B×n)` `y_star` + `(k×B)` refit + `(n×B)` residual arrays. Acceptable for the few-cluster regime; for large-`n`/large-`B`, chunk `_t_star` over draws or precompute the `r`-independent cluster-level pieces (restricted residuals are linear in `r`). | `utils.py::wild_bootstrap_se._t_star` | #543 | Mid | Low |
 | `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 |
-| 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 |
 | Per-cell `treated_units`/`control_units` label arrays (`all_units[positions]`, ~O(n_control) alloc per (g,t) cell) are consumed only by the precomputed-None fallback of the combined-IF assembly, which no in-package caller reaches — build them lazily (or drop from the IF-info dict) to cut per-cell allocation at high cell counts. | `staggered.py::_compute_att_gt_fast` | CS-scaling | Mid | Low |
 

diff --git a/diff_diff/spillover.py b/diff_diff/spillover.py
@@ -41,7 +41,7 @@
 from diff_diff.linalg import _rank_guarded_inv, solve_ols
 from diff_diff.results import SpilloverDiDResults
 from diff_diff.two_stage import _compute_gmm_corrected_meat
-from diff_diff.utils import safe_inference
+from diff_diff.utils import _iterative_fe_solve, safe_inference
 
 # Type alias mirroring diff_diff.conley.ConleyMetric so callers can supply
 # any of the built-in identifiers or a user callable returning a pairwise
@@ -1286,10 +1286,11 @@ def _convert_treatment_to_first_treat(
 # Two-stage Gardner inline (Step 3)
 # =============================================================================
 
-# Convergence tolerance for the iterative alternating-projection FE solver
-# (Gauss-Seidel style; same recursion as the shared
+# Convergence budget for the stage-1 FE solve (delegated to the shared
 # `diff_diff.utils._iterative_fe_solve` used by ImputationDiD/TwoStageDiD).
-_FE_ITER_MAX = 100
+# max_iter aligned to the shared 10,000 convention (the R fixest/pyfixest
+# budget; historical local cap was 100 - see the SpilloverDiD REGISTRY note).
+_FE_ITER_MAX = 10_000
 _FE_ITER_TOL = 1e-10
 
 
@@ -1401,26 +1402,33 @@ def _iterative_fe_subset(
     max_iter: int = _FE_ITER_MAX,
     tol: float = _FE_ITER_TOL,
     weights: Optional[np.ndarray] = None,
-) -> Tuple[np.ndarray, np.ndarray, bool]:
-    """Stage-1 iterative-alternating-projection FE solver on the Butts subsample.
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Stage-1 FE solve on the Butts subsample via the shared Gauss-Seidel engine.
 
     Fits ``y[Omega_0] = mu_i + lambda_t + u`` on the untreated-and-unexposed
     rows (``Omega_0_mask`` True). Returns FE arrays indexed by code, with
     ``NaN`` at positions whose unit / time is not represented in the
     subsample (rank-deficient cells).
 
-    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.
+    Thin Butts-subsample wrapper over the shared
+    ``diff_diff.utils._iterative_fe_solve`` (the same bincount Gauss-Seidel
+    recursion ImputationDiD/TwoStageDiD route through): this function owns
+    the SpilloverDiD-specific front door — the empty-Omega_0 and empty
+    positive-weight-Omega_0 ``ValueError`` gates and the subsample
+    extraction — and delegates the iteration, the zero-weight NaN-FE
+    convention, and the ``warn_if_not_converged`` non-convergence signal
+    to the shared engine. Per sweep the shared engine computes the
+    identical group means and convergence metric as the historical local
+    loop (converged fits are bit-identical); the iteration budget is the
+    shared ``max_iter=10_000`` convention (see the SpilloverDiD REGISTRY
+    note), superseding the historical local cap of 100.
 
     **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
     positive weights converges to the same fixed point as the unweighted
     iteration for w == 1). The per-period mean becomes
     ``sum_{i in t} w_i * resid_i / sum_{i in t} w_i`` (weighted bincount
-    numerator over weighted bincount denominator). The ``weights is None``
-    branch is bit-identical to the pre-Wave-E.1 path so the Wave B/C/D
-    no-survey contract is unchanged.
+    numerator over weighted bincount denominator).
 
     Parameters
     ----------
@@ -1434,7 +1442,7 @@ def _iterative_fe_subset(
         True for rows in the stage-1 fit subsample (D_it=0 AND S_it=0).
     weights : ndarray of shape (n_rows,), optional
         Hájek-normalized survey weights (``sum_i w_i = n``). When provided,
-        switches the iteration to WLS-FE; when None, the original unweighted
+        switches the iteration to WLS-FE; when None, the unweighted
         bincount path applies.
 
     Returns
@@ -1443,8 +1451,6 @@ def _iterative_fe_subset(
         Unit FE indexed by code. ``NaN`` for units absent from Omega_0.
     time_fe_arr : ndarray of shape (n_times,)
         Time FE indexed by code. ``NaN`` for periods absent from Omega_0.
-    converged : bool
-        Whether the iterative solver reached ``tol`` within ``max_iter``.
     """
     if omega_0_mask.sum() == 0:
         raise ValueError(
@@ -1480,10 +1486,9 @@ def _iterative_fe_subset(
     n_times = int(time_codes_full.max()) + 1
 
     # Operate on the subset only (faster than masking each iteration).
-    y_sub = y_full[omega_0_effective]
-    unit_sub = unit_codes_full[omega_0_effective]
-    time_sub = time_codes_full[omega_0_effective]
-    n_sub = len(y_sub)
+    y_sub = np.asarray(y_full, dtype=np.float64)[omega_0_effective]
+    unit_sub = np.asarray(unit_codes_full)[omega_0_effective]
+    time_sub = np.asarray(time_codes_full)[omega_0_effective]
 
     # Wave E.1: extract weights subset once outside the iterative loop
     # (mirrors TwoStageDiD's `w_0 = weights[omega_0_mask.values]` cache
@@ -1492,57 +1497,17 @@ def _iterative_fe_subset(
     if weights is not None:
         w_sub = np.asarray(weights, dtype=np.float64)[omega_0_effective]
 
-    alpha = np.zeros(n_sub)
-    beta = np.zeros(n_sub)
-    converged = False
-    for _ in range(max_iter):
-        # beta[t] = (weighted) mean over rows in time-group t of (y - alpha)
-        resid = y_sub - alpha
-        if w_sub is None:
-            time_sums = np.bincount(time_sub, weights=resid, minlength=n_times)
-            time_denoms = np.bincount(time_sub, minlength=n_times).astype(np.float64)
-        else:
-            time_sums = np.bincount(time_sub, weights=w_sub * resid, minlength=n_times)
-            time_denoms = np.bincount(time_sub, weights=w_sub, minlength=n_times)
-        time_means = np.where(time_denoms > 0, time_sums / np.maximum(time_denoms, 1e-300), 0.0)
-        beta_new = time_means[time_sub]
-
-        # alpha[i] = (weighted) mean over rows in unit-group i of (y - beta_new)
-        resid = y_sub - beta_new
-        if w_sub is None:
-            unit_sums = np.bincount(unit_sub, weights=resid, minlength=n_units)
-            unit_denoms = np.bincount(unit_sub, minlength=n_units).astype(np.float64)
-        else:
-            unit_sums = np.bincount(unit_sub, weights=w_sub * resid, minlength=n_units)
-            unit_denoms = np.bincount(unit_sub, weights=w_sub, minlength=n_units)
-        unit_means = np.where(unit_denoms > 0, unit_sums / np.maximum(unit_denoms, 1e-300), 0.0)
-        alpha_new = unit_means[unit_sub]
-
-        max_change = max(
-            float(np.max(np.abs(alpha_new - alpha))) if n_sub > 0 else 0.0,
-            float(np.max(np.abs(beta_new - beta))) if n_sub > 0 else 0.0,
-        )
-        alpha = alpha_new
-        beta = beta_new
-        if max_change < tol:
-            converged = True
-            break
-
-    # Build FE arrays indexed by code; NaN for unseen units/periods.
-    unit_fe_arr = np.full(n_units, np.nan, dtype=np.float64)
-    time_fe_arr = np.full(n_times, np.nan, dtype=np.float64)
-    # For each code present in the subset, take any row's converged value
-    # (constant within group at convergence). Sort-by-code to make access
-    # deterministic.
-    seen_unit_codes = np.unique(unit_sub)
-    for u_code in seen_unit_codes:
-        idx = np.flatnonzero(unit_sub == u_code)[0]
-        unit_fe_arr[u_code] = alpha[idx]
-    seen_time_codes = np.unique(time_sub)
-    for t_code in seen_time_codes:
-        idx = np.flatnonzero(time_sub == t_code)[0]
-        time_fe_arr[t_code] = beta[idx]
-    return unit_fe_arr, time_fe_arr, converged
+    return _iterative_fe_solve(
+        y_sub,
+        unit_sub,
+        time_sub,
+        n_units,
+        n_times,
+        weights=w_sub,
+        max_iter=max_iter,
+        tol=tol,
+        method_name="SpilloverDiD stage-1 FE (Butts Omega_0 subsample)",
+    )
 
 
 def _residualize_butts(
@@ -2685,21 +2650,16 @@ def fit(
         # Step 11: stage 1 — fit FE on Omega_0. Wave E.1 threads Hájek-
         # normalized survey weights when survey_design was supplied.
         y_full = np.asarray(data[outcome].values, dtype=np.float64)
-        unit_fe_arr, time_fe_arr, converged = _iterative_fe_subset(
+        # Non-convergence surfaces via the shared engine's
+        # warn_if_not_converged (labelled with the SpilloverDiD stage-1
+        # method name), replacing the historical caller-side warning.
+        unit_fe_arr, time_fe_arr = _iterative_fe_subset(
             y_full,
             np.asarray(unit_codes_full),
             np.asarray(time_codes_full),
             omega_0_mask,
             weights=survey_weights,
         )
-        if not converged:
-            warnings.warn(
-                "SpilloverDiD stage-1 iterative FE solver did not converge "
-                f"within {_FE_ITER_MAX} iterations (tol={_FE_ITER_TOL}). "
-                "Results may be unreliable.",
-                UserWarning,
-                stacklevel=2,
-            )
         stage1_n_obs = int(omega_0_effective.sum())
 
         # Step 12: residualize ALL observations.

diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md
@@ -3997,6 +3997,8 @@ where `D_it` is the treatment indicator (1 if unit `i` is treated by time `t`) a
 
 **Note:** Stage-1 fits unit + time FE on Butts' STRICTER subsample `Omega_0 = {D_it = 0 AND S_it = 0}` (untreated AND unexposed) — the clean far-away control group. This differs from `TwoStageDiD`'s `Omega_0 = {D_it = 0}` (untreated; includes near-controls in post-treatment periods). The stricter Butts subsample prevents spillover-contaminated near-controls from biasing the time FE; near-controls post-treatment carry `delta_j` variation that the ring covariates pick up at stage 2.
 
+**Note (shared FE engine, 2026-07):** The stage-1 solver `_iterative_fe_subset` is a thin Butts-subsample wrapper over the shared bincount Gauss-Seidel helper `diff_diff.utils._iterative_fe_solve` (the same engine ImputationDiD/TwoStageDiD route through) — the wrapper owns the SpilloverDiD front door (empty-Omega_0 and empty positive-weight-Omega_0 `ValueError` gates, subsample extraction) and delegates the iteration, the zero-weight/positive-weight NaN-FE convention, and the `warn_if_not_converged` non-convergence `UserWarning` (labelled "SpilloverDiD stage-1 FE (Butts Omega_0 subsample)") to the shared engine. Per sweep the shared engine computes the identical group means and convergence metric as the historical local loop, so converged fits are bit-identical; `max_iter` is aligned from the historical local cap of 100 to the shared 10,000 convention (the R fixest/pyfixest budget already used by ImputationDiD/TwoStageDiD) — fits that previously exhausted 100 iterations and warned may now converge instead (strictly more accurate FE; `tol=1e-10` unchanged).
+
 **Note (Omega_0 row-level identification — period strict, unit warn-and-drop, plus connectivity):** Every period must have at least one Omega_0 row (else time FE is structurally unidentified for that period, and dropping it would lose all units' cross-time identification) — hard `ValueError`. Units lacking Omega_0 rows (e.g. baseline-treated units with `D_it = 1` at every observed `t`) are warned-and-dropped: their unit FE is NaN, residualization writes NaN on their rows, and the downstream finite-mask path at stage 2 excludes them from estimation. This mirrors `TwoStageDiD`'s always-treated unit handling (`two_stage.py:294-336`) and Gardner's framework, which identifies effects from supported observations rather than requiring every unit estimable. **Connectivity:** the supported-units bipartite graph (supported units linked by shared Omega_0 periods) must form a single connected component. If the graph splits into K > 1 components, the iterative FE solver identifies (`mu_i`, `lambda_t`) only up to component-specific constants, and residualization combines `mu_i` from one component with `lambda_t` from another — silently corrupting `y_tilde` and downstream `tau_total` / `delta_j`. Balanced panel + per-unit/per-period Omega_0 coverage is NECESSARY but NOT SUFFICIENT; connectivity is the load-bearing identification condition. Under the current absorbing-treatment regime the disconnected case is plausibly unreachable in practice (we were unable to construct an example surviving the upstream validators), but `_check_omega_0_connectivity` is in place as defense-in-depth and future-proofs Wave B follow-ups (event-study, survey-design integration, reversible-treatment relaxation if ever added).
 
 **Note (Gardner identity, non-staggered):** Under non-staggered timing, the two-stage Gardner residualize-then-fit with the Omega_0-restricted stage 1 is **empirically bit-identical** to the single-stage TWFE ring regression on the full sample using the time-varying form `y_it ~ mu_i + lambda_t + tau * D_it + sum_j delta_j * (1 - D_it) * Ring_{it,j}`. This is the non-staggered ring estimator from Butts Equations 4-6. The empirical equivalence is verified by a 20-seed deterministic regression test (`TestSpilloverDiDNonStaggeredFEEquivalence`) at `atol=1e-10`. The Omega_0 restriction is therefore innocent for the non-staggered point estimate — it only changes the variance composition (which is why the stage-1 GMM correction enters at stage 2 in the staggered case). Reported `tau_total` for non-staggered timing IS the Butts Eqs. 4-6 estimator.