diff --git a/CHANGELOG.md b/CHANGELOG.md index a669e3c4..4fdd3f40 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -118,6 +118,27 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 numerical, or public-API change. ### Fixed +- **Structural (non-covariate) matrix inverses are now rank-guarded.** The internal design-Gram + bread inversions in `ContinuousDiD` (ACRT-variance `Psi'WPsi`), `TwoStageDiD` (Stage-2 + `X_2'WX_2`, both the analytical and multiplier-bootstrap surfaces), `SpilloverDiD` (Wave D + `A_22`), and the Conley spatial-HAC variance (`X'WX`) sat on a `LinAlgError`-only fallback: + `np.linalg.inv`/`solve` raise only on an *exactly* singular matrix, so a **near**-singular + Gram returned a garbage inverse (~1e13) straight into the SE (and `ContinuousDiD`'s exact- + singular fallback was a *silent* minimum-norm `pinv`; Conley *raised* `ValueError`). All now + route through the shared `_rank_guarded_inv` (`diff_diff/linalg.py`) — the same generalized + inverse already used for the covariate IF SEs — which truncates redundant directions on the + equilibrated Gram to give a finite SE on the identified subspace (the well-conditioned + near-collinear limit, not minimum-norm; NaN only at rank 0) and emits a `UserWarning` when a + direction is dropped. For the per-coefficient reporters (TwoStageDiD event/group, SpilloverDiD + rings, Conley), a dropped (unidentified) named coefficient is reported with **NaN** SE (not the + zero-filled `0`); linear-combination reporters (ContinuousDiD dose curves) are unaffected, since + a dropped direction correctly contributes 0 there. These invert *internal* bases users cannot + perturb with `covariates=`. + **Behavior change:** a rank-deficient Conley design no longer raises — it rank-reduces with a + warning. Well-conditioned designs are unchanged (the fast path is `np.linalg.solve(A, I)`, + R-parity preserved). (`HeterogeneousAdoptionDiD`'s non-symmetric IV bread and `ImputationDiD` + — whose vcov is already rank-guarded upstream via `solve_ols` — were assessed and excluded; + see `TODO.md`.) - **`CallawaySantAnna` / `StaggeredTripleDifference` covariate outcome-regression is now scale-robust.** The covariate OR nuisance fits — `_compute_all_att_gt_covariate_reg` and `_doubly_robust` (CS) and `_compute_or` (StaggeredTripleDifference) — previously used an diff --git a/TODO.md b/TODO.md index 7300de2d..bcaea156 100644 --- a/TODO.md +++ b/TODO.md @@ -30,7 +30,6 @@ The `Origin` column (Actionable tables) and the `PR` column (Deferred tables) bo |-------|----------|--------|--------|----------| | `SyntheticControl` cv: thread an `"infeasible"` reason-code from `_outer_solve_V_cv()` / `_placebo_fit_unit()` so `in_space_placebo()` / `leave_one_out()` distinguish a structural cv-refit exclusion (donor-indistinguishable re-aggregated window) from a genuine inner-solver non-convergence — mirror the split `in_time_placebo()` already emits. Warnings already distinguish the two causes; only the machine-readable status/count is missing. | `synthetic_control.py`, `synthetic_control_results.py` | follow-up | Mid | Low | | `CallawaySantAnna`: materialize NaN entries for non-estimable `(g,t)` cells in `group_time_effects` (currently omitted with a consolidated warning); requires updating downstream consumers (event study, `balance_e`, aggregation). | `staggered.py` | #256 | Mid | Low | -| Adopt the shared `_rank_guarded_inv` for the *structural* (non-covariate) matrix inverses sharing the `LinAlgError`-only fallback that can go near-singular: `continuous_did.py:1056` (dose B-spline), `spillover.py:3371` (ring-solve, partially guarded), `two_stage.py:3154` (TSL Stage-2 variance), `imputation.py:2403`, `had.py:2413`, `conley.py:1109`. These invert internal bases users cannot perturb with `covariates=` (distinct from the already-fixed covariate-triggered SE bug); the helper is the seam. | `continuous_did.py`, `spillover.py`, `two_stage.py`, `imputation.py`, `had.py`, `conley.py` | dr-or-se-rank-guard | Mid | Low | | Survey-design resolution / collapse patterns are inconsistent across panel estimators — `ContinuousDiD` rebuilds unit-level design in SE code, `EfficientDiD` builds once in `fit()`, `StackedDiD` re-resolves on stacked data. Extract shared helpers for panel-to-unit collapse, post-filter re-resolution, and metadata recomputation. | `continuous_did.py`, `efficient_did.py`, `stacked_did.py` | #226 | Mid | Low | | `SyntheticControl` remaining ADH-2015 §4 items: the regression-weight `W^reg = X_0'(X_0 X_0')^{-1} X_1` extrapolation diagnostic (flag implied OLS weights outside `[0,1]`) and sparse-SC subset search (`l < J`, holding `V` fixed). LOO, in-time placebo, CV `V`-selection, and inverse-variance `V` have landed; these two are the deferred tail. | `synthetic_control.py`, `synthetic_control_results.py` | ADH-2015 | Mid | Low | | `SyntheticControl` conformal (CWZ 2021) extensions: (a) one-sided / signed-`t` variants (§7); (b) covariates in the conformal proxy (`X_jt`, eqs 4/6 — current proxy is outcomes-only); (c) AR / innovation-permutation path (Lemmas 5-7) for time-series proxies. The joint test, pointwise CIs, and average-effect CI have landed. | `conformal.py`, `synthetic_control_results.py` | CWZ-2021 | Heavy | Low | @@ -150,6 +149,8 @@ Doable in principle, but no current caller and/or explicitly out of paper scope. |----------|----------|----------| | **`bread_inv` reuse not bit-identically achievable.** "Factor `(X'WX)` once, reuse across HC2/HC2-BM" can't be done bit-identically (the bar for a pure perf refactor of the inference path). Internal bread ops solve against *different* RHS (`X.T`, `eye`, `meat`+`temp.T`, `contrasts`); only same-RHS results are bit-reusable. Measured: `lu_solve(lu_factor(A),B)` differs from `solve(A,B)` up to 6.4e-15; the `inv(A)@meat@inv(A)` sandwich differs up to 1.24e-14 — both nonzero and *below* the affected goldens' tolerances (1e-12/1e-10), so a broad reuse would silently shift SEs without tripping the suite. The one genuine bit-identical redundancy (a duplicated `solve(bread, X.T)` in the unweighted one-way `hc2_bm`+`return_dof` path) is dwarfed by that path's dense `M=I−H` build, so the saving is negligible. | `linalg.py::compute_robust_vcov` | 2026-06-01 | | **R-script-per-test consolidation has no CI impact.** No CI workflow installs R, so every R-parity test skips in CI behind a per-file availability gate — consolidating `Rscript` spawns yields zero CI speedup. `test_methodology_twfe.py` already session-caches its R fits. The only residual is a LOCAL-dev micro-opt for `test_methodology_continuous_did.py` / `test_methodology_callaway.py` (re-spawn `library(...)` per call). Low value; retained as a local-dev note. | `tests/test_methodology_continuous_did.py`, `tests/test_methodology_callaway.py` | #139 / 2026-06-07 | +| **`HeterogeneousAdoptionDiD` mass-point IV bread is non-symmetric — `_rank_guarded_inv` inapplicable.** The structural rank-guard sweep (continuous_did / two_stage / spillover / conley) excluded `had.py`'s `ZtWX = Zd'WX` ([1, instrument]' × [1, endogenous]): it is a non-symmetric 2×2 Wald-IV bread (`V = ZtWX_inv @ Omega @ ZtWX_inv.T`), and `_rank_guarded_inv` assumes a **symmetric PSD** Gram (symmetric `D=diag(A)` equilibration + eigendecomposition), so applying it would be methodologically wrong. The existing fallback already returns NaN SE on a singular bread; an IV-appropriate near-singular guard would need a different mechanism. | `had.py` | structural-rank-guard / 2026-06-28 | +| **`ImputationDiD` SE vcov is already rank-guarded upstream.** Excluded from the structural rank-guard sweep: the lead/effect vcov comes from `solve_ols(..., return_vcov=True, rank_deficient_action=...)` at the OLS fit (`imputation.py:~2316`), which already drops rank-deficient columns. The only raw inverse (`solve(V_gamma, gamma)`, `imputation.py:~2530`) is the pretrends **Wald F-test statistic** with a safe `NaN` fallback — a test statistic, not a sandwich bread — so there is no garbage-SE exposure. No structural rank-guard needed. | `imputation.py` | structural-rank-guard / 2026-06-28 | --- diff --git a/diff_diff/conley.py b/diff_diff/conley.py index b43e7c3b..7bde531a 100644 --- a/diff_diff/conley.py +++ b/diff_diff/conley.py @@ -1104,16 +1104,37 @@ def _compute_conley_vcov( _conley_sparse=_conley_sparse, ) - # Sandwich via two solves (mirrors _compute_cr2_bm pattern in linalg.py) - try: - temp = np.linalg.solve(bread_matrix, meat) - vcov = np.linalg.solve(bread_matrix, temp.T).T - except np.linalg.LinAlgError as e: - if "Singular" in str(e): - raise ValueError( - "Design matrix is rank-deficient (singular X'X matrix). " - "Cannot compute Conley spatial HAC variance." - ) from e - raise + # Sandwich via the shared rank-guarded inverse of the design Gram. + # np.linalg.solve only raises on an *exactly* singular bread, so a *near*- + # singular X'WX would otherwise flow a garbage inverse (~1e13) straight into + # the spatial-HAC variance. `_rank_guarded_inv` truncates redundant + # directions on the equilibrated Gram -> a finite SE on the identified + # subspace (NaN only at rank 0), matching the covariate IF rank-guard and the + # other structural bread inversions (ContinuousDiD / TwoStageDiD / + # SpilloverDiD). Lazy import: `linalg` imports this module, so a top-level + # `from diff_diff.linalg import ...` would be circular; resolving at call time + # is safe (linalg is already loaded by the time this runs). + from diff_diff.linalg import _rank_guarded_inv + + bread_inv, n_dropped, _, dropped = _rank_guarded_inv(bread_matrix, return_dropped=True) + if n_dropped: + warnings.warn( + "Conley spatial HAC variance: the design Gram (X'WX) is " + f"rank-deficient ({n_dropped} redundant direction(s) dropped); " + "rank-reducing to a finite SE on the identified subspace " + "(NaN if rank 0). This usually indicates collinear regressors.", + UserWarning, + stacklevel=2, + ) + # vcov = bread^{-1} @ meat @ bread^{-1}; algebraically identical to the prior + # two symmetric solves given `bread` symmetric (holds for any meat). + vcov = bread_inv @ meat @ bread_inv + # A dropped (unidentified) coefficient is zero-filled in bread_inv, which would + # otherwise report se=0 for that named coefficient. NaN its row/col in the + # FINAL vcov so per-coefficient SE extraction yields NaN (not 0) for the + # unidentified directions, while the identified coefficients stay finite. + if dropped.any(): + vcov[dropped, :] = np.nan + vcov[:, dropped] = np.nan return vcov diff --git a/diff_diff/continuous_did.py b/diff_diff/continuous_did.py index 8f433b25..afed61f7 100644 --- a/diff_diff/continuous_did.py +++ b/diff_diff/continuous_did.py @@ -29,7 +29,7 @@ ContinuousDiDResults, DoseResponseCurve, ) -from diff_diff.linalg import solve_ols +from diff_diff.linalg import _rank_guarded_inv, solve_ols from diff_diff.survey import ( ResolvedSurveyDesign, _resolve_survey_for_fit, @@ -1047,21 +1047,32 @@ def _compute_dose_response_gt( # Store bootstrap info for influence function computation # bread = (Psi'WPsi / n_treated)^{-1} when survey, (Psi'Psi / n_treated)^{-1} otherwise + # Bread = (Psi'WPsi / mass)^{-1} via the shared rank-guarded inverse: + # np.linalg.inv only raises on an *exactly* singular Gram, so a *near*- + # singular B-spline design (clustered doses / near-duplicate knots) + # previously returned a garbage inverse (~1e13) -> garbage SE. The prior + # `pinv` fallback was both minimum-norm (not the column-drop / near- + # collinear limit) and *silent*. `_rank_guarded_inv` truncates redundant + # directions on the equilibrated Gram -> finite SE on the identified + # subspace (NaN only at rank 0), matching the covariate IF rank-guard. if w_treated is not None: w_treated_sum = float(np.sum(w_treated)) PtWP = Psi.T @ (Psi * w_treated[:, np.newaxis]) # Normalize bread by weighted mass (not raw count) for consistency # with downstream IF score denominators that also use weighted mass - try: - bread = np.linalg.inv(PtWP / w_treated_sum) - except np.linalg.LinAlgError: - bread = np.linalg.pinv(PtWP / w_treated_sum) + bread, n_dropped, _ = _rank_guarded_inv(PtWP / w_treated_sum) else: PtP = Psi.T @ Psi - try: - bread = np.linalg.inv(PtP / n_treated) - except np.linalg.LinAlgError: - bread = np.linalg.pinv(PtP / n_treated) + bread, n_dropped, _ = _rank_guarded_inv(PtP / n_treated) + if n_dropped: + warnings.warn( + "ContinuousDiD ACRT variance: the B-spline design Gram is " + f"rank-deficient ({n_dropped} redundant direction(s) dropped); " + "rank-reducing to a finite SE on the identified subspace. " + "Analytical SEs reflect the reduced rank (NaN if rank 0).", + UserWarning, + stacklevel=2, + ) # ee_treated: per-unit estimating equation vectors (K-vector per unit) # For WLS (survey weights), the score is w_i * X_i * u_i to match the diff --git a/diff_diff/linalg.py b/diff_diff/linalg.py index f43efed6..7085e78d 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -312,12 +312,33 @@ def _equilibrated_lstsq(X: np.ndarray, y: np.ndarray) -> np.ndarray: return coef_scaled / safe_norms +@overload +def _rank_guarded_inv( + A: np.ndarray, + *, + rcond: float = ..., + tracker: Optional[list] = ..., + return_dropped: Literal[False] = ..., +) -> Tuple[np.ndarray, int, int]: ... + + +@overload +def _rank_guarded_inv( + A: np.ndarray, + *, + rcond: float = ..., + tracker: Optional[list] = ..., + return_dropped: Literal[True], +) -> Tuple[np.ndarray, int, int, np.ndarray]: ... + + def _rank_guarded_inv( A: np.ndarray, *, rcond: float = 1e-10, tracker: Optional[list] = None, -) -> Tuple[np.ndarray, int, int]: + return_dropped: bool = False, +) -> Union[Tuple[np.ndarray, int, int], Tuple[np.ndarray, int, int, np.ndarray]]: """Rank-guarded (generalized) inverse of a symmetric PSD Gram matrix. Influence-function standard errors invert a covariate Gram matrix @@ -384,8 +405,19 @@ def _rank_guarded_inv( unchanged, so well-conditioned fits are numerically unaffected. """ k = A.shape[0] + + def _ret(inv, n_dropped, n_keep, dropped): + # ``dropped`` is a length-k boolean mask of the truncated (unidentified) + # coordinates. Callers that report a PER-COEFFICIENT SE from ``vcov`` + # diagonals must NaN the dropped coordinates in the FINAL vcov — the + # zero-filled inverse would otherwise report ``se=0`` for an unidentified + # coefficient. (Linear-combination consumers — e.g. an ATT or dose + # prediction — keep the default 3-tuple: the dropped direction correctly + # contributes 0 to an identified linear combination.) + return (inv, n_dropped, n_keep, dropped) if return_dropped else (inv, n_dropped, n_keep) + if k == 0: - return np.zeros((0, 0), dtype=float), 0, 0 + return _ret(np.zeros((0, 0), dtype=float), 0, 0, np.zeros(0, dtype=bool)) # Symmetric equilibration: scale row/col i by sqrt(A[i, i]) so the # eigenvalue threshold is scale-invariant. Zero/negative diagonal -> 1.0 @@ -406,7 +438,7 @@ def _rank_guarded_inv( # Fast path: full rank -> exact solve (bit-identical to the prior code). if max_eig > 0.0 and n_keep == k: - return np.linalg.solve(A, np.eye(k)), 0, k + return _ret(np.linalg.solve(A, np.eye(k)), 0, k, np.zeros(k, dtype=bool)) # Rank-deficient: record one condition-number sample for the aggregate # fallback warning (the helper is the sole owner of this append). @@ -415,7 +447,7 @@ def _rank_guarded_inv( tracker.append(float(np.linalg.cond(A))) if n_keep == 0: - return np.full((k, k), np.nan), k, 0 + return _ret(np.full((k, k), np.nan), k, 0, np.ones(k, dtype=bool)) # Column-drop generalized inverse: keep the n_keep most-independent columns # (pivoted QR on the equilibrated Gram), invert that principal submatrix, and @@ -431,7 +463,9 @@ def _rank_guarded_inv( A_eq_ginv[np.ix_(kept, kept)] = np.linalg.inv(A_eq[np.ix_(kept, kept)]) A_ginv = A_eq_ginv * inv_scales[:, None] * inv_scales[None, :] n_dropped = k - n_keep - return A_ginv, n_dropped, n_keep + dropped_mask = np.ones(k, dtype=bool) + dropped_mask[kept] = False + return _ret(A_ginv, n_dropped, n_keep, dropped_mask) def _solve_ols_rust( diff --git a/diff_diff/spillover.py b/diff_diff/spillover.py index 2ceb9fe9..1bac157f 100644 --- a/diff_diff/spillover.py +++ b/diff_diff/spillover.py @@ -38,7 +38,7 @@ _haversine_km, _validate_callable_metric_result, ) -from diff_diff.linalg import solve_ols +from diff_diff.linalg import _rank_guarded_inv, solve_ols from diff_diff.results import SpilloverDiDResults from diff_diff.two_stage import _compute_gmm_corrected_meat from diff_diff.utils import safe_inference @@ -3357,27 +3357,38 @@ def fit( score_pad_mask=score_pad_mask_arg, ) - # Bread sandwich: A_22^{-1} = (X_2' W X_2)^{-1} via `np.linalg.solve` - # with dense lstsq fallback + UserWarning (mirrors the bread-fallback - # pattern at `two_stage.py:1763-1788`). Wave E.1 adds the W diagonal - # under the survey path so the bread aligns with the WLS gamma / - # weighted Psi construction in the meat helper. + # Bread sandwich: A_22^{-1} = (X_2' W X_2)^{-1} via the shared rank-guarded + # generalized inverse `_rank_guarded_inv` (column-drop on a near-singular + # Gram + UserWarning; dropped coordinates are NaN'd in the vcov below). + # Wave E.1 adds the W diagonal under the survey path so the bread aligns + # with the WLS gamma / weighted Psi construction in the meat helper. if survey_weights_fit is not None: A_22_kept = X_2_kept.T @ (X_2_kept * survey_weights_fit[:, None]) else: A_22_kept = X_2_kept.T @ X_2_kept - eye_kept = np.eye(A_22_kept.shape[0]) - try: - bread_kept = np.linalg.solve(A_22_kept, eye_kept) - except np.linalg.LinAlgError: + # np.linalg.solve only raises on an *exactly* singular Gram; a *near*- + # singular A_22 would otherwise flow a garbage inverse (~1e13) into the + # SE. `_rank_guarded_inv` truncates redundant directions on the + # equilibrated Gram -> finite SE on the identified subspace (NaN at + # rank 0), matching the covariate IF rank-guard. A_22_kept is already + # column-dropped upstream; this is the within-kept near-singular guard + # (its rank-0 all-NaN return composes with the (k,k) re-inflation below). + bread_kept, n_dropped, _, dropped = _rank_guarded_inv(A_22_kept, return_dropped=True) + if n_dropped: warnings.warn( - "SpilloverDiD Wave D bread: A_22 = X_2' X_2 is singular; " - "falling back to dense lstsq. SE may be unreliable.", + "SpilloverDiD Wave D bread: A_22 = X_2' W X_2 is rank-deficient; " + "rank-reducing to a finite SE on the identified subspace " + f"({n_dropped} redundant direction(s) dropped, NaN if rank 0).", UserWarning, stacklevel=2, ) - bread_kept = np.linalg.lstsq(A_22_kept, eye_kept, rcond=None)[0] vcov_kept = bread_kept @ meat_kept @ bread_kept + # A within-kept dropped (unidentified) coefficient is zero-filled in + # bread_kept, which would report se=0; NaN its row/col so per-coef SE is + # NaN, not 0. These ride along the (k, k) re-inflation below. + if dropped.any(): + vcov_kept[dropped, :] = np.nan + vcov_kept[:, dropped] = np.nan # Re-inflate to (k, k) with NaN at rank-deficient column positions # so downstream code (which indexes vcov[i, i] for per-coef SE) sees diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index 5255dc9c..6209b3e2 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -36,7 +36,7 @@ _validate_conley_kwargs, _validate_meat_psd, ) -from diff_diff.linalg import solve_ols +from diff_diff.linalg import _rank_guarded_inv, solve_ols from diff_diff.two_stage_bootstrap import TwoStageDiDBootstrapMixin from diff_diff.two_stage_results import ( TwoStageBootstrapResults, # noqa: F401 @@ -3175,29 +3175,36 @@ def _compute_gmm_variance( XtWX_2 = X_2.T @ (X_2 * survey_weights[:, None]) else: XtWX_2 = X_2.T @ X_2 - try: - bread = np.linalg.solve(XtWX_2, np.eye(k)) - except np.linalg.LinAlgError: - # Sibling of finding #17 (axis A) — the TSL-variance bread - # fallback was previously silent. Note: X_2 is the Stage-2 - # indicator design (treatment / horizon / group dummies), not - # user covariates, so the diagnostic guidance points at that - # layer. + # np.linalg.solve only raises on an *exactly* singular Gram; a *near*- + # singular X_2'WX_2 would otherwise flow a garbage inverse (~1e13) + # straight into the SE. `_rank_guarded_inv` truncates redundant + # directions on the equilibrated Gram -> finite SE on the identified + # subspace (NaN only at rank 0), matching the covariate IF rank-guard. + # Sibling of finding #17 (axis A): the prior fallback fired only on an + # exactly-singular matrix. X_2 is the Stage-2 indicator design (not user + # covariates), so the diagnostic guidance points at that layer. + bread, n_dropped, _, dropped = _rank_guarded_inv(XtWX_2, return_dropped=True) + if n_dropped: warnings.warn( "Rank-deficient second-stage design matrix X_2'WX_2 in " - "TwoStageDiD TSL variance; falling back to np.linalg.lstsq " - "for the bread matrix. Analytical SEs may be numerically " - "unstable. The Stage-2 design is built from treatment, " - "event-time, or group indicators, so this typically " + "TwoStageDiD TSL variance; rank-reducing to a finite SE on the " + f"identified subspace ({n_dropped} redundant direction(s) " + "dropped, NaN if rank 0). The Stage-2 design is built from " + "treatment, event-time, or group indicators, so this typically " "indicates a zero-weight or all-zero indicator column " "(e.g. an aggregation path with no qualifying observations).", UserWarning, stacklevel=2, ) - bread = np.linalg.lstsq(XtWX_2, np.eye(k), rcond=None)[0] # 7. V = bread @ meat @ bread V = bread @ meat @ bread + # A dropped (unidentified) Stage-2 coefficient is zero-filled in `bread`, + # which would report se=0 for that named coefficient; NaN its row/col in + # the FINAL vcov so per-coefficient SE extraction yields NaN (not 0). + if dropped.any(): + V[dropped, :] = np.nan + V[:, dropped] = np.nan return V def _build_fe_design( diff --git a/diff_diff/two_stage_bootstrap.py b/diff_diff/two_stage_bootstrap.py index e21aae42..4a2a0b2d 100644 --- a/diff_diff/two_stage_bootstrap.py +++ b/diff_diff/two_stage_bootstrap.py @@ -22,6 +22,7 @@ from diff_diff.bootstrap_utils import ( generate_survey_multiplier_weights_batch as _generate_survey_multiplier_weights_batch, ) +from diff_diff.linalg import _rank_guarded_inv from diff_diff.two_stage_results import TwoStageBootstrapResults # Maximum number of elements before falling back to per-column sparse aggregation. @@ -136,7 +137,7 @@ def _compute_cluster_S_scores( X_2: np.ndarray, cluster_ids: np.ndarray, survey_weights: Optional[np.ndarray] = None, - ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """ Compute per-cluster S_g scores for bootstrap. @@ -145,9 +146,13 @@ def _compute_cluster_S_scores( S : np.ndarray, shape (G, k) Per-cluster influence scores. bread : np.ndarray, shape (k, k) - (X'_2 X_2)^{-1}. + (X'_2 X_2)^{-1} (rank-guarded; zero-filled rows/cols for any dropped, + unidentified Stage-2 coordinate). unique_clusters : np.ndarray Unique cluster identifiers. + dropped : np.ndarray of bool, shape (k,) + Mask of dropped (unidentified) Stage-2 coordinates; callers NaN the + corresponding bootstrap coefficient columns so their SE is NaN, not 0. """ k = X_2.shape[1] @@ -273,27 +278,28 @@ def _compute_cluster_S_scores( XtX_2 = X_2.T @ (X_2 * survey_weights[:, None]) else: XtX_2 = np.dot(X_2.T, X_2) - try: - bread = np.linalg.solve(XtX_2, np.eye(k)) - except np.linalg.LinAlgError: - # Sibling of finding #17 (axis A) — the bootstrap bread - # fallback was previously silent. X_2 is the Stage-2 indicator - # design (treatment / horizon / group dummies), so a singular - # bread typically indicates a zero-weight or all-zero column. + # np.linalg.solve only raises on an *exactly* singular Gram; a *near*- + # singular X_2'WX_2 would otherwise flow a garbage inverse (~1e13) into + # the bootstrap SE. `_rank_guarded_inv` truncates redundant directions on + # the equilibrated Gram -> finite SE on the identified subspace (NaN at + # rank 0) — the cross-surface twin of the analytical TSL bread guard in + # two_stage.py. Sibling of finding #17 (axis A): the prior fallback fired + # only on an exactly-singular matrix. + bread, n_dropped, _, dropped = _rank_guarded_inv(XtX_2, return_dropped=True) + if n_dropped: warnings.warn( "Rank-deficient second-stage design matrix X_2'WX_2 in " - "TwoStageDiD multiplier bootstrap bread; falling back to " - "np.linalg.lstsq. Bootstrap SEs may be numerically " - "unstable. The Stage-2 design is built from treatment, " - "event-time, or group indicators, so this typically " - "indicates a zero-weight or all-zero indicator column " + "TwoStageDiD multiplier bootstrap bread; rank-reducing to a " + f"finite SE on the identified subspace ({n_dropped} redundant " + "direction(s) dropped, NaN if rank 0). The Stage-2 design is " + "built from treatment, event-time, or group indicators, so this " + "typically indicates a zero-weight or all-zero indicator column " "(e.g. an aggregation path with no qualifying observations).", UserWarning, stacklevel=2, ) - bread = np.linalg.lstsq(XtX_2, np.eye(k), rcond=None)[0] - return S, bread, unique_clusters + return S, bread, unique_clusters, dropped def _build_nan_bootstrap_results( self, @@ -404,7 +410,7 @@ def _run_bootstrap( X_2_static = D.reshape(-1, 1) - S_static, bread_static, unique_clusters = self._compute_cluster_S_scores( + S_static, bread_static, unique_clusters, _ = self._compute_cluster_S_scores( df=df, unit=unit, time=time, @@ -553,7 +559,7 @@ def _run_bootstrap( if h_int in horizon_to_col: X_2_es[i, horizon_to_col[h_int]] = 1.0 - S_es, bread_es, _ = self._compute_cluster_S_scores( + S_es, bread_es, _, dropped_es = self._compute_cluster_S_scores( df=df, unit=unit, time=time, @@ -570,6 +576,12 @@ def _run_bootstrap( # boot_coef_es: (B, k_es) boot_coef_es = np.dot(np.dot(all_weights, S_es), bread_es.T) + # A dropped (unidentified) event-time coefficient is zero-filled in + # bread_es -> a 0 bootstrap column -> se=0. NaN it (via the explicit + # dropped mask) so the per-horizon SE is NaN, not 0. (Defensive: a + # coefficient the point estimate also drops already has a NaN effect + # and is skipped below; this guards the inconsistent case.) + boot_coef_es[:, dropped_es] = np.nan event_study_ses = {} event_study_cis = {} @@ -614,7 +626,7 @@ def _run_bootstrap( if g in group_to_col: X_2_grp[i, group_to_col[g]] = 1.0 - S_grp, bread_grp, _ = self._compute_cluster_S_scores( + S_grp, bread_grp, _, dropped_grp = self._compute_cluster_S_scores( df=df, unit=unit, time=time, @@ -630,6 +642,9 @@ def _run_bootstrap( ) boot_coef_grp = np.dot(np.dot(all_weights, S_grp), bread_grp.T) + # NaN any dropped (unidentified) group coefficient (via the explicit + # dropped mask) so its per-group SE is NaN, not 0 (see event-study note). + boot_coef_grp[:, dropped_grp] = np.nan group_ses = {} group_cis = {} diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 20647cab..d8866e4d 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -861,6 +861,8 @@ The guard is fired by `_survey_se_from_group_if` (analytical and replicate) and **R Reference:** `contdid` v0.1.0 (CRAN). +- **Note (rank-guarded ACRT-variance bread):** The ACRT influence-function bread `(Psi'WPsi / mass)^{-1}` (`_compute_dose_response_gt`, `continuous_did.py`) is inverted by the shared `_rank_guarded_inv` (`diff_diff/linalg.py`). `np.linalg.inv` raises only on an *exactly* singular B-spline design Gram; a **near**-singular Gram (clustered doses / near-duplicate knots) previously returned a garbage inverse (~1e13), and the exact-singular fallback was a *silent* minimum-norm `pinv`. The rank-guard truncates redundant directions on the equilibrated Gram → a finite SE on the identified subspace (the well-conditioned near-collinear limit, **not** minimum-norm; NaN only at rank 0), and `fit()` warns when a direction is dropped. See the CallawaySantAnna "rank-guarded IF standard errors" Note for the generalized-inverse semantics. + ### Identification Two levels of parallel trends (following CGBS 2024, Assumptions 1-2): @@ -1425,6 +1427,8 @@ where `W_it(h) = 1[K_it = h]` are lead indicators, estimated on `Omega_0` only. **Primary source:** [Gardner, J. (2022). Two-stage differences in differences. arXiv:2207.05943.](https://arxiv.org/abs/2207.05943) +- **Note (rank-guarded TSL-variance bread):** The Stage-2 variance bread `(X_2'WX_2)^{-1}` is inverted by the shared `_rank_guarded_inv` (`diff_diff/linalg.py`) on BOTH the analytical (`two_stage.py`) and multiplier-bootstrap (`two_stage_bootstrap.py`) surfaces. `np.linalg.solve` raised only on an exactly-singular bread (prior fallback: dense `lstsq`); a **near**-singular `X_2'WX_2` would otherwise flow a garbage inverse (~1e13) into the SE. The rank-guard truncates redundant directions → finite SE on the identified subspace (NaN at rank 0) and warns. A dropped (unidentified) Stage-2 coefficient (event-time / group effect) is reported with **NaN** SE on both the analytical and bootstrap surfaces — not the zero-filled `0`. `X_2` is the Stage-2 indicator design (treatment/event-time/group dummies), not user covariates. See the CallawaySantAnna "rank-guarded IF standard errors" Note. Sibling of axis-A finding #17. + **Key implementation requirements:** *Assumption checks / warnings:* @@ -3668,6 +3672,8 @@ should be a deliberate user choice. **Note:** This estimator is a **documented synthesis** of ingredients — no single published software package implements the exact recipe. `did2s` (R/Stata, Butts & Gardner 2022) implements the Gardner two-stage residualization but does NOT support ring covariates. The Butts (2021/2023) paper proposes the ring estimator in Equations 5/6/8 (non-staggered) and Section 5 / Table 2 (staggered) but does not ship reference software. The diff-diff implementation combines: (a) Butts (2021) Section 5 / Table 2 identification, (b) Gardner (2022) two-stage residualize-then-fit, (c) Wave A's Conley spatial-HAC vcov. +- **Note (rank-guarded Wave D bread):** The Wave D variance bread `A_22 = (X_2'WX_2)^{-1}` is inverted by the shared `_rank_guarded_inv` (`diff_diff/linalg.py`) on the already-column-dropped kept submatrix. `np.linalg.solve` raised only on an exactly-singular bread (prior fallback: dense `lstsq`); a **near**-singular kept Gram would otherwise flow a garbage inverse (~1e13) into the SE. The rank-guard truncates redundant directions → finite SE on the identified subspace (NaN at rank 0), re-inflated to (k, k) with NaN at the dropped positions, and warns. See the CallawaySantAnna "rank-guarded IF standard errors" Note. Sibling of axis-A finding #17. + **Identification spec (committed):** The stage-2 regressor for ring `j` is the **time-varying** form @@ -3903,6 +3909,8 @@ temporally) correlated. Extends White (1980) HC0 by allowing pairwise correlation that decays with geographic distance, plus a within-unit Newey-West-style Bartlett temporal HAC on panel data. +- **Note (rank-guarded design bread):** The spatial-HAC sandwich bread `(X'WX)^{-1}` (`_compute_conley_vcov`, `conley.py`) is inverted by the shared `_rank_guarded_inv` (`diff_diff/linalg.py`). A near-singular design Gram previously returned a garbage inverse (~1e13) and an *exactly* singular one **raised `ValueError`**; the bread now rank-reduces to a finite SE on the identified subspace (NaN only at rank 0) and warns, matching the other structural bread guards. A dropped (unidentified) regression coefficient is reported with **NaN** SE (its row/col in the returned vcov), not the zero-filled `0`. **Behavior change:** a rank-deficient design (collinear regressors) no longer raises — it rank-reduces with a warning. The well-conditioned path is unchanged (`np.linalg.solve(A, I)`). `_rank_guarded_inv` is imported lazily inside the function because `linalg` imports `conley` (one-way), so a top-level import would be circular. See the CallawaySantAnna "rank-guarded IF standard errors" Note. + **Two operating modes:** - **Cross-sectional (Phase 1):** Pass `vcov_type="conley"` plus `conley_coords` (n × 2 array) and `conley_cutoff_km` on direct diff --git a/tests/test_conley_vcov.py b/tests/test_conley_vcov.py index a48f6db9..f66f8bc3 100644 --- a/tests/test_conley_vcov.py +++ b/tests/test_conley_vcov.py @@ -3657,3 +3657,78 @@ def _staggered_panel(seed=51, n_units=36, n_periods=6): y = 0.3 * (u / n_units) + 0.4 * t + 1.5 * post + rng.normal(0, 0.5) recs.append({"unit": u, "time": t, "y": y, "cohort": g, "lat": lat, "lon": lon}) return pd.DataFrame(recs) + + +class TestConleyBreadRankGuard: + """The conley spatial-HAC bread inversion now routes through the shared + ``_rank_guarded_inv``: a near-singular design Gram rank-reduces to a finite + SE on the identified subspace (previously a garbage ~1e13 inverse), and a + singular Gram no longer raises ``ValueError`` (it rank-reduces + warns). + Sibling of the covariate IF rank-guard and the ContinuousDiD / TwoStageDiD / + SpilloverDiD structural bread guards.""" + + @staticmethod + def _cross_section(seed, dup="indep"): + # Fixed rng draw order (x1, noise, y) so ``exact`` and ``near`` share + # x1 / coords / y and differ ONLY in the third column. + rng = np.random.default_rng(seed) + n = 60 + coords = rng.uniform(0.0, 5.0, size=(n, 2)) + x1 = rng.normal(size=n) + noise = rng.normal(size=n) + y = rng.normal(size=n) + if dup == "exact": + x2 = x1.copy() # exactly collinear -> singular X'X + elif dup == "near": + x2 = x1 + 1e-3 * noise # highly collinear but full rank + else: + x2 = noise # independent + X = np.column_stack([np.ones(n), x1, x2]) + beta, *_ = np.linalg.lstsq(X, y, rcond=None) + resid = y - X @ beta + return X, resid, coords + + def test_singular_gram_rank_reduces_not_raises(self): + X, resid, coords = self._cross_section(0, dup="exact") + bread = X.T @ X # exactly singular (x2 == x1) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + V = _compute_conley_vcov(X, resid, coords, 3.0, "euclidean", "bartlett", bread) + # Pre-fix this raised ValueError; now it rank-reduces. The dropped + # (duplicate, unidentified) coefficient gets a NaN row/col — se=NaN, not + # the zero-filled se=0 — while the identified coefficients stay finite. + nan_diag = np.isnan(np.diag(V)) + assert nan_diag.sum() == 1, f"exactly one duplicate coef should be NaN, got {np.diag(V)}" + kept = np.flatnonzero(~nan_diag) + assert 0 in kept, "intercept (identified) must remain finite" + assert np.all(np.isfinite(V[np.ix_(kept, kept)])), "identified block must be finite" + dropped_idx = int(np.flatnonzero(nan_diag)[0]) + assert np.all(np.isnan(V[dropped_idx, :])) and np.all(np.isnan(V[:, dropped_idx])) + msgs = [str(w.message) for w in caught] + assert any( + "Conley spatial HAC variance" in m and "rank-deficient" in m for m in msgs + ), msgs + + def test_rank_zero_gram_returns_nan(self): + X, resid, coords = self._cross_section(1, dup="indep") + Xz = np.zeros_like(X) + bread = Xz.T @ Xz # rank 0 + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + V = _compute_conley_vcov(Xz, resid, coords, 3.0, "euclidean", "bartlett", bread) + assert np.all(np.isnan(V)) + + def test_column_drop_equals_near_collinear_limit(self): + # Column-drop generalized inverse == the near-collinear full-rank limit: + # the identified-subspace SE (intercept) matches between an exactly- + # collinear design (rank-reduced) and a highly-but-not-exactly collinear + # one (full-rank normal inverse). Mirrors the covariate rank-guard's + # verified se_ratio ~ 1 property. + Xe, re_, ce = self._cross_section(7, dup="exact") + Xn, rn, cn = self._cross_section(7, dup="near") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + Ve = _compute_conley_vcov(Xe, re_, ce, 3.0, "euclidean", "bartlett", Xe.T @ Xe) + Vn = _compute_conley_vcov(Xn, rn, cn, 3.0, "euclidean", "bartlett", Xn.T @ Xn) + assert np.isfinite(Ve[0, 0]) and Ve[0, 0] > 0 + np.testing.assert_allclose(Ve[0, 0], Vn[0, 0], rtol=5e-2) diff --git a/tests/test_continuous_did.py b/tests/test_continuous_did.py index bc4d025b..4cffb060 100644 --- a/tests/test_continuous_did.py +++ b/tests/test_continuous_did.py @@ -1528,3 +1528,39 @@ def test_event_study_analytical_se_finite(self): assert 0 <= info["p_value"] <= 1, f"p_value out of range for e={e}" lo, hi = info["conf_int"] assert np.isfinite(lo) and np.isfinite(hi), f"conf_int contains NaN for e={e}" + + +class TestContinuousDiDBreadRankGuard: + """The ContinuousDiD ACRT-variance bread (Psi'WPsi) now routes through the + shared ``_rank_guarded_inv``: a near-singular B-spline design Gram rank- + reduces to a finite SE on the identified subspace and warns (the prior + ``pinv`` fallback was minimum-norm AND silent).""" + + def test_rank_deficient_bspline_gram_warns_and_finite_se(self): + import warnings + from unittest.mock import patch + + import diff_diff.continuous_did as cd_mod + + data = generate_continuous_did_data(n_units=100, n_periods=3, seed=42, noise_sd=0.5) + real_rgi = cd_mod._rank_guarded_inv + + def force_drop(A, **kwargs): + # Finite inverse, but report a dropped direction to exercise the + # rank-reduce warning path deterministically (B-spline Gram rank- + # deficiency is hard to force via dose data alone). + inv, _, rank = real_rgi(A, **kwargs) + return inv, 1, rank + + with patch.object(cd_mod, "_rank_guarded_inv", side_effect=force_drop): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + results = ContinuousDiD().fit( + data, "outcome", "unit", "period", "first_treat", "dose" + ) + msgs = [str(w.message) for w in caught] + assert any( + "ContinuousDiD ACRT variance" in m and "rank-deficient" in m for m in msgs + ), msgs + # rank-reduced bread still yields finite ACRT SEs. + assert np.all(np.isfinite(results.dose_response_acrt.se)) diff --git a/tests/test_linalg.py b/tests/test_linalg.py index 301e91d0..fed4e357 100644 --- a/tests/test_linalg.py +++ b/tests/test_linalg.py @@ -2221,3 +2221,33 @@ def test_empty_matrix(self): assert ginv.shape == (0, 0) assert n_dropped == 0 assert rank == 0 + + def test_return_dropped_mask(self): + # return_dropped=True exposes a length-k boolean mask of the truncated + # (unidentified) coordinates so per-coefficient callers can NaN them + # (a zero-filled dropped coordinate would otherwise report se=0). + rng = np.random.RandomState(7) + n = 100 + x = rng.standard_normal(n) + X = np.column_stack([np.ones(n), x, x]) # cols 1 and 2 exactly collinear + A = self._gram(X) + + # Default call is unchanged (backward-compatible 3-tuple). + assert len(_rank_guarded_inv(A)) == 3 + + ginv, n_dropped, rank, dropped = _rank_guarded_inv(A, return_dropped=True) + assert n_dropped == 1 and rank == 2 + assert dropped.dtype == bool and dropped.shape == (3,) + assert dropped.sum() == 1 + # The dropped coordinate is exactly the zero-filled row/col of the inverse. + assert np.all(ginv[dropped] == 0.0) + assert bool(dropped[0]) is False # intercept stays identified + + # Full rank -> no dropped coordinates. + A_ok = self._gram(rng.standard_normal((n, 3))) + _, nd_ok, _, dropped_ok = _rank_guarded_inv(A_ok, return_dropped=True) + assert nd_ok == 0 and not dropped_ok.any() + + # Rank 0 -> every coordinate dropped. + _, nd0, _, dropped0 = _rank_guarded_inv(np.zeros((3, 3)), return_dropped=True) + assert nd0 == 3 and dropped0.all() diff --git a/tests/test_methodology_two_stage.py b/tests/test_methodology_two_stage.py index 5adc6d80..2406c190 100644 --- a/tests/test_methodology_two_stage.py +++ b/tests/test_methodology_two_stage.py @@ -464,7 +464,7 @@ def test_bootstrap_scores_use_exact_residuals(self) -> None: ) treated = ((df["first_treat"] > 0) & (df["time"] >= df["first_treat"])).values X_2 = treated.astype(float).reshape(-1, 1) - S, bread, _ = est._compute_cluster_S_scores( + S, bread, _, _ = est._compute_cluster_S_scores( df, "unit", "time", @@ -523,7 +523,7 @@ def test_bootstrap_event_study_scores_use_exact_residuals(self) -> None: X_2, cluster_ids, ) - S, bread, _ = est._compute_cluster_S_scores( + S, bread, _, _ = est._compute_cluster_S_scores( df, "unit", "time", diff --git a/tests/test_spillover.py b/tests/test_spillover.py index 5ae261e7..7f8ee611 100644 --- a/tests/test_spillover.py +++ b/tests/test_spillover.py @@ -8826,3 +8826,80 @@ def test_r_drift_goldens(self): assert res.n_psu == 8 assert res.n_strata == 2 assert res.survey_metadata.df_survey == 6 + + +class TestSpilloverDiDBreadRankGuard: + """The SpilloverDiD Wave D bread (A_22 = X_2' W X_2) now routes through the + shared ``_rank_guarded_inv``: a near-singular Stage-2 design Gram rank- + reduces to a finite SE on the identified subspace and warns (was a dense + lstsq fallback that fired only on an exactly-singular bread).""" + + def test_rank_deficient_bread_warns_and_fits(self): + from unittest.mock import patch + + import diff_diff.spillover as sp_mod + + df = _make_butts_2period_dgp(seed=42) + real_rgi = sp_mod._rank_guarded_inv + + def force_drop(A, **kwargs): + # Finite inverse + n_dropped=1 (warning fires), empty dropped mask so + # the identified SE stays finite. Mirrors the helper's return arity. + inv, _, rank = real_rgi(A) + if kwargs.get("return_dropped"): + return inv, 1, rank, np.zeros(A.shape[0], dtype=bool) + return inv, 1, rank + + est = SpilloverDiD(rings=[0.0, 100.0], conley_coords=("lat", "lon")) + with patch.object(sp_mod, "_rank_guarded_inv", side_effect=force_drop): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + result = est.fit(df, outcome="y", unit="unit", time="time", treatment="D") + msgs = [str(w.message) for w in caught] + assert any( + "SpilloverDiD Wave D bread" in m and "rank-deficient" in m for m in msgs + ), msgs + assert est.is_fitted_ and np.isfinite(result.att) + + def test_dropped_ring_coefficient_propagates_nan_inference(self): + """A dropped (unidentified) Wave D coefficient must report NaN se / t_stat / + p_value / CI for its ring effect at the ESTIMATOR level (not the zero-filled + se=0) — the per-coefficient propagation the rank-guard enables (CI codex P2 + test-depth).""" + from unittest.mock import patch + + import diff_diff.spillover as sp_mod + + df = _make_butts_2period_dgp(seed=42) + real_rgi = sp_mod._rank_guarded_inv + + def drop_last(A, **kwargs): + # Genuinely drop the last Wave D coordinate (zero-fill its row/col + + # report it dropped) so the caller NaNs that ring's vcov entry. + inv, _, rank = real_rgi(A) + inv = np.array(inv, dtype=float) + inv[-1, :] = 0.0 + inv[:, -1] = 0.0 + k = A.shape[0] + dropped = np.zeros(k, dtype=bool) + dropped[-1] = True + if kwargs.get("return_dropped"): + return inv, 1, k - 1, dropped + return inv, 1, k - 1 + + est = SpilloverDiD(rings=[0.0, 50.0, 100.0], conley_coords=("lat", "lon")) + with patch.object(sp_mod, "_rank_guarded_inv", side_effect=drop_last): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = est.fit(df, outcome="y", unit="unit", time="time", treatment="D") + eff = res.spillover_effects + nan_rows = eff[np.isnan(eff["se"])] + fin_rows = eff[np.isfinite(eff["se"]) & (eff["se"] > 0)] + assert len(nan_rows) >= 1, ( + f"a dropped Wave D coord should NaN a ring SE; got {eff['se'].tolist()}" + ) + assert len(fin_rows) >= 1, "identified rings should keep finite SE" + # The NaN-SE ring's FULL inference must be NaN, not just se. + for _, r in nan_rows.iterrows(): + assert np.isnan(r["t_stat"]) and np.isnan(r["p_value"]) + assert np.isnan(r["ci_low"]) and np.isnan(r["ci_high"]) diff --git a/tests/test_two_stage.py b/tests/test_two_stage.py index 1e5b29fb..68dc4d74 100644 --- a/tests/test_two_stage.py +++ b/tests/test_two_stage.py @@ -1451,15 +1451,18 @@ def test_iterative_demean_no_warning_on_convergence(self): class TestTwoStageStage2BreadWarning: - """Sibling of STD finding #17: the TwoStage Stage-2 bread fallback - (`X'_2 W X_2` singular) was silent. X_2 is built from treatment/ - event-time/group indicators — not user covariates — so we force the - LinAlgError via patching `np.linalg.solve` rather than data crafting, + """Sibling of STD finding #17: the TwoStage Stage-2 bread (`X'_2 W X_2`) + inversion was silent on a singular design and garbage on a *near*-singular + one. It now routes through the shared `_rank_guarded_inv`, which rank-reduces + to a finite SE on the identified subspace and warns. X_2 is built from + treatment/event-time/group indicators (not user covariates), so we force the + rank-deficiency at the `_rank_guarded_inv` seam rather than via data crafting, per the PR #334 CI review guidance.""" - def test_analytical_bread_lstsq_fallback_warns(self): - """When np.linalg.solve on the Stage-2 bread raises, the analytical - TSL path must warn and still return a finite variance via lstsq.""" + def test_analytical_bread_rank_reduces_and_warns(self): + """When the Stage-2 bread is rank-deficient, the analytical TSL path + rank-reduces via `_rank_guarded_inv`, warns, and still returns a finite + variance.""" from unittest.mock import patch import diff_diff.two_stage as ts_mod @@ -1467,22 +1470,20 @@ def test_analytical_bread_lstsq_fallback_warns(self): data = generate_test_data(n_units=80, n_periods=6, seed=77) est = TwoStageDiD() - real_solve = np.linalg.solve - - def raise_for_square_eye(a, b): - # Identify the Stage-2 bread call by shape: b is np.eye(k) - # (square identity). All other solve calls in the codepath - # have different b shapes. - if ( - isinstance(b, np.ndarray) - and b.ndim == 2 - and b.shape[0] == b.shape[1] - and np.allclose(b, np.eye(b.shape[0])) - ): - raise np.linalg.LinAlgError("forced by test") - return real_solve(a, b) + real_rgi = ts_mod._rank_guarded_inv + + def force_drop(A, **kwargs): + # Force a rank-deficiency *report* (finite inverse, n_dropped=1) to + # exercise the rank-reduce warning path deterministically, with an + # empty dropped mask so the identified ATT SE stays finite (the + # NaN-for-dropped behavior is covered by the conley direct-call and + # _rank_guarded_inv unit tests). Mirrors the helper's return arity. + inv, _, rank = real_rgi(A) + if kwargs.get("return_dropped"): + return inv, 1, rank, np.zeros(A.shape[0], dtype=bool) + return inv, 1, rank - with patch.object(ts_mod.np.linalg, "solve", side_effect=raise_for_square_eye): + with patch.object(ts_mod, "_rank_guarded_inv", side_effect=force_drop): with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") result = est.fit( @@ -1494,18 +1495,18 @@ def raise_for_square_eye(a, b): ) fallback = [w for w in caught if "TwoStageDiD TSL variance" in str(w.message)] assert len(fallback) >= 1, ( - "Expected TSL-variance bread fallback warning when np.linalg.solve " - f"was forced to raise; got warnings: " - f"{[str(w.message) for w in caught]}" + "Expected TSL-variance rank-reduce warning when the Stage-2 bread is " + f"rank-deficient; got warnings: {[str(w.message) for w in caught]}" ) msg = str(fallback[0].message) - assert "np.linalg.lstsq" in msg + assert "rank-reducing" in msg assert "X_2'WX_2" in msg - # lstsq fallback must still produce a finite SE. + # rank-reduced bread must still produce a finite SE. assert np.isfinite(result.overall_se) - def test_bootstrap_bread_lstsq_fallback_warns(self): - """Same contract for the multiplier-bootstrap bread path.""" + def test_bootstrap_bread_rank_reduces_and_warns(self): + """Same contract for the multiplier-bootstrap bread path (the + cross-surface twin in two_stage_bootstrap.py).""" from unittest.mock import patch import diff_diff.two_stage_bootstrap as tsb_mod @@ -1513,19 +1514,17 @@ def test_bootstrap_bread_lstsq_fallback_warns(self): data = generate_test_data(n_units=80, n_periods=6, seed=77) est = TwoStageDiD(n_bootstrap=10, seed=0) - real_solve = np.linalg.solve + real_rgi = tsb_mod._rank_guarded_inv - def raise_for_square_eye(a, b): - if ( - isinstance(b, np.ndarray) - and b.ndim == 2 - and b.shape[0] == b.shape[1] - and np.allclose(b, np.eye(b.shape[0])) - ): - raise np.linalg.LinAlgError("forced by test") - return real_solve(a, b) + def force_drop(A, **kwargs): + # n_dropped=1 (warning fires), empty dropped mask so the identified + # bootstrap SEs stay finite. Mirrors the helper's return arity. + inv, _, rank = real_rgi(A) + if kwargs.get("return_dropped"): + return inv, 1, rank, np.zeros(A.shape[0], dtype=bool) + return inv, 1, rank - with patch.object(tsb_mod.np.linalg, "solve", side_effect=raise_for_square_eye): + with patch.object(tsb_mod, "_rank_guarded_inv", side_effect=force_drop): with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") est.fit( @@ -1537,14 +1536,105 @@ def raise_for_square_eye(a, b): ) fallback = [w for w in caught if "TwoStageDiD multiplier bootstrap bread" in str(w.message)] assert len(fallback) >= 1, ( - "Expected bootstrap-bread fallback warning when np.linalg.solve " - f"was forced to raise; got warnings: " - f"{[str(w.message) for w in caught]}" + "Expected bootstrap-bread rank-reduce warning when the bread is " + f"rank-deficient; got warnings: {[str(w.message) for w in caught]}" ) msg = str(fallback[0].message) - assert "np.linalg.lstsq" in msg + assert "rank-reducing" in msg assert "X_2'WX_2" in msg + @staticmethod + def _drop_last_rgi(real_rgi): + """Mock factory: genuinely drop the last Stage-2 coordinate (zero-fill its + row/col in the inverse + report it dropped) so the caller NaNs it.""" + + def drop_last(A, **kwargs): + inv, _, rank = real_rgi(A) + inv = np.array(inv, dtype=float) + inv[-1, :] = 0.0 + inv[:, -1] = 0.0 + k = A.shape[0] + dropped = np.zeros(k, dtype=bool) + dropped[-1] = True + if kwargs.get("return_dropped"): + return inv, 1, k - 1, dropped + return inv, 1, k - 1 + + return drop_last + + @pytest.mark.parametrize( + "aggregate,attr", + [("event_study", "event_study_effects"), ("group", "group_effects")], + ) + def test_dropped_coefficient_propagates_nan_inference(self, aggregate, attr): + """A dropped (unidentified) Stage-2 coefficient must report NaN se / t_stat + / p_value / conf_int at the ESTIMATOR level (not the zero-filled se=0) — the + per-coefficient propagation the rank-guard enables (CI codex P2 test-depth). + Covers the analytical event-study AND group surfaces.""" + from unittest.mock import patch + + import diff_diff.two_stage as ts_mod + + data = generate_test_data(n_units=80, n_periods=6, seed=77) + with patch.object( + ts_mod, "_rank_guarded_inv", side_effect=self._drop_last_rgi(ts_mod._rank_guarded_inv) + ): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = TwoStageDiD().fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate=aggregate, + ) + effects = getattr(res, attr) + nan_k = [k for k, e in effects.items() if np.isnan(e["se"])] + fin_k = [k for k, e in effects.items() if np.isfinite(e["se"]) and e["se"] > 0] + assert nan_k, ( + f"a dropped Stage-2 coordinate should yield a NaN-se {aggregate} effect; " + f"got {{k: e['se'] for k, e in effects.items()}}" + ) + assert fin_k, "identified effects should keep finite SE" + # A dropped coefficient's FULL inference tuple must be NaN, not just se. + for k in nan_k: + e = effects[k] + assert np.isnan(e["t_stat"]) and np.isnan(e["p_value"]) + assert all(np.isnan(c) for c in e["conf_int"]) + + def test_bootstrap_dropped_coefficient_propagates_nan_inference(self): + """Same estimator-level NaN propagation on the multiplier-bootstrap surface: + a dropped Stage-2 coordinate NaNs the bootstrap coefficient column, so the + affected event-time SE is NaN while identified horizons stay finite.""" + from unittest.mock import patch + + import diff_diff.two_stage as ts_mod + import diff_diff.two_stage_bootstrap as tsb_mod + + data = generate_test_data(n_units=80, n_periods=6, seed=77) + with patch.object( + tsb_mod, "_rank_guarded_inv", side_effect=self._drop_last_rgi(tsb_mod._rank_guarded_inv) + ): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = ts_mod.TwoStageDiD(n_bootstrap=25, seed=0).fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="all", + ) + es = res.event_study_effects + nan_h = [h for h, e in es.items() if np.isnan(e["se"])] + fin_h = [h for h, e in es.items() if np.isfinite(e["se"]) and e["se"] > 0] + assert nan_h, ( + "a dropped Stage-2 coordinate should yield a NaN-se bootstrap horizon; " + f"got {{h: e['se'] for h, e in es.items()}}" + ) + assert fin_h, "identified bootstrap horizons should keep finite SE" + # ============================================================================= # TestTwoStageDiDWaveE3ParityAlwaysTreated