"""Stage-2 aggregation + GMM variance engine for the Gardner two-stage estimator.

Extracted verbatim from ``diff_diff/two_stage.py`` (and, for the two shared

static helpers, ``two_stage_bootstrap.py``) for the M-022/M-119 post-fit

``aggregate()`` migration: ``two_stage.py`` imports ``two_stage_results.py``

and ``two_stage_bootstrap.py`` (which imports ``two_stage_results.py`` too),

so the results module can import neither -- the shared machinery lives here,

an import-leaf module both sides can reach (the

``efficient_did_aggregation.py`` / ``imputation_aggregation.py`` precedent).

Contents:

- module helpers ``_SPARSE_DENSE_THRESHOLD``, ``_LSMRUnconvergedError`` and

``_lsmr_certified_normal_solve`` (``two_stage.py`` re-imports all three --

``spillover.py`` and the bootstrap module's lazy imports keep working);

- :class:`_TwoStageAggregationMixin` -- the three Stage-2 aggregation levels

(static / event-study / group), the joint GMM sandwich they recompute

through, the Stage-1 helpers the replicate replay refits with, and the

replicate-weight inference override replay. Inherited by ``TwoStageDiD``

(fit-time behavior byte-identical) and hosted post-fit by the throwaway

``_TwoStageKitAggregator`` (``two_stage_results.py``).

"""

import warnings

from typing import Any, Dict, List, Optional, Set, Tuple

import numpy as np

import pandas as pd

from scipy import sparse

from scipy.sparse.linalg import factorized as sparse_factorized

from diff_diff.linalg import _rank_guarded_inv, solve_ols

from diff_diff.utils import _iterative_fe_solve, demean_by_groups, safe_inference

# Maximum number of elements before falling back to per-column sparse aggregation.

# 10M float64 elements ≈ 80 MB peak allocation. Above this, per-column .getcol()

# trades throughput for bounded memory. Keep in sync with two_stage_bootstrap.py.

_SPARSE_DENSE_THRESHOLD = 10_000_000

class _LSMRUnconvergedError(RuntimeError):

"""LSMR could not certify the Stage-1 normal-equation solve; the

variance boundary converts this to NaN inference (fail-closed)."""

def _lsmr_certified_normal_solve(

gram_csc, rhs: np.ndarray, context: str = "TwoStageDiD GMM sandwich"

) -> np.ndarray:

"""Least-squares solve of the (possibly singular) sparse Stage-1 Gram

system ``gram @ out = rhs`` via per-column LSMR — no dense

materialization of the ``(p_1, p_1)`` normal matrix (`O((U+T+K)^2)`

OOM risk on large panels; the pattern the ImputationDiD LSMR fix

closed, ported here after the consumer-invariance analysis).

OUTPUT-PRESERVING despite the min-norm ambiguity on singular systems:

least-squares solutions differ only by a ``null(X'X) = null(X_10)``

component (weighted: ``null(X'WX) = null(W^{1/2}X_10)`` — zero-weight

rows are inert in every weighted consumer because Psi/score/residual

contributions carry the same ``W`` factor), and EVERY ``gamma_hat``

consumer is an ``X_10``-range functional. One ``theta_exact`` consumer

(the bootstrap exact-residual helper's ``X_1_sparse @ theta_exact``)

evaluates theta on TREATED rows where a ``null(X_10)`` component would

NOT annihilate — parity there holds for a second reason: both dense

``lstsq`` (SVD) and LSMR return the MIN-NORM least-squares solution, so

the two solvers agree on the whole vector (to iterative tolerance), not

just on range functionals; the fit-level singular-design parity test

locks this at the V/SE level. The remaining consumers are the

``X_10``-range functionals — ``Psi_stage1 = X_10 @ gamma_hat``, the GMM

score correction ``c_g' gamma_hat`` with ``c_g = X_{10,g}' eps_{10,g}``

in ``rowspace(X_10)``, and Stage-1 residuals ``y - X_10 theta`` — so

every null component annihilates. Locked by the singular-system parity

test against a dense-lstsq oracle.

CONVERGENCE IS VALIDATED (fail-closed): ``istop`` in ``{0, 1, 2, 4, 5}``

certifies a solution / least-squares solution within tolerance (4/5 are

the machine-precision analogues of 1/2 per SciPy); anything else gets

ONE retry with an uncapped condition limit, then raises

:class:`_LSMRUnconvergedError` — converted to NaN inference at the

variance boundary rather than feeding an unverified solution into the

GMM sandwich.

"""

import scipy.sparse.linalg as spla

_certified = (0, 1, 2, 4, 5)

rhs_2d = np.atleast_2d(np.asarray(rhs, dtype=np.float64))

if rhs_2d.shape[0] == 1 and np.asarray(rhs).ndim == 1:

rhs_2d = rhs_2d.T

dim = gram_csc.shape[0]

out = np.empty((dim, rhs_2d.shape[1]))

for j in range(rhs_2d.shape[1]):

result = spla.lsmr(gram_csc, rhs_2d[:, j], atol=1e-14, btol=1e-14)

z, istop = result[0], int(result[1])

if istop not in _certified or not np.all(np.isfinite(z)):

result = spla.lsmr(

gram_csc,

rhs_2d[:, j],

atol=1e-14,

btol=1e-14,

conlim=1e16,

maxiter=max(50 * dim, 10_000),

)

z, istop = result[0], int(result[1])

if istop not in _certified or not np.all(np.isfinite(z)):

warnings.warn(

f"{context}: the LSMR fallback solve of the "

f"Stage-1 normal equations did not converge (istop={istop}); "

"the affected variance is reported as NaN rather than from "

"an unverified solution.",

UserWarning,

stacklevel=3,

)

raise _LSMRUnconvergedError(f"LSMR uncertified (istop={istop})")

out[:, j] = z

return out

class _TwoStageAggregationMixin:

"""Shared Stage-2/GMM methods (moved verbatim from ``TwoStageDiD``).

HOST-ATTRIBUTE CONTRACT -- the complete ``self.`` surface the moved

methods read (typed class-level declarations for ``mypy diff_diff`` at

zero errors on both hosts). Zero methods WRITE to ``self`` -- the

post-fit throwaway host exists for estimator-mutation isolation only.

"""

alpha: float

pretrends: bool

horizon_max: Optional[int]

rank_deficient_action: str

def _iterative_fe(

self,

y: np.ndarray,

unit_vals: np.ndarray,

time_vals: np.ndarray,

idx: pd.Index,

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 (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

-------

unit_fe : dict

Mapping from unit -> unit fixed effect.

time_fe : dict

Mapping from time -> time fixed effect.

"""

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

def _fit_untreated_model(

self,

df: pd.DataFrame,

outcome: str,

unit: str,

time: str,

covariates: Optional[List[str]],

omega_0_mask: pd.Series,

weights: Optional[np.ndarray] = None,

) -> Tuple[

Dict[Any, float], Dict[Any, float], float, Optional[np.ndarray], Optional[np.ndarray]

]:

"""

Stage 1: Estimate unit + time FE on untreated observations.

Parameters

----------

weights : np.ndarray, optional

Full-panel survey weights (same length as df). The untreated subset

is extracted internally via omega_0_mask. When None, unweighted.

Returns

-------

unit_fe, time_fe, grand_mean, delta_hat, kept_cov_mask

"""

df_0 = df.loc[omega_0_mask]

w_0 = weights[omega_0_mask.values] if weights is not None else None

if covariates is None or len(covariates) == 0:

y = df_0[outcome].values.copy()

unit_fe, time_fe = self._iterative_fe(

y, df_0[unit].values, df_0[time].values, df_0.index, weights=w_0

)

return unit_fe, time_fe, 0.0, None, None

else:

y = df_0[outcome].values.copy()

X_raw = df_0[covariates].values.copy()

units = df_0[unit].values

times = df_0[time].values

# 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,

y_dm,

return_vcov=False,

rank_deficient_action=self.rank_deficient_action,

column_names=covariates,

weights=w_0,

)

delta_hat = result[0]

kept_cov_mask = np.isfinite(delta_hat)

delta_hat_clean = np.where(np.isfinite(delta_hat), delta_hat, 0.0)

y_adj = y - np.dot(X_raw, delta_hat_clean)

unit_fe, time_fe = self._iterative_fe(y_adj, units, times, df_0.index, weights=w_0)

return unit_fe, time_fe, 0.0, delta_hat_clean, kept_cov_mask

def _residualize(

self,

df: pd.DataFrame,

outcome: str,

unit: str,

time: str,

covariates: Optional[List[str]],

unit_fe: Dict[Any, float],

time_fe: Dict[Any, float],

grand_mean: float,

delta_hat: Optional[np.ndarray],

) -> np.ndarray:

"""

Compute residualized outcome y_tilde for ALL observations.

y_tilde_i = y_i - mu_hat_i - eta_hat_t [- X_i @ delta_hat]

"""

alpha_i = df[unit].map(unit_fe).values

beta_t = df[time].map(time_fe).values

# Handle missing FE (NaN for units/periods not in untreated sample)

alpha_i = np.where(pd.isna(alpha_i), np.nan, alpha_i).astype(float)

beta_t = np.where(pd.isna(beta_t), np.nan, beta_t).astype(float)

y_hat = grand_mean + alpha_i + beta_t

if delta_hat is not None and covariates:

y_hat = y_hat + np.dot(df[covariates].values, delta_hat)

y_tilde = df[outcome].values - y_hat

return y_tilde

@staticmethod

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). ``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())

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

def _stage2_static(

self,

df: pd.DataFrame,

unit: str,

time: str,

first_treat: str,

covariates: Optional[List[str]],

omega_0_mask: pd.Series,

omega_1_mask: pd.Series,

unit_fe: Dict[Any, float],

time_fe: Dict[Any, float],

grand_mean: float,

delta_hat: Optional[np.ndarray],

cluster_var: str,

kept_cov_mask: Optional[np.ndarray],

survey_weights: Optional[np.ndarray] = None,

survey_weight_type: str = "pweight",

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.

Returns (att, se).

"""

y_tilde = df["_y_tilde"].values.copy()

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)

D[nan_mask] = 0.0

# X_2: treatment indicator (no intercept)

X_2 = D.reshape(-1, 1)

# Avoid degenerate case where all treated obs have NaN y_tilde

if D.sum() == 0:

return np.nan, np.nan

# Stage 2 OLS for point estimate (discard naive SE)

coef, residuals, _ = solve_ols(

X_2,

y_tilde,

return_vcov=False,

weights=survey_weights,

weight_type=survey_weight_type,

)

att = float(coef[0])

# GMM sandwich variance

# An uncertified LSMR Stage-1 fallback solve fails closed:

# NaN vcov -> NaN SE/t/p/CI (the helper already warned).

try:

V = self._compute_gmm_variance(

df=df,

unit=unit,

time=time,

covariates=covariates,

omega_0_mask=omega_0_mask,

unit_fe=unit_fe,

time_fe=time_fe,

delta_hat=delta_hat,

kept_cov_mask=kept_cov_mask,

X_2=X_2,

cluster_ids=df[cluster_var].values,

survey_weights=survey_weights,

resolved_survey=resolved_survey,

score_pad_mask=score_pad_mask,

cluster_ids_full=cluster_ids_full,

)

except _LSMRUnconvergedError:

V = np.full((X_2.shape[1], X_2.shape[1]), np.nan)

se = float(np.sqrt(max(V[0, 0], 0.0)))

return att, se

def _stage2_event_study(

self,

df: pd.DataFrame,

unit: str,

time: str,

first_treat: str,

covariates: Optional[List[str]],

omega_0_mask: pd.Series,

omega_1_mask: pd.Series,

unit_fe: Dict[Any, float],

time_fe: Dict[Any, float],

grand_mean: float,

delta_hat: Optional[np.ndarray],

cluster_var: str,

treatment_groups: List[Any],

ref_period: int,

balance_e: Optional[int],

kept_cov_mask: Optional[np.ndarray],

survey_weights: Optional[np.ndarray] = None,

survey_weight_type: str = "pweight",

survey_df: Optional[int] = None,

resolved_survey=None,

score_pad_mask: Optional[np.ndarray] = None,

cluster_ids_full: Optional[np.ndarray] = None,

warn_nan: bool = True,

) -> Tuple[Dict[int, Dict[str, Any]], Optional[np.ndarray], Optional[List[int]]]:

"""Event study Stage 2: OLS of y_tilde on relative-time dummies.

Returns ``(effects, vcov, vcov_index)``: the per-horizon effects

dict, the full GMM variance-covariance matrix over the ESTIMATED

horizon coefficients, and the horizon labels ordering its

rows/columns. The reference period and Proposition-5 horizons are

never regression columns, so they appear in ``effects`` but not in

``vcov_index``; all-filtered horizons (n_obs == 0) ARE columns,

with NaN-filled rows/columns from the rank guard. ``(dict, None,

None)`` on the degenerate early returns that fit no Stage-2

regression.

"""

y_tilde = df["_y_tilde"].values.copy()

nan_mask = self._mask_nan_ytilde(y_tilde, warn=warn_nan)

rel_times = df["_rel_time"].values

n = len(df)

# Get all horizons — include pre-periods when pretrends=True

if self.pretrends:

evt_rel = rel_times[~df["_never_treated"].values]

else:

evt_rel = rel_times[omega_1_mask.values]

all_horizons = sorted(set(int(h) for h in evt_rel if np.isfinite(h)))

# Apply horizon_max filter

if self.horizon_max is not None:

all_horizons = [h for h in all_horizons if abs(h) <= self.horizon_max]

# Apply balance_e filter

if balance_e is not None:

cohort_rel_times = self._build_cohort_rel_times(df, first_treat)

balanced_cohorts = set()

if all_horizons:

max_h = max(all_horizons)

required_range = set(range(-balance_e, max_h + 1))

for g, horizons in cohort_rel_times.items():

if required_range.issubset(horizons):

balanced_cohorts.add(g)

if not balanced_cohorts:

warnings.warn(

f"No cohorts satisfy balance_e={balance_e} requirement. "

"Event study results will contain only the reference period. "

"Consider reducing balance_e.",

UserWarning,

stacklevel=2,

)

return (

{

ref_period: {

"effect": 0.0,

"se": 0.0,

"t_stat": np.nan,

"p_value": np.nan,

"conf_int": (0.0, 0.0),

"n_obs": 0,

}

},

None,

None,

)

balance_mask = df[first_treat].isin(balanced_cohorts).values

else:

balance_mask = np.ones(n, dtype=bool)

# Check Proposition 5: no never-treated units

has_never_treated = df["_never_treated"].any()

h_bar = np.inf

if not has_never_treated and len(treatment_groups) > 1:

h_bar = max(treatment_groups) - min(treatment_groups)

# Identify Prop 5 horizons and compute their actual treated obs counts.

# Treated obs have NaN y_tilde at these horizons (counterfactual

# unidentified), but actual_n counts them to distinguish from truly

# empty horizons. rel_times is NaN for untreated/never-treated obs

# (line ~653), so (rel_times == h) is False for them.

prop5_horizons = []

prop5_effects: Dict[int, Dict[str, Any]] = {}

if h_bar < np.inf:

for h in all_horizons:

if h == ref_period:

continue

if h >= h_bar:

actual_n = int(np.sum((rel_times == h) & omega_1_mask.values & balance_mask))

if actual_n > 0:

prop5_horizons.append(h)

prop5_effects[h] = {

"effect": np.nan,

"se": np.nan,

"t_stat": np.nan,

"p_value": np.nan,

"conf_int": (np.nan, np.nan),

"n_obs": actual_n,

}

# Remove reference period AND Prop 5 horizons from estimation

prop5_set = set(prop5_horizons)

est_horizons = [h for h in all_horizons if h != ref_period and h not in prop5_set]

if len(est_horizons) == 0:

# No horizons to estimate — return the reference row PLUS any

# Proposition-5 rows (local-review fix, 2(b) PR-3b): when EVERY

# non-reference horizon is Prop-5-unidentified, the rows must

# still surface as all-NaN with n_obs > 0 and the consolidated

# warning, exactly as on the normal path below — dropping them

# here reported real treated horizons as absent instead of

# unidentified (contra REGISTRY Prop-5 contract).

if prop5_horizons:

warnings.warn(

f"Horizons {prop5_horizons} are not identified without "

f"never-treated units (Proposition 5). Set to NaN.",

UserWarning,

stacklevel=2,

)

ref_only: Dict[int, Dict[str, Any]] = {

ref_period: {

"effect": 0.0,

"se": 0.0,

"t_stat": np.nan,

"p_value": np.nan,

"conf_int": (0.0, 0.0),

"n_obs": 0,

}

}

ref_only.update(prop5_effects)

return (ref_only, None, None)

# Build Stage 2 design: one column per horizon (no intercept)

# Never-treated obs get all-zero rows (undefined relative time -> NaN)

# With no intercept, they contribute zero to X'_2 X_2 and X'_2 y_tilde

horizon_to_col = {h: j for j, h in enumerate(est_horizons)}

k = len(est_horizons)

X_2 = np.zeros((n, k))

for i in range(n):

if not balance_mask[i]:

continue

if nan_mask[i]:

continue # NaN y_tilde -> don't include in event study

h = rel_times[i]

if np.isfinite(h):

h_int = int(h)

if h_int in horizon_to_col:

X_2[i, horizon_to_col[h_int]] = 1.0

# Stage 2 OLS

coef, residuals, _ = solve_ols(

X_2,

y_tilde,

return_vcov=False,

weights=survey_weights,

weight_type=survey_weight_type,

)

# GMM variance for full coefficient vector

# An uncertified LSMR Stage-1 fallback solve fails closed:

# NaN vcov -> NaN SE/t/p/CI (the helper already warned).

try:

V = self._compute_gmm_variance(

df=df,

unit=unit,

time=time,

covariates=covariates,

omega_0_mask=omega_0_mask,

unit_fe=unit_fe,

time_fe=time_fe,

delta_hat=delta_hat,

kept_cov_mask=kept_cov_mask,

X_2=X_2,

cluster_ids=df[cluster_var].values,

survey_weights=survey_weights,

resolved_survey=resolved_survey,

score_pad_mask=score_pad_mask,

cluster_ids_full=cluster_ids_full,

)

except _LSMRUnconvergedError:

V = np.full((X_2.shape[1], X_2.shape[1]), np.nan)

# Build results dict

event_study_effects: Dict[int, Dict[str, Any]] = {}

# Reference period marker

event_study_effects[ref_period] = {

"effect": 0.0,

"se": 0.0,

"t_stat": np.nan,

"p_value": np.nan,

"conf_int": (0.0, 0.0),

"n_obs": 0,

}

for h in est_horizons:

j = horizon_to_col[h]

n_obs = int(np.sum(X_2[:, j]))

if n_obs == 0:

event_study_effects[h] = {

"effect": np.nan,

"se": np.nan,

"t_stat": np.nan,

"p_value": np.nan,

"conf_int": (np.nan, np.nan),

"n_obs": 0,

}

continue

effect = float(coef[j])

se = float(np.sqrt(max(V[j, j], 0.0)))

t_stat, p_val, ci = safe_inference(effect, se, alpha=self.alpha, df=survey_df)

event_study_effects[h] = {

"effect": effect,

"se": se,

"t_stat": t_stat,

"p_value": p_val,

"conf_int": ci,

"n_obs": n_obs,

}

# Add Proposition 5 entries (unidentified horizons with n_obs > 0)

event_study_effects.update(prop5_effects)

if prop5_horizons:

warnings.warn(

f"Horizons {prop5_horizons} are not identified without "

f"never-treated units (Proposition 5). Set to NaN.",

UserWarning,

stacklevel=2,

)

return event_study_effects, V, [int(h) for h in est_horizons]

def _stage2_group(

self,

df: pd.DataFrame,

unit: str,

time: str,

first_treat: str,

covariates: Optional[List[str]],

omega_0_mask: pd.Series,

omega_1_mask: pd.Series,

unit_fe: Dict[Any, float],

time_fe: Dict[Any, float],

grand_mean: float,

delta_hat: Optional[np.ndarray],

cluster_var: str,

treatment_groups: List[Any],

kept_cov_mask: Optional[np.ndarray],

survey_weights: Optional[np.ndarray] = None,

survey_weight_type: str = "pweight",

survey_df: Optional[int] = None,

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, warn=warn_nan)

n = len(df)

# Build Stage 2 design: one column per cohort (no intercept)

group_to_col = {g: j for j, g in enumerate(treatment_groups)}

k = len(treatment_groups)

X_2 = np.zeros((n, k))

ft_vals = df[first_treat].values

treated_mask = omega_1_mask.values

for i in range(n):

if treated_mask[i] and not nan_mask[i]:

g = ft_vals[i]

if g in group_to_col:

X_2[i, group_to_col[g]] = 1.0

# Stage 2 OLS

coef, residuals, _ = solve_ols(

X_2,

y_tilde,

return_vcov=False,

weights=survey_weights,

weight_type=survey_weight_type,

)

# GMM variance

# An uncertified LSMR Stage-1 fallback solve fails closed:

# NaN vcov -> NaN SE/t/p/CI (the helper already warned).

try:

V = self._compute_gmm_variance(

df=df,

unit=unit,

time=time,

covariates=covariates,

omega_0_mask=omega_0_mask,

unit_fe=unit_fe,

time_fe=time_fe,

delta_hat=delta_hat,

kept_cov_mask=kept_cov_mask,

X_2=X_2,

cluster_ids=df[cluster_var].values,

survey_weights=survey_weights,

resolved_survey=resolved_survey,

score_pad_mask=score_pad_mask,

cluster_ids_full=cluster_ids_full,

)

except _LSMRUnconvergedError:

V = np.full((X_2.shape[1], X_2.shape[1]), np.nan)

group_effects: Dict[Any, Dict[str, Any]] = {}

for g in treatment_groups:

j = group_to_col[g]

n_obs = int(np.sum(X_2[:, j]))

if n_obs == 0:

group_effects[g] = {

"effect": np.nan,

"se": np.nan,

"t_stat": np.nan,

"p_value": np.nan,

"conf_int": (np.nan, np.nan),

"n_obs": 0,

}

continue

effect = float(coef[j])

se = float(np.sqrt(max(V[j, j], 0.0)))

t_stat, p_val, ci = safe_inference(effect, se, alpha=self.alpha, df=survey_df)

group_effects[g] = {

"effect": effect,

"se": se,

"t_stat": t_stat,

"p_value": p_val,

"conf_int": ci,

"n_obs": n_obs,

}

return group_effects

@staticmethod

def _compute_gmm_scores(

c_by_cluster: np.ndarray,

gamma_hat: np.ndarray,

s2_by_cluster: np.ndarray,

) -> np.ndarray:

"""

Compute per-cluster GMM scores S_g = gamma_hat' c_g - X'_{2g} eps_{2g}.

Handles NaN/overflow from rank-deficient FE by wrapping in errstate

and replacing non-finite values with 0.

Parameters

----------

c_by_cluster : np.ndarray, shape (G, p)

Per-cluster Stage 1 scores.

gamma_hat : np.ndarray, shape (p, k)

Cross-moment correction matrix.

s2_by_cluster : np.ndarray, shape (G, k)

Per-cluster Stage 2 scores.

Returns

-------

np.ndarray, shape (G, k)

Per-cluster influence scores.

"""

with np.errstate(invalid="ignore", divide="ignore", over="ignore"):

correction = np.dot(c_by_cluster, gamma_hat)

np.nan_to_num(correction, copy=False, nan=0.0, posinf=0.0, neginf=0.0)

return correction - s2_by_cluster

def _compute_gmm_variance(

self,

df: pd.DataFrame,

unit: str,

time: str,

covariates: Optional[List[str]],

omega_0_mask: pd.Series,

unit_fe: Dict[Any, float],

time_fe: Dict[Any, float],

delta_hat: Optional[np.ndarray],

kept_cov_mask: Optional[np.ndarray],

X_2: np.ndarray,

cluster_ids: np.ndarray,

survey_weights: Optional[np.ndarray] = None,

resolved_survey=None,

score_pad_mask: Optional[np.ndarray] = None,

cluster_ids_full: Optional[np.ndarray] = None,

) -> np.ndarray:

"""

Compute GMM sandwich variance (Butts & Gardner 2022).

Matches the R `did2s` source code implementation: uses the GLOBAL

Hessian inverse (not per-cluster) and NO finite-sample adjustments.

The per-observation influence function is:

IF_i = (X'_2 X_2)^{-1} [gamma_hat' x_{10i} eps_{10i} - x_{2i} eps_{2i}]

where gamma_hat = (X'_{10} X_{10})^{-1} (X'_1 X_2) uses the GLOBAL

cross-moment.

The cluster-robust variance is:

V = (X'_2 X_2)^{-1} (sum_g S_g S'_g) (X'_2 X_2)^{-1}

S_g = gamma_hat' c_g - X'_{2g} eps_{2g}

c_g = X'_{10g} eps_{10g}

With survey weights W (diagonal):

Bread: (X'_2 W X_2)^{-1}

gamma_hat: (X'_{10} W X_{10})^{-1} (X'_1 W X_2)

c_g = sum_{i in g} w_i * x_{10i} * eps_{10i}

s2_g = sum_{i in g} w_i * x_{2i} * eps_{2i}

Parameters

----------

X_2 : np.ndarray, shape (n, k)

Stage 2 design matrix (treatment indicators). The Stage-2 residual

``eps_2`` is re-solved internally from the *exact* Stage-1 residuals

(see the exact-residual note below), so it is not a parameter.

cluster_ids : np.ndarray, shape (n,)

Cluster identifiers, fit-sample length. Used for the per-cluster

stage-1 / stage-2 score aggregation (OLS path).

survey_weights : np.ndarray, optional

Survey weights of shape (n,). When None, unweighted (identical

to current code).

resolved_survey : ResolvedSurveyDesign, optional

Resolved survey design. Under Wave E.3 parity (PR #482 SpilloverDiD

precedent) the design retains full-domain `n_psu` / `n_strata` /

`df_survey` / `strata` / `fpc` / `psu` arrays even when the

always-treated drop removes rows from the OLS sample. The

zero-padded per-cluster scores expand onto the full-domain PSU

list before stratified-meat dispatch. R `survey::svyrecvar(subset())`

convention (Lumley 2010 §2.5); mirrors `imputation.py:2175-2183`

(PreTrendsImputation) and `prep.py:1401-1432` (DCDH cell variance).

score_pad_mask : np.ndarray of shape (n_full,), bool, optional

Wave E.3 parity zero-pad mask. When supplied, indicates which

FULL-DOMAIN rows are present in the fit sample (True = kept

for OLS). Requires `n == int(np.sum(score_pad_mask))`. Co-supplied

with `cluster_ids_full`. Per-cluster stage-1 / stage-2 score

aggregates computed at fit-length are expanded onto the

full-domain unique-PSU list; PSUs absent from the fit sample

(e.g. PSUs containing only always-treated rows) get zero score

rows but still count toward `G_full` for `n_psu` / `df_survey`.

None (default) → no padding, exact pre-PR behavior.

cluster_ids_full : np.ndarray of shape (n_full,), optional

Full-domain PSU labels. Co-supplied with `score_pad_mask`. Must

share the same length. Provides the full-domain unique-PSU list

used both for score zero-pad expansion and for downstream

strata/FPC `obs_idx` lookups against the full-domain

`resolved_survey.strata` / `.fpc` arrays. None (default) → no

padding, exact pre-PR behavior.

Returns

-------

np.ndarray, shape (k, k)

Variance-covariance matrix.

"""

n = len(df)

k = X_2.shape[1]

# Exclude rank-deficient covariates

cov_list = covariates

if covariates and kept_cov_mask is not None and not np.all(kept_cov_mask):

cov_list = [c for c, k_ in zip(covariates, kept_cov_mask) if k_]

# Build sparse FE design matrices X_1 (all obs) and X_10 (untreated only)

X_1_sparse, X_10_sparse, unit_to_idx, time_to_idx = self._build_fe_design(

df, unit, time, cov_list, omega_0_mask

)

p = X_1_sparse.shape[1]

# eps_10 = Y - X_10 @ gamma_hat

# Untreated: stage 1 residual (Y - fitted). Treated: Y (X_10 rows = 0).

# Reconstruct Y from y_tilde: Y = y_tilde + fitted_stage1. Because

# y_tilde = Y - fitted_1, the iterative FE in fitted_1 cancel exactly, so

# y_vals == Y (independent of the iterative solver's tolerance).

alpha_i = df[unit].map(unit_fe).values

beta_t = df[time].map(time_fe).values

# Identification mask: obs whose unit AND time FE are both identified by the

# untreated Stage-1 fit. Rank-deficient / Proposition-5 obs (NaN FE) keep the

# iterative-residual behavior; only identified obs get the exact residuals.

identified = np.isfinite(np.asarray(alpha_i, dtype=float)) & np.isfinite(

np.asarray(beta_t, dtype=float)

)

alpha_i = np.where(pd.isna(alpha_i), 0.0, alpha_i).astype(float)

beta_t = np.where(pd.isna(beta_t), 0.0, beta_t).astype(float)

fitted_1 = alpha_i + beta_t

if delta_hat is not None and cov_list:

if kept_cov_mask is not None and not np.all(kept_cov_mask):

fitted_1 = fitted_1 + np.dot(df[cov_list].values, delta_hat[kept_cov_mask])

else:

fitted_1 = fitted_1 + np.dot(df[cov_list].values, delta_hat)

y_tilde = df["_y_tilde"].values

y_vals = y_tilde + fitted_1 # reconstruct Y

y_vals_clean = np.nan_to_num(y_vals, nan=0.0)

omega_0 = omega_0_mask.values

# 1. gamma_hat = (X'_{10} W X_{10})^{-1} (X'_1 W X_2) [p x k]

# With survey weights, both cross-products need W. We reuse the SAME

# factorization of (X'_{10} W X_{10}) to also solve the exact Stage-1 FE

# coefficients theta_exact (see exact-residual note below).

if survey_weights is not None:

XtWX_10 = X_10_sparse.T @ X_10_sparse.multiply(survey_weights[:, None])

Xt1_WX2 = X_1_sparse.T @ (X_2 * survey_weights[:, None])

rhs_fe = X_10_sparse.T @ (survey_weights * y_vals_clean)

else:

XtWX_10 = X_10_sparse.T @ X_10_sparse # (p x p) sparse

Xt1_WX2 = X_1_sparse.T @ X_2 # (p x k) dense

rhs_fe = X_10_sparse.T @ y_vals_clean # (p,) X'_{10} W Y

try:

solve_XtX = sparse_factorized(XtWX_10.tocsc())

if Xt1_WX2.ndim == 1:

gamma_hat = solve_XtX(Xt1_WX2).reshape(-1, 1)

else:

gamma_hat = np.column_stack(

[solve_XtX(Xt1_WX2[:, j]) for j in range(Xt1_WX2.shape[1])]

)

theta_exact = np.asarray(solve_XtX(np.asarray(rhs_fe).ravel())).ravel()

except RuntimeError as exc:

# Singular matrix — fall back to certified sparse LSMR. Silent-failure

# audit axis C: emit a UserWarning on fallback instead of swallowing.

warnings.warn(

"TwoStageDiD GMM sandwich: sparse factorization of "

f"(X'_{{10}} W X_{{10}}) failed ({type(exc).__name__}); falling "

"back to sparse LSMR. This may indicate a rank-deficient or "

"near-singular Stage 1 design matrix and SE estimates may be "

"less reliable.",

UserWarning,

stacklevel=2,

)

XtWX_10_csc = XtWX_10.tocsc()

gamma_hat = _lsmr_certified_normal_solve(XtWX_10_csc, Xt1_WX2)

theta_exact = _lsmr_certified_normal_solve(

XtWX_10_csc, np.asarray(rhs_fe).ravel()

).ravel()

# Exact Stage-1 / Stage-2 residuals. The point-estimate path uses the

# iterative alternating-projection FE solver (`_iterative_fe`), which

# converges only to ~1e-7 on unbalanced untreated panels; that error is

# negligible for the ATT but perturbs the variance by ~1% relative to the

# analytical GMM sandwich. The variance therefore re-solves the Stage-1 FE

# EXACTLY using the sparse normal equations already factorized for gamma_hat

# (theta_exact), matching R `did2s` to ~1e-7 and mirroring ImputationDiD's

# exact-sparse variance path. The shared `_exact_gmm_residuals` helper is

# used by BOTH this analytical path and the multiplier bootstrap

# (`_compute_cluster_S_scores`) so the influence function is single-sourced.

eps_10, eps_2 = self._exact_gmm_residuals(

X_1_sparse,

theta_exact,

y_vals_clean,

identified,

omega_0,

y_tilde,

X_2,

survey_weights,

)

# 2. Per-cluster Stage 1 scores: c_g = sum_{i in g} w_i * x_{10i} * eps_{10i}

# Only untreated obs have non-zero X_10 rows

# With survey weights: multiply eps_10 by survey_weights before sparse multiply

if survey_weights is not None:

weighted_eps_10 = survey_weights * eps_10

else:

weighted_eps_10 = eps_10

weighted_X10 = X_10_sparse.multiply(weighted_eps_10[:, None]) # sparse element-wise

unique_clusters, cluster_indices = np.unique(cluster_ids, return_inverse=True)

G = len(unique_clusters)

n_elements = weighted_X10.shape[0] * weighted_X10.shape[1]

c_by_cluster = np.zeros((G, p))

if n_elements > _SPARSE_DENSE_THRESHOLD:

# Per-column path: limits peak memory for large FE matrices

weighted_X10_csc = weighted_X10.tocsc()

for j_col in range(p):