From 01f4c72c507593397e81cf1a9305ac72f4fa22e3 Mon Sep 17 00:00:00 2001 From: igerber Date: Tue, 7 Jul 2026 16:49:37 -0400 Subject: [PATCH] refactor(spillover): route stage-1 FE through shared _iterative_fe_solve _iterative_fe_subset becomes a thin Butts-subsample wrapper over diff_diff.utils._iterative_fe_solve (the shared bincount Gauss-Seidel engine ImputationDiD/TwoStageDiD already use), taking the library's FE-solver copy count from 2 to 1. The wrapper keeps the SpilloverDiD front door (empty-Omega_0 / empty positive-weight-Omega_0 ValueError gates + subsample extraction); the shared engine owns the iteration, the zero-weight NaN-FE convention, and warn_if_not_converged (labelled 'SpilloverDiD stage-1 FE (Butts Omega_0 subsample)', replacing the caller-side warning). Per sweep the shared engine computes identical group means and convergence metric, so converged fits are bit-identical; max_iter aligned 100 -> 10,000 per the shared fixest/pyfixest convention with a REGISTRY note (tol=1e-10 unchanged). Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 14 +++++ TODO.md | 1 - diff_diff/spillover.py | 118 ++++++++++++----------------------- docs/methodology/REGISTRY.md | 2 + 4 files changed, 55 insertions(+), 80 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2f7c8dab..4dc439e0 100644 --- 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 index e55b68c0..d17aebf3 100644 --- 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 index cb343f81..926a30ae 100644 --- 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 index 01e72b64..4e729992 100644 --- 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.