From 88965eba90ccc025fdeefb1e94528c18bef1f91c Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 4 Jul 2026 10:12:44 -0400 Subject: [PATCH] perf(imputation,two_stage): route demeaning through shared MAP engine + fix zero-weight replicate bug Delete the private pandas _iterative_demean twins; covariate/lead paths now call demean_by_groups (factorize-once + bincount + optional Rust kernel, group order [time, unit] preserving the historical sweep). Replace both _iterative_fe bodies with a shared bincount Gauss-Seidel helper (utils._iterative_fe_solve, modeled on spillover._iterative_fe_subset); zero-total-weight groups get NaN FE (keys retained) instead of 0/0 NaN poisoning. Fixes covariate + zero-weight replicate designs (JK1/BRR): previously ALL replicate refits failed -> NaN SE + non-convergence warning storm; main fits with zero-weight rows raised opaque ValueError. TwoStageDiD stage-2 nan-ytilde warning suppressed inside replicate closures only (warn_nan=False); main-fit warning unchanged. max_iter modernized 100 -> 10_000. REGISTRY notes rewritten; TODO row resolved with two follow-up rows (spillover migration, lead-path snap guard). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X4AzrFUMqJxcUumSH31mSr --- CHANGELOG.md | 32 +++++ TODO.md | 3 +- diff_diff/imputation.py | 212 ++++++++++---------------------- diff_diff/spillover.py | 8 +- diff_diff/two_stage.py | 200 +++++++++++-------------------- diff_diff/utils.py | 112 +++++++++++++++++ docs/methodology/REGISTRY.md | 6 +- docs/performance-plan.md | 5 +- tests/test_imputation.py | 226 ++++++++++++++++++++++++++++------- tests/test_two_stage.py | 131 +++++++++++++++++--- tests/test_utils.py | 146 ++++++++++++++++++++++ 11 files changed, 739 insertions(+), 342 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9fdf8cb3..358a76e5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `survey_design=SurveyDesign(psu=)`. 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 diff --git a/TODO.md b/TODO.md index 69c80cfa..98724632 100644 --- a/TODO.md +++ b/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 --git a/diff_diff/imputation.py b/diff_diff/imputation.py index e363b581..9e703073 100644 --- a/diff_diff/imputation.py +++ b/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 --git a/diff_diff/spillover.py b/diff_diff/spillover.py index 1bac157f..cb343f81 100644 --- a/diff_diff/spillover.py +++ b/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 diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index 6209b3e2..4335eab0 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -42,7 +42,7 @@ TwoStageBootstrapResults, # noqa: F401 TwoStageDiDResults, ) # noqa: F401 (re-export) -from diff_diff.utils import safe_inference, warn_if_not_converged +from diff_diff.utils import _iterative_fe_solve, demean_by_groups, safe_inference if TYPE_CHECKING: # Forward reference for the Wave E.1 survey-design path. Imported under @@ -1871,6 +1871,7 @@ def _refit_ts(w_r): kept_cov_mask=kcm_r, survey_weights=w_r_fit, survey_weight_type="pweight", + warn_nan=False, ) results.append(att_r) @@ -1895,6 +1896,7 @@ def _refit_ts(w_r): survey_weights=w_r_fit, survey_weight_type="pweight", survey_df=None, + warn_nan=False, ) for e in _sorted_es_periods_ts: results.append(es_r[e]["effect"] if e in es_r else np.nan) @@ -1918,6 +1920,7 @@ def _refit_ts(w_r): survey_weights=w_r_fit, survey_weight_type="pweight", survey_df=None, + warn_nan=False, ) for g in _sorted_groups_ts: results.append(grp_r[g]["effect"] if g in grp_r else np.nan) @@ -2126,18 +2129,26 @@ 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. + 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. + 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 ------- @@ -2146,112 +2157,28 @@ def _iterative_fe( time_fe : dict Mapping from time -> time fixed effect. """ - n = len(y) - alpha = np.zeros(n) - beta = np.zeros(n) - - 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 - ) - - 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, "TwoStageDiD 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( + "TwoStageDiD: 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="TwoStageDiD 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). - - 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() - - 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, "TwoStageDiD iterative demean", max_iter, tol) - return result - def _fit_untreated_model( self, df: pd.DataFrame, @@ -2292,15 +2219,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) - - 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) - ] + + # 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) result = solve_ols( X_dm, @@ -2360,22 +2296,27 @@ def _residualize( # ========================================================================= @staticmethod - def _mask_nan_ytilde(y_tilde): + def _mask_nan_ytilde(y_tilde, warn: bool = True): """Mask non-finite y_tilde values and warn if any found. Returns the boolean mask of non-finite values. Modifies y_tilde in-place - (sets NaN values to 0.0). + (sets NaN values to 0.0). ``warn=False`` suppresses the UserWarning - + used ONLY by the replicate-refit closures, where zero-weight replicate + designs (JK1/BRR) make NaN FE for zeroed-out PSUs expected mechanics + (the main-fit warning still fires once; per-replicate repeats would + emit up to ~3x n_replicates copies of the same message). """ nan_mask = ~np.isfinite(y_tilde) if nan_mask.any(): n_nan = int(nan_mask.sum()) - warnings.warn( - f"{n_nan} observation(s) have non-finite imputed outcomes " - f"(y_tilde) from unidentified fixed effects. These " - f"observations are excluded from ATT estimation.", - UserWarning, - stacklevel=3, - ) + if warn: + warnings.warn( + f"{n_nan} observation(s) have non-finite imputed outcomes " + f"(y_tilde) from unidentified fixed effects. These " + f"observations are excluded from ATT estimation.", + UserWarning, + stacklevel=3, + ) y_tilde[nan_mask] = 0.0 return nan_mask @@ -2399,6 +2340,7 @@ def _stage2_static( resolved_survey=None, score_pad_mask: Optional[np.ndarray] = None, cluster_ids_full: Optional[np.ndarray] = None, + warn_nan: bool = True, ) -> Tuple[float, float]: """ Static (simple ATT) Stage 2: OLS of y_tilde on D_it. @@ -2406,7 +2348,7 @@ def _stage2_static( Returns (att, se). """ y_tilde = df["_y_tilde"].values.copy() - nan_mask = self._mask_nan_ytilde(y_tilde) + nan_mask = self._mask_nan_ytilde(y_tilde, warn=warn_nan) D = omega_1_mask.values.astype(float) # Zero out treatment indicator for NaN y_tilde obs (don't count in ATT) @@ -2475,10 +2417,11 @@ def _stage2_event_study( resolved_survey=None, score_pad_mask: Optional[np.ndarray] = None, cluster_ids_full: Optional[np.ndarray] = None, + warn_nan: bool = True, ) -> Dict[int, Dict[str, Any]]: """Event study Stage 2: OLS of y_tilde on relative-time dummies.""" y_tilde = df["_y_tilde"].values.copy() - nan_mask = self._mask_nan_ytilde(y_tilde) + nan_mask = self._mask_nan_ytilde(y_tilde, warn=warn_nan) rel_times = df["_rel_time"].values n = len(df) @@ -2695,10 +2638,11 @@ def _stage2_group( resolved_survey=None, score_pad_mask: Optional[np.ndarray] = None, cluster_ids_full: Optional[np.ndarray] = None, + warn_nan: bool = True, ) -> Dict[Any, Dict[str, Any]]: """Group (cohort) Stage 2: OLS of y_tilde on cohort dummies.""" y_tilde = df["_y_tilde"].values.copy() - nan_mask = self._mask_nan_ytilde(y_tilde) + nan_mask = self._mask_nan_ytilde(y_tilde, warn=warn_nan) n = len(df) # Build Stage 2 design: one column per cohort (no intercept) diff --git a/diff_diff/utils.py b/diff_diff/utils.py index fc0c0a3b..ebd5c39f 100644 --- a/diff_diff/utils.py +++ b/diff_diff/utils.py @@ -2977,6 +2977,118 @@ def demean_by_groups( return pd.concat([data, new_block], axis=1), n_effects +def _iterative_fe_solve( + y: np.ndarray, + unit_codes: np.ndarray, + time_codes: np.ndarray, + n_units: int, + n_times: int, + weights: Optional[np.ndarray] = None, + max_iter: int = 10_000, + tol: float = 1e-10, + *, + method_name: str, +) -> Tuple[np.ndarray, np.ndarray]: + """Two-way Gauss-Seidel FE solver on integer-coded factors via bincount. + + Estimates ``y = alpha_i + beta_t + u`` by alternating projections + (beta from ``y - alpha``, alpha from ``y - beta_new``), the same + recursion as the historical per-estimator pandas loops in ImputationDiD + and TwoStageDiD but with each factor coded once and each sweep two O(n) + ``np.bincount`` passes (same engine family as ``_demean_map_numpy``; + same accumulation-order caveat vs pandas' compensated grouped mean, see + REGISTRY "Absorbed Fixed Effects"). SpilloverDiD's + ``_iterative_fe_subset`` is the structural sibling. + + Zero-weight convention (mirrors ``spillover._iterative_fe_subset`` and + the SpilloverDiD REGISTRY contract): rows with ``weights == 0`` are + outside the WLS estimating sample, so any unit/period whose rows ALL + carry zero weight has no identifying contribution and surfaces as + ``NaN`` FE — never a silent finite ``0.0``. The iteration itself runs + on the positive-weight subset, which also keeps a zero-total-weight + group from NaN-poisoning the convergence check (the historical pandas + loops divided 0/0 there and could never converge). + + Parameters + ---------- + y : ndarray of shape (n,) + Outcome values (float64). + unit_codes, time_codes : ndarray of shape (n,) + Non-negative integer factor codes (``pd.factorize`` output). + n_units, n_times : int + Number of factor levels (lengths of the factorize uniques). + weights : ndarray of shape (n,), optional + Survey weights; weighted group means ``sum(w*x)/sum(w)``. + max_iter, tol : convergence budget; warns via ``warn_if_not_converged`` + (labelled ``method_name``) when exhausted. + method_name : str + Caller label for the non-convergence warning (keyword-only). + + Returns + ------- + unit_fe_arr : ndarray of shape (n_units,) + Unit FE indexed by code; NaN for codes absent from the + (positive-weight) estimating sample. + time_fe_arr : ndarray of shape (n_times,) + Time FE indexed by code; NaN likewise. + """ + if weights is not None: + w_arr = np.asarray(weights, dtype=np.float64) + keep = w_arr > 0 + y_sub = y[keep] + unit_sub = unit_codes[keep] + time_sub = time_codes[keep] + w_sub: Optional[np.ndarray] = w_arr[keep] + else: + y_sub = y + unit_sub = unit_codes + time_sub = time_codes + w_sub = None + + # Per-group denominators are invariant across iterations. + if w_sub is None: + unit_denoms = np.bincount(unit_sub, minlength=n_units).astype(np.float64) + time_denoms = np.bincount(time_sub, minlength=n_times).astype(np.float64) + else: + unit_denoms = np.bincount(unit_sub, weights=w_sub, minlength=n_units) + time_denoms = np.bincount(time_sub, weights=w_sub, minlength=n_times) + unit_seen = unit_denoms > 0 + time_seen = time_denoms > 0 + unit_safe = np.maximum(unit_denoms, 1e-300) + time_safe = np.maximum(time_denoms, 1e-300) + + alpha = np.zeros(n_units) + beta = np.zeros(n_times) + converged = False + for _ in range(max_iter): + resid = y_sub - alpha[unit_sub] + wx = resid if w_sub is None else w_sub * resid + beta_new = np.where( + time_seen, np.bincount(time_sub, weights=wx, minlength=n_times) / time_safe, 0.0 + ) + + resid = y_sub - beta_new[time_sub] + wx = resid if w_sub is None else w_sub * resid + alpha_new = np.where( + unit_seen, np.bincount(unit_sub, weights=wx, minlength=n_units) / unit_safe, 0.0 + ) + + max_change = max( + float(np.max(np.abs(alpha_new - alpha))), + float(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, method_name, max_iter, tol) + + unit_fe_arr = np.where(unit_seen, alpha, np.nan) + time_fe_arr = np.where(time_seen, beta, np.nan) + return unit_fe_arr, time_fe_arr + + def pre_demean_norms( data: pd.DataFrame, regressors: List[str], diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index a9d07b7a..8f4e26b6 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -1410,7 +1410,8 @@ where `W_it(h) = 1[K_it = h]` are lead indicators, estimated on `Omega_0` only. - **Bootstrap inference:** Uses multiplier bootstrap on the Theorem 3 influence function: `psi_i = sum_t v_it * epsilon_tilde_it`. Cluster-level psi sums are pre-computed for each aggregation target (overall, per-horizon, per-group), then perturbed with multiplier weights (Rademacher by default; configurable via `bootstrap_weights` parameter to use Mammen or Webb weights, matching CallawaySantAnna). This is a library extension (not in the paper) consistent with CallawaySantAnna/SunAbraham bootstrap patterns. - **Auxiliary residuals (Equation 8):** Implements the paper's *unit-clustered* Equation 8 aggregator, `tau_tilde_g = sum_i (sum_{t in G_g,i} v_it)(sum_{t in G_g,i} v_it * tau_hat_it) / sum_i (sum_{t in G_g,i} v_it)^2` (Borusyak-Jaravel-Spiess 2024, eq. 8, p. 3272; minimal-excess-variance derivation in Supplementary Appendix A.8): for each unit form the within-unit weight sum `a_{i,g}` and weighted-effect sum `b_{i,g}` over the unit's observations in group `g`, then combine across units. Groups partition `Omega_1` via `aux_partition` (default `"cohort_horizon"` = cohort × event-time; also `"cohort"` / `"horizon"`). Unimputable (NaN `tau_hat`) and off-target observations carry `v_it = 0` and are excluded from the aggregation — exact for finite `tau_hat` (a zero-weight row adds 0 to both `a` and `b`) and NaN-safe; a group with no contributing observations falls back to the unweighted group mean (a variance no-op, since `psi_g = sum_t v_it * eps_tilde_it = 0` there). - **Note (deviation from R):** R `didimputation::did_imputation` computes the auxiliary aggregator as `sum(v_it^2 * tau_hat_it) / sum(v_it^2)` grouped by cohort × event-time only (no partition control is exposed). At that partition each unit contributes at most one observation per group, so the paper's unit-clustered Equation 8 reduces exactly to `sum(v^2 * tau)/sum(v^2)` — i.e. diff-diff matches R at the default `aux_partition="cohort_horizon"` (pinned in `tests/test_methodology_imputation.py::TestImputationDiDParityR`). diff-diff additionally offers the coarser `aux_partition="cohort"` / `"horizon"` groupings (where a unit may contribute several observations to a group), which have no R analogue and are validated by hand-calculation. Both implement the same paper Equation 8; only the available partition granularity differs. The earlier observation-level mean `sum(v*tau)/sum(v)` (pre-3.5.x) coincided with this only under uniform within-group weights; it was corrected to the exact unit-clustered form during the ImputationDiD methodology validation. -- **Note:** Both the iterative FE solver (`_iterative_fe`, Step 1) and the iterative alternating-projection demeaning helper (`_iterative_demean`, used in covariate residualization and the pre-trend test) emit `UserWarning` when `max_iter` exhausts without reaching `tol`, via `diff_diff.utils.warn_if_not_converged`. Silent return of the current iterate was classified as a silent failure under the Phase 2 audit and replaced with an explicit signal to match the logistic/Poisson IRLS pattern in `linalg.py`. +- **Note:** The Step-1 iterative FE solver (`_iterative_fe`) routes through the shared bincount Gauss-Seidel helper `diff_diff.utils._iterative_fe_solve`, and the covariate/pre-trend within-transformation routes through the shared MAP engine `diff_diff.utils.demean_by_groups` (factorize-once + `np.bincount`, optional Rust kernel; group order `[time, unit]` preserving the historical time-then-unit sweep) — the same convergence contract, accumulation-order numerics (~1e-10 vs the pre-3.7 pandas loops, not bit-for-bit), and `max_iter=10_000` budget documented under "Absorbed Fixed Effects with Survey Weights". Both surfaces emit `UserWarning` via `diff_diff.utils.warn_if_not_converged` when `max_iter` exhausts without reaching `tol` (the demean warning now carries the shared-engine label naming the affected variables rather than the estimator name). Silent return of the current iterate was classified as a silent failure under the Phase 2 audit and replaced with an explicit signal to match the logistic/Poisson IRLS pattern in `linalg.py`. +- **Note:** Zero-total-weight groups (e.g. whole PSUs zeroed by JK1/BRR replicate weights, which reach Step 1 unmasked — `keep_mask` only drops always-treated units): a unit/period whose observations ALL carry zero weight has no identifying contribution and surfaces as `NaN` FE (key retained for the rank-condition membership check; matches the SpilloverDiD `_iterative_fe_subset` REGISTRY contract — never a silent finite `0.0`), and the shared demean engine's inert-row guard leaves those rows un-demeaned instead of NaN-poisoning the column. Before 3.7 the pandas loops divided 0/0 there: the covariate replicate path NaN-poisoned `y_dm`/`X_dm`, failed EVERY replicate refit inside `solve_ols(check_finite=True)`, and returned NaN SEs after a non-convergence warning storm; a main fit with zero-weight rows + covariates raised the same opaque `ValueError`. Both now produce finite results. `TwoStageDiD._mask_nan_ytilde`'s "non-finite imputed outcomes" `UserWarning` is suppressed (via `warn_nan=False`) ONLY inside the replicate-refit closures, where NaN FE for zeroed PSUs is expected mechanics — the main-fit warning is unchanged. - **Note:** `vcov_type` is permanently narrow to `{"hc1"}` per the Theorem 3 IF-based variance decomposition. Analytical-sandwich families `{classical, hc2, hc2_bm}` are rejected at `__init__` — the per-unit influence function aggregation has no equivalent single design matrix on which hat-matrix leverage or Bell-McCaffrey Satterthwaite DOF can be defined. `cluster=` invokes per-cluster IF summation (Theorem 3 equation 7 conservative variance, `sigma_sq = (cluster_psi_sums**2).sum()` — plain CR1 with no Stata-style `(n-1)/(n-p)` finite-sample factor because the IF has no design-matrix `p` in the OLS sense); `cluster=None` (the default) routes the SAME Theorem 3 cluster-summed IF variance with `cluster_var = unit` (the unit column passed to `fit()`), so the summary renders `"CR1 cluster-robust at , G="` rather than the generic `"HC1"` label; `survey_design=` invokes TSL on the combined IF. Under bootstrap (`n_bootstrap > 0`) the analytical variance-family label is suppressed in `summary()` because `fit()` overwrites the reported SE/CI/p-value with bootstrap_results (mirrors the canonical `DiDResults` gate at `results.py:213-226`). `vcov_type='conley'` is deferred to the ImputationDiD Conley follow-up row in TODO.md. - **Note:** `cluster=` combined with a replicate-weight `SurveyDesign` raises `NotImplementedError` at `fit()`. Replicate-weight variance ignores PSU/cluster structure entirely (replicates encode the design implicitly), so honoring `cluster=` would silently no-op while populating `cluster_name`/`n_clusters` on Results dishonestly. Either omit `cluster=` (the replicate weights encode the design structure implicitly) or use a non-replicate survey design (with explicit strata/psu/fpc). Mirrors the `CallawaySantAnna` and `TripleDifference` fail-closed guards. - **Note:** Bootstrap path returns NaN SE when fewer than 2 independent clusters/PSUs are available (`n_clusters < 2` for the analytical-cluster bootstrap path, `n_psu < 2` for the survey-PSU bootstrap path). Without this guard the multiplier bootstrap SE collapses to ≈0 from BLAS roundoff (NOT NaN), and downstream zero-SE guards check exact 0 and miss the degenerate-design case. NaN propagates to all inference fields (SE/CI/p-value) plus per-horizon and per-group bootstrap dicts. @@ -1494,7 +1495,8 @@ Our implementation uses multiplier bootstrap on the GMM influence function: clus - **Zero-observation cohorts in group effects:** If all treated observations for a cohort have NaN `y_tilde` (excluded from estimation), that cohort's group effect is NaN with n_obs=0. - **Note:** Survey weights in TwoStageDiD GMM sandwich via weighted cross-products: bread uses (X'_2 W X_2)^{-1}, gamma_hat uses (X'_{10} W X_{10})^{-1}(X'_1 W X_2), per-cluster scores multiply by survey weights. PSU clustering, stratification, and FPC are fully supported in the meat matrix via `_compute_stratified_meat_from_psu_scores()`. When strata or FPC are present, the meat computation replaces `S' S` with the stratified formula `sum_h (1 - f_h) * (n_h/(n_h-1)) * centered_h' centered_h`. Strata also enters survey df (n_PSU - n_strata) for t-distribution inference. Bootstrap + survey supported (Phase 6) via PSU-level multiplier weights. - **Note (documented synthesis — Wave E.3 parity, full-domain survey design under always-treated drop):** when the always-treated handler drops units that lack untreated observations, TwoStageDiD preserves the FULL-DOMAIN resolved survey design (`n_psu`, `n_strata`, `df_survey`, `strata`, `fpc`, `psu`) for variance estimation. Per-cluster stage-1 / stage-2 score aggregates are computed at the post-drop fit-sample length and then zero-padded onto the full-domain unique-PSU list via `score_pad_mask` + `cluster_ids_full` kwargs on `_compute_gmm_variance`; PSUs that contain only always-treated rows get zero score rows but still count toward `G_full` for `n_psu` / `df_survey` accounting. Stage-1 / stage-2 OLS solve continues to operate on the post-drop sample (`survey_weights` subsetted for OLS arithmetic; bread `(X'_2 W X_2)^{-1}` unchanged because dropped rows contribute zero score under zero-padded weights). Mirrors SpilloverDiD Wave E.3 (PR #482, merge 24de9062) and adopts the canonical "zero-pad scores to full panel + retain full-design resolved survey" convention from R `survey::svyrecvar(subset())` (Lumley 2010 §2.5 "Domains and subpopulations") and the in-library precedents at `imputation.py:2175-2183` (PreTrendsImputation) and `prep.py:1401-1432` (DCDH cell variance). Cluster-injection (`_inject_cluster_as_psu`) operates on the FULL-DOMAIN cluster column (sourced from `data` pre-drop, not the post-drop `df`) so `resolved_survey.strata` and the injected `psu` array stay length-aligned. Pre-PR, the always-treated drop physically subsetted `resolved_survey.weights / strata / psu / fpc / replicate_weights` via `replace(resolved_survey, ...)` and recomputed `n_psu` / `n_strata` on the post-drop sample, producing artificially-deflated `df_survey` when a PSU contained only always-treated rows; tests at `tests/test_two_stage.py::TestTwoStageDiDWaveE3ParityAlwaysTreated` lock the parity contract. -- **Note:** Both the iterative FE solver (`_iterative_fe`, Stage 1) and the iterative alternating-projection demeaning helper (`_iterative_demean`, used in covariate residualization) emit `UserWarning` when `max_iter` exhausts without reaching `tol`, via `diff_diff.utils.warn_if_not_converged`. Silent return of the current iterate was classified as a silent failure under the Phase 2 audit and replaced with an explicit signal to match the logistic/Poisson IRLS pattern in `linalg.py`. +- **Note:** The Stage-1 iterative FE solver (`_iterative_fe`) routes through the shared bincount Gauss-Seidel helper `diff_diff.utils._iterative_fe_solve`, and the covariate within-transformation routes through the shared MAP engine `diff_diff.utils.demean_by_groups` (factorize-once + `np.bincount`, optional Rust kernel; group order `[time, unit]` preserving the historical time-then-unit sweep) — the same convergence contract, accumulation-order numerics (~1e-10 vs the pre-3.7 pandas loops, not bit-for-bit), and `max_iter=10_000` budget documented under "Absorbed Fixed Effects with Survey Weights". Both surfaces emit `UserWarning` via `diff_diff.utils.warn_if_not_converged` when `max_iter` exhausts without reaching `tol` (the demean warning now carries the shared-engine label naming the affected variables rather than the estimator name). Silent return of the current iterate was classified as a silent failure under the Phase 2 audit and replaced with an explicit signal to match the logistic/Poisson IRLS pattern in `linalg.py`. +- **Note:** Zero-total-weight groups (e.g. whole PSUs zeroed by JK1/BRR replicate weights, which reach Stage 1 unmasked — `keep_mask` only drops always-treated units): a unit/period whose observations ALL carry zero weight surfaces as `NaN` FE (key retained for the rank-condition membership check; matches the SpilloverDiD `_iterative_fe_subset` REGISTRY contract — never a silent finite `0.0`), and the shared demean engine's inert-row guard leaves those rows un-demeaned instead of NaN-poisoning the column. Before 3.7 the pandas loops divided 0/0 there: the covariate replicate path NaN-poisoned `y_dm`/`X_dm`, failed EVERY replicate refit inside `solve_ols(check_finite=True)`, and returned NaN SEs after a non-convergence warning storm. It now produces finite replicate SEs. `_mask_nan_ytilde`'s "non-finite imputed outcomes" `UserWarning` is suppressed (via `warn_nan=False`) ONLY inside the replicate-refit closures, where NaN FE for zeroed PSUs is expected mechanics — the main-fit warning is unchanged. - **Note:** When the Stage-2 bread `X'_2 W X_2` is singular, both the analytical TSL variance (`two_stage.py`) and the multiplier-bootstrap bread (`two_stage_bootstrap.py`) now emit a `UserWarning` before falling back to `np.linalg.lstsq`. Previously this fallback was silent. Sibling of axis-A finding #17 in the Phase 2 silent-failures audit; surfaced by the repo-wide lstsq-fallback pattern grep that accompanied the StaggeredTripleDifference fix. - **Note:** The GMM sandwich and bootstrap paths both use `scipy.sparse.linalg.factorized` for the Stage 1 normal-equations solve `(X'_{10} W X_{10}) gamma = X'_1 W X_2` and fall back to dense `lstsq` when the sparse factorization raises `RuntimeError` on a near-singular matrix. Both fallback sites emit a `UserWarning` (silent-failure audit axis C) so callers know SE estimates came from the degraded path rather than the fast sparse path. - **Note:** The GMM sandwich re-solves the Stage-1 unit+time fixed effects **exactly** (sparse OLS reusing the `scipy.sparse.linalg.factorized` factorization of `(X'_{10} W X_{10})` already computed for `gamma_hat`), rather than reusing the iterative alternating-projection FE (`_iterative_fe`) that produces the point estimate. The iterative solver converges only to ~1e-7 on unbalanced *untreated* panels — negligible for the ATT, but enough to perturb the variance by ~1% relative to the analytical GMM sandwich. The exact re-solve makes the analytical GMM SE match R `did2s` to ~1e-7 (`tests/test_methodology_two_stage.py::TestTwoStageDiDParityR`), mirroring ImputationDiD's exact-sparse variance path (`imputation.py` `_build_A_sparse`). It also required adding an **intercept column** to `_build_fe_design` so the first-stage column space spans the constant (the grand mean): the prior intercept-free `[unit_1.., time_1..]` layout (drop first unit + first time, no intercept) silently omitted the grand mean, which the exact residual is first-order sensitive to (the iterative point-estimate solver absorbs the grand mean into its mean-based FE, so the point estimate was unaffected). Obs whose unit or time FE are unidentified (NaN; rank-deficient / Proposition-5) fall back to the iterative residual, so those edge cases are unchanged. The reported `overall_att` still uses the iterative FE (preserving point-estimate equivalence with ImputationDiD at 1e-10); only the variance uses the exact residuals. diff --git a/docs/performance-plan.md b/docs/performance-plan.md index e2a245b8..048b63aa 100644 --- a/docs/performance-plan.md +++ b/docs/performance-plan.md @@ -170,8 +170,9 @@ copy elision is a smaller, broadly-applicable win on the TWFE-family fit path. S ~12.7-15.0 GB across repeats), so the table reports medians - the reduction ratios are stable. Deferred (see `TODO.md`): tiling the *stratified* survey-PSU weight generator (few PSUs, rarely -OOMs), and vectorizing `ImputationDiD` / `TwoStageDiD` `_iterative_demean` (a speed/churn win -that is not bit-identical on pandas 3.0). +OOMs). The `ImputationDiD` / `TwoStageDiD` `_iterative_demean` vectorization deferred here +shipped later as the demean-modernization PR (both estimators now route through the shared +`demean_by_groups` MAP engine and the shared `_iterative_fe_solve` bincount FE solver). --- diff --git a/tests/test_imputation.py b/tests/test_imputation.py index 0245fc42..2d7f65e2 100644 --- a/tests/test_imputation.py +++ b/tests/test_imputation.py @@ -1847,25 +1847,48 @@ def test_overall_se_with_partial_nan_tau_hat(self): ), f"overall_se should be finite with {n_finite} finite and {n_nan} NaN tau_hat" assert np.isfinite(results.overall_att) - def test_iterative_demean_balanced_matches_one_pass(self): - """Test _iterative_demean matches one-pass for balanced panels.""" - rng = np.random.default_rng(42) - n_units, n_periods = 20, 5 - units = np.repeat(np.arange(n_units), n_periods) - times = np.tile(np.arange(n_periods), n_units) - vals = rng.standard_normal(n_units * n_periods) - idx = pd.RangeIndex(len(vals)) + def test_covariate_delta_matches_full_dummy_ols(self): + """Step-A/B covariate coefficients match explicit unit+time dummy OLS. - result_iter = ImputationDiD._iterative_demean(vals, units, times, idx) + Estimator-level lock on the shared-engine within-transform path + (replaces the direct `_iterative_demean` balanced-one-pass test; the + private per-estimator demean loops were consolidated into + `diff_diff.utils.demean_by_groups`). + """ + rng = np.random.default_rng(42) + n_units, n_periods = 15, 6 + rows = [] + for i in range(n_units): + for t in range(n_periods): + if rng.random() < 0.2: # unbalanced + continue + rows.append({"unit": i, "time": t}) + df = pd.DataFrame(rows) + n = len(df) + x = rng.standard_normal((n, 2)) + df["x1"], df["x2"] = x[:, 0], x[:, 1] + u_fe = rng.standard_normal(n_units) + t_fe = np.linspace(0, 1, n_periods) + df["outcome"] = ( + u_fe[df["unit"]] + + t_fe[df["time"]] + + x @ np.array([0.7, -0.4]) + + rng.standard_normal(n) * 0.1 + ) + df["first_treat"] = 0 # all never-treated -> omega_0 = everything - # One-pass for balanced panel - s = pd.DataFrame({"val": vals, "unit": units, "time": times}) - gm = s["val"].mean() - um = s.groupby("unit")["val"].transform("mean").values - tm = s.groupby("time")["val"].transform("mean").values - result_onepass = vals - um - tm + gm + est = ImputationDiD() + omega_0 = pd.Series(True, index=df.index) + _, _, _, delta_hat, _ = est._fit_untreated_model( + df, "outcome", "unit", "time", ["x1", "x2"], omega_0 + ) - np.testing.assert_allclose(result_iter, result_onepass, atol=1e-8) + # Explicit full-dummy OLS: [all unit dummies, time dummies (drop 1), X] + u_d = pd.get_dummies(df["unit"]).values.astype(float) + t_d = pd.get_dummies(df["time"]).values.astype(float)[:, 1:] + X_full = np.column_stack([u_d, t_d, x]) + coef = np.linalg.lstsq(X_full, df["outcome"].values, rcond=None)[0] + np.testing.assert_allclose(delta_hat, coef[-2:], atol=1e-8) def test_unbalanced_panel_fe_correctness(self): """Test FE estimates match OLS for unbalanced panel.""" @@ -2162,31 +2185,41 @@ def test_iterative_fe_no_warning_on_convergence(self): est._iterative_fe(y, units, times, idx) assert not any("did not converge" in str(x.message) for x in w) - def test_iterative_demean_warns_on_nonconvergence(self): - """Silent-failure audit axis B: _iterative_demean must warn when max_iter exhausts.""" - rng = np.random.default_rng(42) - n_units, n_periods = 8, 5 - units = np.repeat(np.arange(n_units), n_periods) - times = np.tile(np.arange(n_periods), n_units) - vals = rng.standard_normal(n_units * n_periods) - idx = pd.RangeIndex(len(vals)) - - with pytest.warns(UserWarning, match="did not converge"): - ImputationDiD._iterative_demean(vals, units, times, idx, max_iter=1, tol=1e-15) - - def test_iterative_demean_no_warning_on_convergence(self): - """Silent-failure audit axis B: no warning on well-behaved convergent input.""" + # NOTE (intentional coverage narrowing): the direct `_iterative_demean` + # non-convergence warn tests were retired with the method itself - the + # covariate within-transform now routes through the shared MAP engine + # with max_iter=10_000 hardcoded at the call sites, so demean + # non-convergence is no longer forceable THROUGH this estimator. + # Engine-level warning coverage lives in + # tests/test_utils.py::TestDemeanByGroups; the `_iterative_fe` warn + # tests above still exercise this estimator's FE-solver warning path. + + def test_iterative_fe_zero_weight_unit_gets_nan_fe(self): + """A unit whose rows ALL carry zero weight surfaces as NaN FE. + + Locks the shared-solver zero-weight contract (spillover precedent: + never a silent finite 0.0) AND that the solver still converges + cleanly - the historical pandas loop divided 0/0 there and burned + max_iter iterations before warning. + """ rng = np.random.default_rng(42) n_units, n_periods = 8, 5 units = np.repeat(np.arange(n_units), n_periods) times = np.tile(np.arange(n_periods), n_units) - vals = rng.standard_normal(n_units * n_periods) - idx = pd.RangeIndex(len(vals)) + y = rng.standard_normal(n_units * n_periods) + w = np.ones(n_units * n_periods) + w[units == 3] = 0.0 # zero out one whole unit + idx = pd.RangeIndex(len(y)) + est = ImputationDiD() - with warnings.catch_warnings(record=True) as w: + with warnings.catch_warnings(record=True) as rec: warnings.simplefilter("always") - ImputationDiD._iterative_demean(vals, units, times, idx) - assert not any("did not converge" in str(x.message) for x in w) + unit_fe, time_fe = est._iterative_fe(y, units, times, idx, weights=w) + assert not any("did not converge" in str(x.message) for x in rec) + + assert np.isnan(unit_fe[3]) # zero-weight unit: NaN, key retained + assert all(np.isfinite(v) for u, v in unit_fe.items() if u != 3) + assert all(np.isfinite(v) for v in time_fe.values()) # ============================================================================= @@ -2700,18 +2733,19 @@ def test_pretrends_hc1_bit_equal_with_cluster(self): # the same Theorem 3 variance machinery — values across default vs # explicit hc1 must agree to within iterative-FE-solver convergence # tolerance. Sub-ULP differences come from BLAS non-associativity in - # `_iterative_demean` across distinct estimator instances and are not - # methodological divergence under the narrow vcov_type contract. + # the shared MAP demean engine (`demean_by_groups`) across distinct + # estimator instances and are not methodological divergence under + # the narrow vcov_type contract. pt_default = r_default.pretrend_test() pt_explicit = r_explicit.pretrend_test() assert np.isclose(pt_default["f_stat"], pt_explicit["f_stat"], rtol=0, atol=1e-12) assert np.isclose(pt_default["p_value"], pt_explicit["p_value"], rtol=0, atol=1e-12) # Event-study SE (computed during fit() via Theorem 3 machinery on - # iterative-demeaned residuals; pretrends=True path includes the + # within-transformed residuals; pretrends=True path includes the # pre-period horizons). Sub-ULP differences come from BLAS - # non-associativity in `_iterative_demean` across distinct estimator - # instances and are not methodological divergence under the narrow - # vcov_type contract. + # non-associativity in the shared MAP demean engine + # (`demean_by_groups`) across distinct estimator instances and are + # not methodological divergence under the narrow vcov_type contract. assert r_default.event_study_effects is not None assert r_explicit.event_study_effects is not None for h in r_default.event_study_effects: @@ -2912,3 +2946,113 @@ def test_imputation_did_convenience_func_threads_vcov_type(self): vcov_type="hc1", ) assert r.vcov_type == "hc1" + + +# ============================================================================= +# TestZeroWeightGroups — shared-engine migration behavioral locks +# ============================================================================= + + +class TestZeroWeightGroups: + """Zero-total-weight groups on the Step-1 paths (shared-engine migration). + + JK1/plain-BRR replicate weights zero whole PSUs and reach Step 1 unmasked. + Before the shared-engine migration the pandas loops divided 0/0 there: + with covariates, y_dm/X_dm NaN-poisoned, EVERY replicate refit failed + inside solve_ols(check_finite=True), and the fit returned NaN SEs after a + non-convergence warning storm. These tests lock the fixed contract. + """ + + @staticmethod + def _with_covariates(data, seed=7): + rng = np.random.default_rng(seed) + d = data.copy() + x = rng.standard_normal((len(d), 2)) + d["x1"], d["x2"] = x[:, 0], x[:, 1] + d["outcome"] = d["outcome"] + x @ np.array([0.6, -0.3]) + return d + + def test_replicate_covariates_zero_weight_psus_finite_se(self): + """Covariates + JK1 zeroed-PSU replicates -> finite SE, no warning storm.""" + data, rep_cols = _imputation_replicate_panel() + data = self._with_covariates(data) + design = SurveyDesign( + weights="weight", + replicate_weights=rep_cols, + replicate_method="JK1", + weight_type="pweight", + ) + with warnings.catch_warnings(record=True) as rec: + warnings.simplefilter("always") + r = ImputationDiD().fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + covariates=["x1", "x2"], + survey_design=design, + ) + messages = [str(w.message) for w in rec] + assert not any("replicate refits failed" in m for m in messages) + assert not any("did not converge" in m for m in messages) + assert np.isfinite(r.overall_att) + assert np.isfinite(r.overall_se) and r.overall_se > 0 + + def test_main_fit_zero_weight_treated_unit_covariates(self): + """Main fit with a zero-weight treated unit + covariates. + + Before: opaque ValueError from solve_ols (NaN in demeaned design). + After: fit succeeds; the zero-weight unit's FE is NaN, its cohort + cell (it is the ONLY cohort-2 unit) goes NaN across ALL inference + fields, and the overall ATT stays finite. + """ + rng = np.random.default_rng(11) + n_units, n_periods = 30, 6 + units = np.repeat(np.arange(n_units), n_periods) + times = np.tile(np.arange(n_periods), n_units) + first_treat = np.zeros(n_units, dtype=int) + first_treat[0] = 2 # the ONLY cohort-2 unit — will carry zero weight + first_treat[10:20] = 3 + ft = np.repeat(first_treat, n_periods) + x = rng.standard_normal((len(units), 2)) + post = (ft > 0) & (times >= ft) + outcome = ( + np.repeat(rng.standard_normal(n_units), n_periods) + + 0.2 * times + + 1.5 * post + + x @ np.array([0.6, -0.3]) + + rng.standard_normal(len(units)) * 0.3 + ) + data = pd.DataFrame( + { + "unit": units, + "time": times, + "outcome": outcome, + "first_treat": ft, + "x1": x[:, 0], + "x2": x[:, 1], + "w": np.where(units == 0, 0.0, 1.0), + } + ) + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + r = ImputationDiD().fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + covariates=["x1", "x2"], + survey_design=SurveyDesign(weights="w"), + aggregate="group", + ) + from tests.conftest import assert_nan_inference + + assert np.isfinite(r.overall_att) + assert r.group_effects is not None + # Cohort 2 contains only the zero-weight unit: NaN across ALL fields. + assert_nan_inference(r.group_effects[2]) + assert np.isnan(r.group_effects[2]["effect"]) + # Cohort 3 is unaffected. + assert np.isfinite(r.group_effects[3]["effect"]) diff --git a/tests/test_two_stage.py b/tests/test_two_stage.py index 68dc4d74..21e3af80 100644 --- a/tests/test_two_stage.py +++ b/tests/test_two_stage.py @@ -1414,31 +1414,40 @@ def test_iterative_fe_no_warning_on_convergence(self): est._iterative_fe(y, units, times, idx) assert not any("did not converge" in str(x.message) for x in w) - def test_iterative_demean_warns_on_nonconvergence(self): - """Silent-failure audit axis B: _iterative_demean must warn when max_iter exhausts.""" - rng = np.random.default_rng(42) - n_units, n_periods = 8, 5 - units = np.repeat(np.arange(n_units), n_periods) - times = np.tile(np.arange(n_periods), n_units) - vals = rng.standard_normal(n_units * n_periods) - idx = pd.RangeIndex(len(vals)) - - with pytest.warns(UserWarning, match="did not converge"): - TwoStageDiD._iterative_demean(vals, units, times, idx, max_iter=1, tol=1e-15) - - def test_iterative_demean_no_warning_on_convergence(self): - """Silent-failure audit axis B: no warning on well-behaved convergent input.""" + # NOTE (intentional coverage narrowing): the direct `_iterative_demean` + # non-convergence warn tests were retired with the method itself - the + # covariate within-transform now routes through the shared MAP engine + # with max_iter=10_000 hardcoded at the call site, so demean + # non-convergence is no longer forceable THROUGH this estimator. + # Engine-level warning coverage lives in + # tests/test_utils.py::TestDemeanByGroups; the `_iterative_fe` warn + # tests above still exercise this estimator's FE-solver warning path. + + def test_iterative_fe_zero_weight_unit_gets_nan_fe(self): + """A unit whose rows ALL carry zero weight surfaces as NaN FE. + + Locks the shared-solver zero-weight contract (spillover precedent: + never a silent finite 0.0) AND clean convergence - the historical + pandas loop divided 0/0 there and burned max_iter iterations. + """ rng = np.random.default_rng(42) n_units, n_periods = 8, 5 units = np.repeat(np.arange(n_units), n_periods) times = np.tile(np.arange(n_periods), n_units) - vals = rng.standard_normal(n_units * n_periods) - idx = pd.RangeIndex(len(vals)) + y = rng.standard_normal(n_units * n_periods) + w = np.ones(n_units * n_periods) + w[units == 2] = 0.0 + idx = pd.RangeIndex(len(y)) + est = TwoStageDiD() - with warnings.catch_warnings(record=True) as w: + with warnings.catch_warnings(record=True) as rec: warnings.simplefilter("always") - TwoStageDiD._iterative_demean(vals, units, times, idx) - assert not any("did not converge" in str(x.message) for x in w) + unit_fe, time_fe = est._iterative_fe(y, units, times, idx, weights=w) + assert not any("did not converge" in str(x.message) for x in rec) + + assert np.isnan(unit_fe[2]) # zero-weight unit: NaN, key retained + assert all(np.isfinite(v) for u, v in unit_fe.items() if u != 2) + assert all(np.isfinite(v) for v in time_fe.values()) # --------------------------------------------------------------------------- @@ -2368,3 +2377,87 @@ def test_n_clusters_counts_nan_cluster_like_the_variance(self): # so G strictly exceeds the distinct non-NaN cluster count. assert r.n_clusters > n_valid assert r.to_dict()["n_clusters"] == expected_g + + +# ============================================================================= +# TestZeroWeightGroups — shared-engine migration behavioral locks +# ============================================================================= + + +class TestZeroWeightGroups: + """Zero-total-weight groups on the Stage-1 paths (shared-engine migration). + + JK1/plain-BRR replicate weights zero whole PSUs and reach Stage 1 + unmasked (keep_mask only drops always-treated units). Before the + shared-engine migration the pandas loops divided 0/0 there: with + covariates, y_dm/X_dm NaN-poisoned, EVERY replicate refit failed inside + solve_ols(check_finite=True), and the fit returned NaN SEs after a + non-convergence warning storm. These tests lock the fixed contract, + including the warn_nan=False suppression of the per-replicate + "non-finite imputed outcomes" warning (main-fit warning unchanged). + """ + + @staticmethod + def _with_covariates(data, seed=7): + rng = np.random.default_rng(seed) + d = data.copy() + x = rng.standard_normal((len(d), 2)) + d["x1"], d["x2"] = x[:, 0], x[:, 1] + d["outcome"] = d["outcome"] + x @ np.array([0.6, -0.3]) + return d + + def test_replicate_covariates_zero_weight_psus_finite_se(self): + """Covariates + JK1 zeroed-PSU replicates -> finite SE, no warning storm.""" + data, rep_cols = _add_survey_cols(generate_test_data(n_units=60, seed=21)) + data = self._with_covariates(data) + design = SurveyDesign(weights="w", replicate_weights=rep_cols, replicate_method="JK1") + with warnings.catch_warnings(record=True) as rec: + warnings.simplefilter("always") + r = TwoStageDiD().fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + covariates=["x1", "x2"], + survey_design=design, + ) + messages = [str(w.message) for w in rec] + assert not any("replicate refits failed" in m for m in messages) + assert not any("did not converge" in m for m in messages) + # Main-fit weights are all positive here, and the per-replicate + # nan-ytilde warning is suppressed (warn_nan=False): zero copies. + assert not any("non-finite imputed outcomes" in m for m in messages) + assert np.isfinite(r.overall_att) + assert np.isfinite(r.overall_se) and r.overall_se > 0 + + def test_main_fit_zero_weight_unit_warns_once_and_fits(self): + """Main-fit zero-weight treated unit: nan-ytilde warning UNCHANGED. + + The warn_nan=False suppression applies ONLY inside replicate-refit + closures - a main fit that produces NaN y_tilde (zero-weight unit -> + NaN FE) must still surface the practitioner-facing warning, and the + fit must succeed (before the migration it raised an opaque + ValueError from solve_ols on the NaN-poisoned demeaned design). + """ + data = self._with_covariates(generate_test_data(n_units=40, seed=22)) + data["w"] = 1.0 + treated_units = data.loc[data["first_treat"] > 0, "unit"].unique() + data.loc[data["unit"] == treated_units[0], "w"] = 0.0 + with warnings.catch_warnings(record=True) as rec: + warnings.simplefilter("always") + r = TwoStageDiD().fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + covariates=["x1", "x2"], + survey_design=SurveyDesign(weights="w"), + ) + nan_ytilde = [ + m for m in (str(w.message) for w in rec) if "non-finite imputed outcomes" in m + ] + assert len(nan_ytilde) >= 1 # main-fit warning still fires + assert np.isfinite(r.overall_att) + assert np.isfinite(r.overall_se) and r.overall_se > 0 diff --git a/tests/test_utils.py b/tests/test_utils.py index ae214456..91e1ac5e 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -20,6 +20,7 @@ from diff_diff.utils import ( _compute_outcome_changes, + _iterative_fe_solve, _project_simplex, check_parallel_trends, check_parallel_trends_robust, @@ -1918,3 +1919,148 @@ def test_identified_low_variation_covariate_not_snapped(self): ) assert snapped == [] np.testing.assert_array_equal(out["xnear_dm"].values, before) + + +# ============================================================================= +# TestIterativeFeSolve — shared bincount Gauss-Seidel FE solver +# ============================================================================= + + +def _frozen_pandas_iterative_fe(y, unit_vals, time_vals, max_iter=10_000, tol=1e-10, weights=None): + """Byte-for-byte copy of the PRE-3.7 per-estimator pandas FE loop + (ImputationDiD/TwoStageDiD `_iterative_fe`, retired by the shared-engine + migration). NOT bit-identical to the bincount solver (pandas' grouped + mean is Kahan-compensated; bincount is naive accumulation). Kept as a + drift-bound guard: the shared solver must agree with this loop at + atol=1e-9. Positive weights only — the old loop divides 0/0 on + zero-total-weight groups (the bug the migration fixed). + """ + idx = pd.RangeIndex(len(y)) + n = len(y) + alpha = np.zeros(n) + beta = np.zeros(n) + 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 + with np.errstate(invalid="ignore", divide="ignore"): + for _ 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 + ) + 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: + break + 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() + return unit_fe, time_fe + + +class TestIterativeFeSolve: + """Drift-bound + contract tests for the shared bincount FE solver.""" + + @staticmethod + def _panel(kind, seed=42): + rng = np.random.default_rng(seed) + n_units, n_periods = 12, 7 + units, times = [], [] + for i in range(n_units): + for t in range(n_periods): + if kind == "unbalanced" and rng.random() < 0.25: + continue + units.append(i) + times.append(t) + units = np.asarray(units) + times = np.asarray(times) + y = ( + rng.standard_normal(n_units)[units] + + np.linspace(0, 1, n_periods)[times] + + rng.standard_normal(len(units)) * 0.3 + ) + return y, units, times + + def _solve_shared(self, y, units, times, weights=None): + unit_codes, unit_uniques = pd.factorize(units, sort=False) + time_codes, time_uniques = pd.factorize(times, sort=False) + u_arr, t_arr = _iterative_fe_solve( + y, + unit_codes.astype(np.intp), + time_codes.astype(np.intp), + len(unit_uniques), + len(time_uniques), + weights=weights, + method_name="test FE solver", + ) + return dict(zip(unit_uniques, u_arr)), dict(zip(time_uniques, t_arr)) + + @pytest.mark.parametrize("kind", ["balanced", "unbalanced"]) + @pytest.mark.parametrize("weighted", [False, True]) + def test_matches_frozen_pandas_loop(self, kind, weighted): + y, units, times = self._panel(kind) + rng = np.random.default_rng(7) + w = 0.5 + rng.exponential(0.5, len(y)) if weighted else None + + u_new, t_new = self._solve_shared(y, units, times, weights=w) + u_old, t_old = _frozen_pandas_iterative_fe(y, units, times, weights=w) + + assert set(u_new) == set(u_old) and set(t_new) == set(t_old) + for k in u_old: + np.testing.assert_allclose(u_new[k], u_old[k], rtol=0, atol=1e-9) + for k in t_old: + np.testing.assert_allclose(t_new[k], t_old[k], rtol=0, atol=1e-9) + + def test_zero_weight_group_nan_fe_and_convergence(self): + """Zero-total-weight unit: NaN FE, key retained, clean convergence.""" + y, units, times = self._panel("unbalanced") + w = np.ones(len(y)) + w[units == 5] = 0.0 + with warnings.catch_warnings(record=True) as rec: + warnings.simplefilter("always") + u_fe, t_fe = self._solve_shared(y, units, times, weights=w) + assert not any("did not converge" in str(x.message) for x in rec) + assert 5 in u_fe and np.isnan(u_fe[5]) + assert all(np.isfinite(v) for k, v in u_fe.items() if k != 5) + assert all(np.isfinite(v) for v in t_fe.values()) + # Positive-weight FE are unchanged by the zero-weight rows' presence + # (they are outside the WLS estimating sample). + keep = units != 5 + u_sub, t_sub = self._solve_shared(y[keep], units[keep], times[keep]) + for k, v in u_sub.items(): + np.testing.assert_allclose(u_fe[k], v, rtol=0, atol=1e-12) + + def test_warns_on_nonconvergence_with_label(self): + y, units, times = self._panel("unbalanced") + unit_codes, unit_uniques = pd.factorize(units, sort=False) + time_codes, time_uniques = pd.factorize(times, sort=False) + with pytest.warns(UserWarning, match="my solver label did not converge"): + _iterative_fe_solve( + y, + unit_codes.astype(np.intp), + time_codes.astype(np.intp), + len(unit_uniques), + len(time_uniques), + max_iter=1, + tol=1e-15, + method_name="my solver label", + )