Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
igerber merged 1 commit into from
Jul 4, 2026
+739 −342
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 32 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
`survey_design=SurveyDesign(psu=<cluster_col>)`. No behavior change for unclustered fits.

### Fixed
- **ImputationDiD/TwoStageDiD covariate fits with zero-weight replicate designs (JK1/plain
BRR) now produce finite SEs.** Replicate weights that zero out whole PSUs reach Step 1
unmasked; the previous per-estimator pandas demeaning loops divided 0/0 on zero-total-weight
groups, NaN-poisoning the demeaned design so EVERY replicate refit failed inside
`solve_ols` ("All replicate refits failed. Returning NaN variance." after a
non-convergence warning storm — measured: 198 warnings and SE=NaN on a 350k-row panel
with 16 JK1 replicates, now 0 warnings and a finite SE). A main fit combining zero-weight
rows with covariates previously raised an opaque `ValueError`; it now fits, with the
zero-weight unit's FE surfacing as NaN (spillover convention — keys retained for the
rank-condition check, never a silent finite 0.0) and its unidentified cells NaN across
all inference fields. TwoStageDiD's per-replicate "non-finite imputed outcomes" warning
is suppressed inside replicate-refit closures only (`warn_nan=False`); the main-fit
warning is unchanged.
- **Dube, Girardi, Jordà & Taylor (2025) citation corrected to *J. Applied Econometrics*
40(**7**):741-758** (was cited as issue 5) across `docs/references.rst`,
`docs/methodology/REGISTRY.md`, `docs/methodology/papers/dube-2025-review.md`,
Expand All @@ -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
Expand Down
3 changes: 2 additions & 1 deletion TODO.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
212 changes: 66 additions & 146 deletions diff_diff/imputation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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
-------
Expand All @@ -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,
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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
Expand Down
8 changes: 5 additions & 3 deletions diff_diff/spillover.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
Loading
Loading