From a3553dd102bcefd85e5c7c4f8bd73c6b0d83b552 Mon Sep 17 00:00:00 2001 From: igerber Date: Mon, 6 Jul 2026 13:15:30 -0400 Subject: [PATCH] perf(efficient_did): omega_ridge stabilization + fused tiled GEMM analytical path Ridge-regularize the Omega* inversion behind the efficient weights (omega_ridge param, default 1e-6 relative to trace/H; 0 = exact legacy inv/pinv path, entire legacy code path preserved verbatim) and rewrite the conditional analytical path as a fused unit-tiled GEMM pass (kernel-covariance tables with (tpre_j, tpre_k) dedup, per-group kernel matrices reused across cells, batched ridge solves replacing the per-unit SVD + pseudoinverse loop). nocov twin: Gram-form omega + the same ridge on global weights; per-cell pseudoinverse warnings consolidate into one fit-level warning. The ridge path drops the identically-zero (g'=g, t_pre=t) self-pair so pre-treatment placebos stay data-driven. Motivation beyond speed: PT-All's telescoping overidentified moments make sample Omega* numerically singular (cond 1e17-1e22, 100% of units), so the legacy pseudoinverse redrew per-cell ATT(g,t) at ~1e-2 rel under ANY floating-point change (1-ulp input perturbation on the shipped code: 1.2e-4). After: 1-ulp per-cell stability ~3e-9. Default calibrated against the HRS Table 6 anchors (all unchanged, worst 0.0257 SE) and Monte Carlo (bias/RMSE/coverage identical to legacy). Conditional fit: 337.3s -> 7.8s at 2k units (43x), 57.7s at 10k, and a 100k-unit fit now completes at 4.0 GB peak (was memory-killed). Survey-weighted 1k: 95.3s -> 3.3s. Docs: REGISTRY Omega* ridge Note (calibration evidence, one-time value shift, degenerate-pair exclusion, diagnostics rationale), CHANGELOG, performance-plan tables, TODO row swap, llms-full line, pseudoinverse-mention sweep. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01X4AzrFUMqJxcUumSH31mSr --- CHANGELOG.md | 30 ++ TODO.md | 2 +- diff_diff/efficient_did.py | 209 ++++++++++--- diff_diff/efficient_did_covariates.py | 372 +++++++++++++++++++++++- diff_diff/efficient_did_results.py | 4 + diff_diff/efficient_did_weights.py | 179 +++++++++++- diff_diff/guides/llms-full.txt | 1 + docs/methodology/REGISTRY.md | 5 +- docs/performance-plan.md | 52 ++++ tests/test_efficient_did.py | 347 +++++++++++++++++++++- tests/test_methodology_efficient_did.py | 281 +++++++++++++++++- 11 files changed, 1427 insertions(+), 55 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 96d4f1cfc..6681fdb99 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,6 +37,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 survey/subpopulation design leaving `< 2` positive-weight `d_L` units all raise; multi-cohort and `covariates=` × `lowest_dose` raise `NotImplementedError` (deferred). The default `never_treated` / `not_yet_treated` paths are unchanged. +- **`EfficientDiD` `omega_ridge` parameter** (default `1e-6`) — ridge-regularizes the Omega* + inversion behind the efficient weights: solves `(Omega* + omega_ridge * max(trace/H, 0) * I) x = 1` + instead of pseudo-inverting the numerically singular Omega* that PT-All's telescoping + overidentified moments produce (measured cond 1e17–1e22 for 100% of units on realistic panels). + Makes per-cell `ATT(g,t)` numerically well-defined and platform-stable: a 1-ulp input + perturbation now moves per-cell values by ≤3e-9 relative (previously ~1e-4 through the + pseudoinverse's rcond-cutoff cliff, meaning per-cell values silently depended on BLAS/platform). + Calibrated against the HRS Table 6 replication anchors (all unchanged; worst deviation + 0.0257 SE, shift ≤ 0.0001 SE) and Monte Carlo (bias/RMSE/SE-calibration/coverage statistically + identical to legacy). **One-time value shift:** covariate-path per-cell and event-study values + move within their pre-existing indeterminacy band on upgrade (worst observed post-treatment + cell shift ~0.6 of its own SE at n=500, shrinking with n; overall-ATT shift 1.6e-2 → 1.1e-3 + relative from n=500 → 2k); the no-covariates path is essentially unchanged (~1e-7). + `omega_ridge=0` restores the entire legacy code path bit-for-bit. The ridge path also drops the + degenerate pre-treatment self-pair `(g'=g, t_pre=t)` (an identically-zero moment) so + pre-treatment placebos stay data-driven, and consolidates the legacy per-cell + condition-number warnings into one fit-level warning. See the Omega* ridge Note in + `docs/methodology/REGISTRY.md`. - **`ContinuousDiD` discrete-treatment saturated regression** (`treatment_type="discrete"`) for multi-valued / discrete dose (CGBS 2024 Eq. 4.1). Each distinct dose level gets its own effect coefficient — `ATT(d_j) = mean_{D=d_j}(ΔY) − control` (a per-level 2×2 DiD) — instead of a B-spline @@ -191,6 +209,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 supersedes it. ### Changed +- **`EfficientDiD` analytical covariate path rewritten as a fused unit-tiled GEMM pass** + (kernel-covariance tables with `(t_pre_j, t_pre_k)` dedup, per-group kernel matrices reused + across all `(g,t)` cells, batched ridge solves replacing the per-unit SVD + pseudoinverse loop; + the no-covariates Omega* gets a Gram-form twin). Measured (3-rep medians, 20 periods, + 5 cohorts, 5 covariates, `aggregate="all"`): fit 35.9s → 1.8s at 500 units (20x), + 95.4s → 3.5s at 1k (27x), 337.3s → 7.8s at 2k (43x), ~2.3h (extrapolated O(n²)) → 58s at 10k; + a 100k-unit fit now completes in ~86 min at 4.0 GB peak RSS (previously memory-killed). + Survey-weighted: 95.3s → 3.3s at 1k units. Memory is tile-bounded + (`_TARGET_OMEGA_TILE_BYTES` = 256 MB). Point estimates under the default ridge differ from + the legacy pseudoinverse path as described under `omega_ridge` above; at `omega_ridge=0` + results are bit-identical to the previous release. The conditional-path scalability warning + now fires once per fit at n > 50,000 (previously per cell at n > 5,000). - **CallawaySantAnna multiplier bootstrap rewritten as a fused, column-tiled scatter-GEMM** (EfficientDiD's bootstrap routes through the same kernel). The former loop sliced the `(block × n_units)` weight matrix twice per (g,t) cell per weight block — at a 40-period, diff --git a/TODO.md b/TODO.md index 4f0da2157..52b16feda 100644 --- a/TODO.md +++ b/TODO.md @@ -47,7 +47,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| -| `EfficientDiD` analytical stage has a pathological scaling hot loop: 339s for ONE fit at just 2k units × 20 periods with 5 covariates (15+ CPU-min at 10k units; memory pressure at 100k), while CS handles 100k units in ~2s. Process sampling attributes it to a broadcasted NumPy subtract/multiply inside a Python-level loop (EIF/weights construction). Profile properly and vectorize; until then EfficientDiD is impractical beyond a few thousand units. Found while benchmarking the Phase 2 bootstrap rewrite (its bootstrap stage is 0.018s at 2k units — the analytical stage is the whole problem). | `efficient_did.py` / `efficient_did_weights.py` | CS-scaling | Heavy | Medium | +| `EfficientDiD` conditional path at very large n: the fused tiled GEMM rewrite (v3.7; 43x at 2k units, 100k now completes in ~86 min) computes kernel-covariance tables per cell per tile. Cross-cell table hoisting — Term 5 tables are (g,t)-independent and Term 2 tables depend only on t — would cut the intrinsic per-cell GEMM factor by ~10x at 100k-unit scale if practitioners need routine 100k covariate fits. Follow-up lever from the v3.7 flop analysis; not needed below ~50k units. | `efficient_did_covariates.py::compute_conditional_cells_tiled` | CS-scaling | Mid | Low | | `ImputationDiD` dense `(A0'A0).toarray()` scales `O((U+T+K)^2)` — OOM risk on large panels (only triggers when the sparse solver fails). Needs an alternative dense fallback or richer sparse strategy. | `imputation.py` | #141 | Heavy | Medium | | CR2 Bell-McCaffrey DOF uses a naive `O(n²k)` per-coefficient loop over cluster pairs; Pustejovsky-Tipton (2018) Appendix B has a scores-based formulation avoiding the full `n×n` `M`. Switch when a user hits a large-`n` cluster-robust design. | `linalg.py::_compute_cr2_bm` | Phase 1a | Heavy | Low | | Rust-backend HC2: the Rust path only supports HC1; HC2 and CR2 Bell-McCaffrey fall through to NumPy. Noticeable for large-`n` fits. | `rust/src/linalg.rs` | Phase 1a | Mid | Low | diff --git a/diff_diff/efficient_did.py b/diff_diff/efficient_did.py index 0c0fadd11..29861c97a 100644 --- a/diff_diff/efficient_did.py +++ b/diff_diff/efficient_did.py @@ -33,6 +33,9 @@ EfficientDiDBootstrapMixin, ) from diff_diff.efficient_did_covariates import ( + OMEGA_RIDGE_DEFAULT, + _silverman_bandwidth, + compute_conditional_cells_tiled, compute_eif_cov, compute_generated_outcomes_cov, compute_omega_star_conditional, @@ -43,6 +46,7 @@ ) from diff_diff.efficient_did_results import EfficientDiDResults, HausmanPretestResult from diff_diff.efficient_did_weights import ( + _omega_star_nocov_gram, compute_efficient_weights, compute_eif_nocov, compute_generated_outcomes_nocov, @@ -284,6 +288,18 @@ class EfficientDiD(EfficientDiDBootstrapMixin): kernel_bandwidth : float or None Bandwidth for Gaussian kernel in conditional Omega* estimation. None = Silverman's rule-of-thumb (automatic). + omega_ridge : float, default ``OMEGA_RIDGE_DEFAULT`` (1e-6) + Relative ridge for the Omega* inversion behind the efficient + weights: solves ``(Omega* + omega_ridge * max(trace/H, 0) * I) x = 1`` + instead of inverting the numerically singular Omega* that PT-All's + telescoping overidentified moments produce. Stabilizes per-cell + ATT(g,t) against floating-point-level input/BLAS changes (1-ulp + stability ~1e-9 vs ~1e-4 for the legacy pseudoinverse) without + changing overall-ATT bias/RMSE/coverage (see REGISTRY.md). + ``omega_ridge=0`` restores the exact legacy inv/pinv code path + bit-for-bit - including its per-cell condition-number warnings and + the slow O(n^2 H^2) conditional-Omega* loops, so expect the legacy + runtime as well. Examples -------- @@ -309,6 +325,7 @@ def __init__( sieve_criterion: str = "bic", ratio_clip: float = 20.0, kernel_bandwidth: Optional[float] = None, + omega_ridge: float = OMEGA_RIDGE_DEFAULT, ): self.pt_assumption = pt_assumption self.alpha = alpha @@ -323,6 +340,7 @@ def __init__( self.sieve_criterion = sieve_criterion self.ratio_clip = ratio_clip self.kernel_bandwidth = kernel_bandwidth + self.omega_ridge = omega_ridge self.is_fitted_ = False self.results_: Optional[EfficientDiDResults] = None self._unit_resolved_survey = None @@ -361,6 +379,11 @@ def _validate_params(self) -> None: f"sieve_k_max must be a positive integer (or None for auto), " f"got {self.sieve_k_max}" ) + if not (np.isfinite(self.omega_ridge) and self.omega_ridge >= 0): + raise ValueError( + f"omega_ridge must be finite and >= 0 (0 = legacy inv/pinv path), " + f"got {self.omega_ridge}" + ) self._validate_vcov_type(self.vcov_type) @staticmethod @@ -424,6 +447,7 @@ def get_params(self) -> Dict[str, Any]: "sieve_criterion": self.sieve_criterion, "ratio_clip": self.ratio_clip, "kernel_bandwidth": self.kernel_bandwidth, + "omega_ridge": self.omega_ridge, } def set_params(self, **params: Any) -> "EfficientDiD": @@ -819,6 +843,36 @@ def fit( eif_by_gt: Dict[Tuple[Any, Any], np.ndarray] = {} stored_weights: Dict[Tuple[Any, Any], np.ndarray] = {} stored_cond: Dict[Tuple[Any, Any], float] = {} + # Ridge path: ill-conditioned cells consolidate into ONE fit-level + # warning (legacy omega_ridge=0 keeps the per-cell pinv warnings). + _ill_conditioned_cells: List[Tuple[Any, Any, float]] = [] + # Ridge-path covariate cells are deferred to a fused unit-tiled pass 2 + # (compute_conditional_cells_tiled) after the sieve caches are fully + # populated; each deferred cell holds an order-preserving placeholder + # in group_time_effects until finalization. + _cond_cell_specs: List[Dict[str, Any]] = [] + + def _finalize_cell(g: Any, att_gt: float, eif_vals: np.ndarray) -> Dict[str, Any]: + """Per-cell SE + inference record, shared by the inline and + deferred (pass-2) paths. + + Analytical SE = sqrt(mean(EIF^2) / n) [paper p.21]; with survey: + TSL variance via compute_survey_vcov. + """ + if self._unit_resolved_survey is not None: + se_gt = self._compute_survey_eif_se(eif_vals) + else: + se_gt = _compute_se_from_eif(eif_vals, n_units, unit_cluster_indices, n_clusters) + t_stat, p_val, ci = safe_inference(att_gt, se_gt, alpha=self.alpha, df=self._survey_df) + return { + "effect": att_gt, + "se": se_gt, + "t_stat": t_stat, + "p_value": p_val, + "conf_int": ci, + "n_treated": n_treated_per_g[g], + "n_control": n_control_count, + } for g in treatment_groups: # Under PT-Post, use per-group baseline Y_{g-1-anticipation} @@ -882,6 +936,21 @@ def fit( > 0 ] + # Ridge path: drop the degenerate same-cohort self-pair + # (g' = g, t_pre = t), which arises only for PRE-treatment + # cells (t < g). Its generated outcome telescopes to 0 = 0 + # identically (zero-information moment; the exact-null + # direction of Omega*). The legacy pseudoinverse truncates + # that direction, spreading weight over noisy moments; a + # ridge would instead load ~all weight on the zero-variance + # moment, collapsing pre-treatment placebos to a + # deterministic 0 and silently disabling the pre-trend + # diagnostic. Dropping the pair restores data-driven + # placebos. NOT applied at omega_ridge=0 (bit-identical + # legacy behavior). + if self.omega_ridge > 0 and pairs: + pairs = [(gp, tpre) for gp, tpre in pairs if not (gp == g and tpre == t)] + if not pairs: warnings.warn( f"No valid comparison pairs for (g={g}, t={t}). " "ATT will be NaN.", @@ -957,6 +1026,41 @@ def fit( basis_cache=sieve_basis_cache, ) + # Inverse propensity estimation (algorithm step 4) + # s_hat_{g'}(X) = 1/p_{g'}(X) for Eq 3.12 scaling. + # Populated BEFORE the ridge-path deferral so pass 2 sees + # complete caches (order swap vs the legacy sequence is + # inert: gen_out does not consume s_hat and vice versa). + for group_id in {g, np.inf} | {gp for gp, _ in pairs}: + if group_id not in s_hat_cache: + group_mask_s = ( + never_treated_mask if np.isinf(group_id) else cohort_masks[group_id] + ) + s_hat_cache[group_id] = estimate_inverse_propensity_sieve( + covariate_matrix, + group_mask_s, + k_max=self.sieve_k_max, + criterion=self.sieve_criterion, + unit_weights=unit_level_weights, + basis_cache=sieve_basis_cache, + ) + + if self.omega_ridge > 0: + # Fused tiled GEMM path: defer omega/weights/EIF to + # pass 2; placeholder preserves results ordering. + _cond_cell_specs.append( + { + "g": g, + "t": t, + "pairs": pairs, + "t_col": t_col_val, + "y1_col": effective_p1_col, + } + ) + group_time_effects[(g, t)] = None # type: ignore[assignment] + continue + + # ----- Legacy (omega_ridge=0) covariate path ----- # Per-unit DR generated outcomes: shape (n_units, H) gen_out = compute_generated_outcomes_cov( target_g=g, @@ -974,22 +1078,6 @@ def fit( y_hat = np.mean(gen_out, axis=0) # shape (H,) - # Inverse propensity estimation (algorithm step 4) - # s_hat_{g'}(X) = 1/p_{g'}(X) for Eq 3.12 scaling - for group_id in {g, np.inf} | {gp for gp, _ in pairs}: - if group_id not in s_hat_cache: - group_mask_s = ( - never_treated_mask if np.isinf(group_id) else cohort_masks[group_id] - ) - s_hat_cache[group_id] = estimate_inverse_propensity_sieve( - covariate_matrix, - group_mask_s, - k_max=self.sieve_k_max, - criterion=self.sieve_criterion, - unit_weights=unit_level_weights, - basis_cache=sieve_basis_cache, - ) - # Conditional Omega*(X) with per-unit propensities (Eq 3.12) omega_cond = compute_omega_star_conditional( target_g=g, @@ -1008,7 +1096,7 @@ def fit( ) # Per-unit weights: (n_units, H) - per_unit_w = compute_per_unit_weights(omega_cond) + per_unit_w = compute_per_unit_weights(omega_cond, omega_ridge=self.omega_ridge) # ATT = (survey-)weighted mean of per-unit DR scores if per_unit_w.shape[1] > 0: @@ -1025,8 +1113,14 @@ def fit( eif_vals = compute_eif_cov(per_unit_w, gen_out, att_gt, n_units) eif_by_gt[(g, t)] = eif_vals else: - # No-covariates path (closed-form) - omega = compute_omega_star_nocov( + # No-covariates path (closed-form). Ridge path uses the + # Gram/GEMM twin of the O(H^2) _sample_cov loop + # (reassociation-level agreement); omega_ridge=0 keeps + # the legacy loop verbatim. + _omega_fn = ( + _omega_star_nocov_gram if self.omega_ridge > 0 else compute_omega_star_nocov + ) + omega = _omega_fn( target_g=g, target_t=t, valid_pairs=pairs, @@ -1039,10 +1133,14 @@ def fit( unit_weights=unit_level_weights, ) - weights, _, cond_num = compute_efficient_weights(omega) + weights, _, cond_num = compute_efficient_weights( + omega, omega_ridge=self.omega_ridge + ) stored_weights[(g, t)] = weights if omega.size > 0: stored_cond[(g, t)] = cond_num + if self.omega_ridge > 0 and cond_num > 1e12: + _ill_conditioned_cells.append((g, t, cond_num)) y_hat = compute_generated_outcomes_nocov( target_g=g, @@ -1074,28 +1172,53 @@ def fit( ) eif_by_gt[(g, t)] = eif_vals - # Analytical SE = sqrt(mean(EIF^2) / n) [paper p.21] - # With survey: use TSL variance via compute_survey_vcov - if self._unit_resolved_survey is not None: - se_gt = self._compute_survey_eif_se(eif_vals) - else: - se_gt = _compute_se_from_eif( - eif_vals, n_units, unit_cluster_indices, n_clusters - ) + group_time_effects[(g, t)] = _finalize_cell(g, att_gt, eif_vals) - t_stat, p_val, ci = safe_inference( - att_gt, se_gt, alpha=self.alpha, df=self._survey_df - ) - - group_time_effects[(g, t)] = { - "effect": att_gt, - "se": se_gt, - "t_stat": t_stat, - "p_value": p_val, - "conf_int": ci, - "n_treated": int(np.sum(cohort_masks[g])), - "n_control": int(np.sum(never_treated_mask)), - } + # ----- Pass 2: fused tiled conditional path (ridge, covariates) ----- + if _cond_cell_specs: + assert covariate_matrix is not None + bandwidth = ( + self.kernel_bandwidth + if self.kernel_bandwidth is not None + else _silverman_bandwidth(covariate_matrix, unit_level_weights) + ) + cell_estimates = compute_conditional_cells_tiled( + _cond_cell_specs, + outcome_wide=outcome_wide, + covariate_matrix=covariate_matrix, + cohort_masks=cohort_masks, + never_treated_mask=never_treated_mask, + period_to_col=period_to_col, + cohort_fractions=cohort_fractions, + m_hat_cache=m_hat_cache, + r_hat_cache=r_hat_cache, + s_hat_cache=s_hat_cache, + bandwidth=bandwidth, + omega_ridge=self.omega_ridge, + unit_weights=unit_level_weights, + ) + for spec in _cond_cell_specs: + g = spec["g"] + t = spec["t"] + att_gt, eif_vals = cell_estimates[(g, t)] + eif_by_gt[(g, t)] = eif_vals + group_time_effects[(g, t)] = _finalize_cell(g, att_gt, eif_vals) + + # Fit-level consolidation of ill-conditioned Omega* cells (ridge path). + # Legacy (omega_ridge=0) warns per cell inside compute_efficient_weights. + if _ill_conditioned_cells: + _max_g, _max_t, _max_cond = max(_ill_conditioned_cells, key=lambda r: r[2]) + warnings.warn( + f"Omega* was ill-conditioned (cond > 1e12) in " + f"{len(_ill_conditioned_cells)} of {len(stored_cond)} (g, t) cells " + f"(max cond {_max_cond:.2e} at (g={_max_g}, t={_max_t})); the " + f"omega_ridge={self.omega_ridge:g} regularization handled these " + f"cells (expected under PT-All, whose overidentified moment set " + f"contains telescoping near-duplicate moments). Set omega_ridge=0 " + f"to restore the legacy pseudoinverse path.", + UserWarning, + stacklevel=2, + ) if not group_time_effects: raise ValueError( @@ -1271,6 +1394,7 @@ def fit( sieve_criterion=self.sieve_criterion, ratio_clip=self.ratio_clip, kernel_bandwidth=self.kernel_bandwidth, + omega_ridge=self.omega_ridge, survey_metadata=( self._recompute_unit_survey_metadata(survey_metadata) if survey_metadata is not None @@ -1736,7 +1860,8 @@ def hausman_pretest( alpha : float Significance level for the test. **nuisance_kwargs - Passed to both fits (e.g. ``sieve_k_max``, ``ratio_clip``). + Passed to both fits (e.g. ``sieve_k_max``, ``ratio_clip``, + ``omega_ridge``). Returns ------- diff --git a/diff_diff/efficient_did_covariates.py b/diff_diff/efficient_did_covariates.py index 393f98d84..0158db26a 100644 --- a/diff_diff/efficient_did_covariates.py +++ b/diff_diff/efficient_did_covariates.py @@ -28,6 +28,21 @@ from diff_diff.linalg import solve_ols +# Default ridge for the Omega* inversion (see ``compute_per_unit_weights`` / +# ``compute_efficient_weights``). Under PT-All the overidentified moment set +# contains telescoping near-duplicate moments, so sample Omega* is numerically +# singular (cond ~1e17-1e22 on realistic panels); the legacy pseudoinverse sits +# on the rcond-cutoff cliff, where any last-digit change to Omega* (BLAS +# reordering, platform change, 1-ulp data perturbation) moves per-cell weights +# and ATT(g,t) at the ~1e-2 relative level. The ridge +# ``Omega + lam * max(trace/H, 0) * I`` damps the statistically-null +# directions smoothly: 1-ulp per-cell stability improves from ~1e-4 to ~1e-9 +# at this default, while overall-ATT bias/RMSE/coverage are unchanged +# (Monte Carlo, see REGISTRY.md). Calibrated 2026-07 against 1-ulp stability +# (target <=1e-6, achieved ~3e-9) and the HRS Table 6 anchors (all within the +# published-value tolerance at every candidate lambda). +OMEGA_RIDGE_DEFAULT = 1e-6 + # --------------------------------------------------------------------------- # Outcome regression # --------------------------------------------------------------------------- @@ -1080,27 +1095,377 @@ def compute_omega_star_conditional( return omega +# --------------------------------------------------------------------------- +# Fused tiled GEMM path: conditional Omega* + ridge weights + DR scores +# --------------------------------------------------------------------------- + +# Memory cap for the unit-tiled conditional path. Resolved at CALL time inside +# compute_conditional_cells_tiled (never bound as a def-time default) so tests +# can monkeypatch it to force multi-tile execution. +_TARGET_OMEGA_TILE_BYTES = 256 * 1024 * 1024 + + +def _kcov_batch(W: np.ndarray, A: np.ndarray, B: np.ndarray) -> np.ndarray: + """Batched kernel-weighted local covariances via GEMM. + + For row-normalized kernel weights ``W`` (each row sums to 1), + + KCov_i(A, B) = sum_j W_ij (A_j - Abar_i)(B_j - Bbar_i) + = (W @ (A0 * B0))_i - (W @ A0)_i * (W @ B0)_i + + with ``A0 = A - mean(A)`` (global pre-centering keeps the + uncentered-product form cancellation-safe; any constant shift cancels in + the identity). Row normalization is per-row, so the identity holds for + survey-weighted kernels and for any row-tile of ``W``. + + Parameters + ---------- + W : ndarray, shape (n_rows, n_group) + Row-normalized kernel weight matrix (or a row-tile of one). + A, B : ndarray, shape (n_group, m) + Column-matched inputs; column ``j`` of the output is + ``KCov_i(A[:, j], B[:, j])``. + + Returns + ------- + ndarray, shape (n_rows, m) + """ + A0 = A - A.mean(axis=0) + B0 = B - B.mean(axis=0) + return W @ (A0 * B0) - (W @ A0) * (W @ B0) + + +def compute_conditional_cells_tiled( + cell_specs: List[Dict], + outcome_wide: np.ndarray, + covariate_matrix: np.ndarray, + cohort_masks: Dict[float, np.ndarray], + never_treated_mask: np.ndarray, + period_to_col: Dict[float, int], + cohort_fractions: Dict[float, float], + m_hat_cache: Dict[Tuple, np.ndarray], + r_hat_cache: Dict[Tuple[float, float], np.ndarray], + s_hat_cache: Dict[float, np.ndarray], + bandwidth: float, + omega_ridge: float, + unit_weights: Optional[np.ndarray] = None, + never_treated_val: float = np.inf, + tile_bytes: Optional[int] = None, +) -> Dict[Tuple[float, float], Tuple[float, np.ndarray]]: + """Fused conditional-path estimation for all (g, t) cells, unit-tiled. + + Replaces the legacy per-cell chain (dense ``compute_omega_star_conditional`` + H^2 loop -> per-unit SVD/pinv weights -> EIF) with, per unit-tile: + + 1. ONE row-normalized kernel weight matrix per comparison group + (``_kernel_weights_matrix``), reused across ALL cells - the matrices + depend only on covariates/bandwidth/group, not on (g, t). + 2. Per cell, the five Eq 3.12 terms as kernel-covariance TABLES via + :func:`_kcov_batch` - Term 2/5 covariances depend only on the + ``(t_pre_j, t_pre_k)`` column combination, so the H(H+1)/2 pair loop + dedups to ~T^2 GEMM columns - then a gather-assembly of the + ``(tile, H, H)`` Omega* block. + 3. Ridge-regularized batched weights (:func:`_ridge_solve_weights`; + requires ``omega_ridge > 0`` - the legacy ``omega_ridge=0`` path never + reaches this function). + 4. DR generated outcomes on the tile rows (reusing + :func:`compute_generated_outcomes_cov` on row-sliced views - gen_out + is row-separable across units) and per-unit weighted scores. + + ``att_gt`` is finalized only after ALL tiles complete (survey-weighted + mean of the full score vector), and ``EIF = scores - att_gt`` is computed + once from the finalized value - never per tile. + + Cells with ``H == 1`` skip kernel/omega assembly entirely (weights are + trivially 1); the results are omega-independent, matching the legacy path. + + Parameters + ---------- + cell_specs : list of dict + One per estimable cell: ``{"g", "t", "pairs", "t_col", "y1_col"}`` + (``y1_col`` varies per cohort under ``pt_assumption="post"``). + ``pairs`` must be non-empty. + tile_bytes : int, optional + Memory budget for per-tile state; ``None`` resolves the module + constant ``_TARGET_OMEGA_TILE_BYTES`` at call time. + (remaining parameters as in ``compute_omega_star_conditional``) + + Returns + ------- + dict mapping ``(g, t)`` to ``(att_gt, eif_values)``. + """ + n_units = outcome_wide.shape[0] + if not cell_specs: + return {} + + if n_units > 50_000: + warnings.warn( + f"Conditional Omega* estimation with n={n_units} units is " + f"expensive (the kernel weight matrices are intrinsically " + f"O(n^2) in memory traffic, computed in bounded tiles).", + UserWarning, + stacklevel=2, + ) + + if tile_bytes is None: + tile_bytes = _TARGET_OMEGA_TILE_BYTES + + h_max = max(len(spec["pairs"]) for spec in cell_specs) + # Per-tile-row footprint: (H, H) omega block + kernel-matrix rows across + # all groups (sum of group sizes <= n_units, counted twice for the cdist + # temp) + small per-cell vectors. + row_bytes = 8 * (h_max * h_max + 2 * n_units + 4 * h_max) + tile_units = int(max(1, min(n_units, tile_bytes // max(1, row_bytes)))) + + # Group-level (non-tiled) data, shared across cells and tiles + group_keys: set = set() + for spec in cell_specs: + group_keys.add(spec["g"]) + group_keys.add(never_treated_val) + for gp, _ in spec["pairs"]: + group_keys.add(never_treated_val if np.isinf(gp) else gp) + + def _mask_for(gkey: float) -> np.ndarray: + return never_treated_mask if np.isinf(gkey) else cohort_masks[gkey] + + y_grp: Dict[float, np.ndarray] = {} + x_grp: Dict[float, np.ndarray] = {} + w_grp: Dict[float, Optional[np.ndarray]] = {} + s_all: Dict[float, np.ndarray] = {} + for gkey in group_keys: + mask = _mask_for(gkey) + y_grp[gkey] = outcome_wide[mask] + x_grp[gkey] = covariate_matrix[mask] + w_grp[gkey] = unit_weights[mask] if unit_weights is not None else None + s_all[gkey] = s_hat_cache.get( + gkey, + np.full(n_units, 1.0 / max(cohort_fractions.get(gkey, 1e-10), 1e-10)), + ) + + scores: Dict[Tuple[float, float], np.ndarray] = { + (spec["g"], spec["t"]): np.empty(n_units) for spec in cell_specs + } + + for lo in range(0, n_units, tile_units): + hi = min(lo + tile_units, n_units) + + # Kernel weight matrices for this tile, one per group, lazily built + # and shared across ALL cells. + w_tile_cache: Dict[float, np.ndarray] = {} + + def _w_tile(gkey: float) -> np.ndarray: + if gkey not in w_tile_cache: + w_tile_cache[gkey] = _kernel_weights_matrix( + covariate_matrix[lo:hi], + x_grp[gkey], + bandwidth, + group_weights=w_grp[gkey], + ) + return w_tile_cache[gkey] + + # Row-sliced views of the nuisance caches / masks for gen_out reuse + masks_tile = {k: v[lo:hi] for k, v in cohort_masks.items()} + nt_mask_tile = never_treated_mask[lo:hi] + m_hat_tile = {k: v[lo:hi] for k, v in m_hat_cache.items()} + r_hat_tile = {k: v[lo:hi] for k, v in r_hat_cache.items()} + outcome_tile = outcome_wide[lo:hi] + + for spec in cell_specs: + g = spec["g"] + t = spec["t"] + pairs = spec["pairs"] + t_col = spec["t_col"] + y1_col = spec["y1_col"] + H = len(pairs) + + gen_out_tile = compute_generated_outcomes_cov( + target_g=g, + target_t=t, + valid_pairs=pairs, + outcome_wide=outcome_tile, + cohort_masks=masks_tile, + never_treated_mask=nt_mask_tile, + period_to_col=period_to_col, + period_1_col=y1_col, + cohort_fractions=cohort_fractions, + m_hat_cache=m_hat_tile, + r_hat_cache=r_hat_tile, + never_treated_val=never_treated_val, + ) + + if H == 1: + scores[(g, t)][lo:hi] = gen_out_tile[:, 0] + continue + + pcols = sorted({period_to_col[tp] for _, tp in pairs}) + pindex = {c: i for i, c in enumerate(pcols)} + n_p = len(pcols) + j_p = [pindex[period_to_col[tp]] for _, tp in pairs] + j_gp = [gp for gp, _ in pairs] + + iu, ju = np.triu_indices(n_p) + combo_idx = np.zeros((n_p, n_p), dtype=int) + combo_idx[iu, ju] = np.arange(len(iu)) + combo_idx[ju, iu] = combo_idx[iu, ju] + + s_g_tile = s_all[g][lo:hi] + s_inf_tile = s_all[never_treated_val][lo:hi] + wg = _w_tile(g) + winf = _w_tile(never_treated_val) + + y_g = y_grp[g] + y_inf = y_grp[never_treated_val] + a_g = y_g[:, t_col] - y_g[:, y1_col] + + # Term 1: s_g * KCov(a, a | g) - shared by every (j, k) + term1 = s_g_tile * _kcov_batch(wg, a_g[:, None], a_g[:, None])[:, 0] + + # Term 2 table: s_inf * KCov(Y_t - Y_pj, Y_t - Y_pk | inf) + u_inf = y_inf[:, [t_col]] - y_inf[:, pcols] + t2 = s_inf_tile[:, None] * _kcov_batch(winf, u_inf[:, iu], u_inf[:, ju]) + + # Term 3/4 table: s_g * KCov(a, Y_p - Y_1 | g), only if any + # same-cohort pair exists + t34 = None + if any(gp == g for gp in j_gp): + b_g = y_g[:, pcols] - y_g[:, [y1_col]] + t34 = s_g_tile[:, None] * _kcov_batch(wg, np.repeat(a_g[:, None], n_p, axis=1), b_g) + + # Term 5 tables per distinct comparison group + t5: Dict[float, np.ndarray] = {} + for gp in set(j_gp): + gkey = never_treated_val if np.isinf(gp) else gp + if gkey in t5: + continue + y_c = y_grp[gkey] + v_c = y_c[:, pcols] - y_c[:, [y1_col]] + t5[gkey] = s_all[gkey][lo:hi][:, None] * _kcov_batch( + _w_tile(gkey), v_c[:, iu], v_c[:, ju] + ) + + omega_tile = np.empty((hi - lo, H, H)) + for j in range(H): + pj = j_p[j] + gpj = j_gp[j] + for k in range(j, H): + pk = j_p[k] + gpk = j_gp[k] + val = term1 + t2[:, combo_idx[pj, pk]] + if gpj == g: + val = val - t34[:, pj] + if gpk == g: + val = val - t34[:, pk] + if gpj == gpk: + gkey = never_treated_val if np.isinf(gpj) else gpj + val = val + t5[gkey][:, combo_idx[pj, pk]] + omega_tile[:, j, k] = val + if j != k: + omega_tile[:, k, j] = val + + w_units = _ridge_solve_weights(omega_tile, omega_ridge) + scores[(g, t)][lo:hi] = np.sum(w_units * gen_out_tile, axis=1) + + out: Dict[Tuple[float, float], Tuple[float, np.ndarray]] = {} + for spec in cell_specs: + gt = (spec["g"], spec["t"]) + s = scores[gt] + if unit_weights is not None: + att_gt = float(np.average(s, weights=unit_weights)) + else: + att_gt = float(np.mean(s)) + out[gt] = (att_gt, s - att_gt) + return out + + # --------------------------------------------------------------------------- # Per-unit efficient weights from conditional Omega* # --------------------------------------------------------------------------- +def _ridge_solve_weights(omega_stack: np.ndarray, omega_ridge: float) -> np.ndarray: + """Batched ridge-regularized efficient weights for a stack of Omega*. + + Solves ``(Omega_i + lam * max(trace(Omega_i)/H, 0) * I) x = 1`` per unit + and normalizes ``w_i = x / (1'x)``. The trace-scaled ridge is + scale-equivariant and O(H) to compute (no SVD); the ``max(..., 0)`` guards + a machine-noise-negative trace from producing a negative ridge. Special + cases match the legacy path: an (approximately) all-zero matrix or a + near-zero denominator yields uniform weights ``1/H``. + + Parameters + ---------- + omega_stack : ndarray, shape (n, H, H) + Stack of covariance matrices (H >= 2). + omega_ridge : float + Relative ridge scale ``lam > 0``. + + Returns + ------- + weights : ndarray, shape (n, H) + """ + n, H, _ = omega_stack.shape + ones = np.ones(H) + weights = np.full((n, H), 1.0 / H) + + # np.allclose(omega_i, 0.0) == elementwise |x| <= atol (1e-8), rtol inert + zero_mask = np.all(np.abs(omega_stack) <= 1e-8, axis=(1, 2)) + rest = np.flatnonzero(~zero_mask) + if rest.size == 0: + return weights + + om = omega_stack[rest] + trace = np.trace(om, axis1=1, axis2=2) + scale = np.maximum(trace / H, 0.0) + om_ridged = om + (omega_ridge * scale)[:, None, None] * np.eye(H)[None] + try: + # gufunc solve treats a 2-D rhs as a matrix; a stack of vectors needs + # shape (n, H, 1) + num = np.linalg.solve(om_ridged, np.ones((rest.size, H, 1)))[..., 0] + except np.linalg.LinAlgError: + # Unreachable for PSD Omega* with lam > 0; kept as a per-unit + # minimum-norm fallback so one pathological unit cannot fail the batch. + num = np.empty((rest.size, H)) + for m in range(rest.size): + try: + num[m] = np.linalg.solve(om_ridged[m], ones) + except np.linalg.LinAlgError: + num[m] = np.linalg.pinv(om_ridged[m]) @ ones + den = num @ ones + ok = np.abs(den) >= 1e-15 + solved = np.full((rest.size, H), 1.0 / H) + solved[ok] = num[ok] / den[ok, None] + weights[rest] = solved + return weights + + def compute_per_unit_weights( omega_conditional: np.ndarray, cond_threshold: float = 1e12, + omega_ridge: float = 0.0, ) -> np.ndarray: """Per-unit efficient weights from conditional Omega* inverse. ``w(X_i) = 1' Omega*(X_i)^{-1} / (1' Omega*(X_i)^{-1} 1)`` - Falls back to pseudoinverse per unit if condition number exceeds threshold. + With ``omega_ridge = 0`` (this helper's default), runs the legacy per-unit + loop: exact inverse, falling back to a pseudoinverse when the condition + number exceeds ``cond_threshold``. On the numerically singular Omega* + produced by PT-All's telescoping moments, that pseudoinverse is + cutoff-cliff sensitive (see ``OMEGA_RIDGE_DEFAULT``). With + ``omega_ridge > 0``, uses the batched ridge solve + ``(Omega_i + lam * max(trace/H, 0) * I) x = 1`` instead - numerically + stable and vectorized over units. The estimator passes its + ``omega_ridge`` parameter (default ``OMEGA_RIDGE_DEFAULT``); standalone + callers of this helper keep exact legacy behavior unless they opt in. Parameters ---------- omega_conditional : ndarray, shape (n_units, H, H) Per-unit conditional covariance matrices. cond_threshold : float - Condition number threshold for pseudoinverse fallback. + Condition number threshold for pseudoinverse fallback (legacy path). + omega_ridge : float + Relative ridge scale; 0 = legacy inv/pinv path. Returns ------- @@ -1113,6 +1478,9 @@ def compute_per_unit_weights( if H == 1: return np.ones((n_units, 1)) + if omega_ridge > 0: + return _ridge_solve_weights(omega_conditional, omega_ridge) + ones = np.ones(H) weights = np.zeros((n_units, H)) diff --git a/diff_diff/efficient_did_results.py b/diff_diff/efficient_did_results.py index 83fdd2fc7..12627b697 100644 --- a/diff_diff/efficient_did_results.py +++ b/diff_diff/efficient_did_results.py @@ -134,6 +134,9 @@ class EfficientDiDResults: Clipping bound for sieve propensity ratios. kernel_bandwidth : float or None Bandwidth used for kernel-smoothed conditional Omega*. + omega_ridge : float + Relative ridge used for the Omega* inversion behind the efficient + weights (0 = legacy exact-inverse/pseudoinverse path). """ group_time_effects: Dict[Tuple[Any, Any], Dict[str, Any]] @@ -184,6 +187,7 @@ class EfficientDiDResults: sieve_criterion: str = "bic" ratio_clip: float = 20.0 kernel_bandwidth: Optional[float] = None + omega_ridge: float = 0.0 # Survey design metadata (SurveyMetadata instance from diff_diff.survey) survey_metadata: Optional[Any] = field(default=None) diff --git a/diff_diff/efficient_did_weights.py b/diff_diff/efficient_did_weights.py index 685feaac8..e1e60a3f3 100644 --- a/diff_diff/efficient_did_weights.py +++ b/diff_diff/efficient_did_weights.py @@ -71,8 +71,11 @@ def enumerate_valid_triples( # Including g'=∞ (never-treated) produces moments where the second # and third terms telescope: y_hat = E[Y_t-Y_1|G=g] - E[Y_t-Y_1|G=∞] # regardless of t_pre. These redundant moments add no information - # beyond the basic 2x2 DiD; Omega*'s pseudoinverse assigns them - # zero effective weight. Retained for implementation simplicity. + # beyond the basic 2x2 DiD; the default omega_ridge damps them (the + # legacy omega_ridge=0 pseudoinverse truncates them instead - see the + # Omega* ridge Note in REGISTRY.md). Retained for implementation + # simplicity. The estimator's ridge path additionally drops the + # degenerate PRE-treatment self-pair (g'=g, t_pre=t) at fit level. candidate_groups: List[float] = [never_treated_val] for gp in treatment_groups: candidate_groups.append(gp) @@ -318,27 +321,177 @@ def compute_omega_star_nocov( return omega +def _sample_cov_cols( + A: np.ndarray, + B: np.ndarray, + w: Optional[np.ndarray] = None, +) -> np.ndarray: + """Column-batched :func:`_sample_cov` (identical semantics per column). + + A, B : (n, m) column-matched; returns (m,) with + ``out[j] = _sample_cov(A[:, j], B[:, j], w)`` - including the guards: + fewer than 2 rows -> 0.0, and (weighted) ``sum(w) <= 1`` -> 0.0. + """ + n, m = A.shape + if n < 2: + return np.zeros(m) + if w is None: + A0 = A - A.mean(axis=0) + B0 = B - B.mean(axis=0) + return np.einsum("ij,ij->j", A0, B0) / (n - 1) + sum_w = float(np.sum(w)) + if sum_w <= 1.0: + return np.zeros(m) + wn = w / sum_w + A0 = A - wn @ A + B0 = B - wn @ B + return np.einsum("i,ij,ij->j", w, A0, B0) / (sum_w - 1.0) + + +def _omega_star_nocov_gram( + target_g: float, + target_t: float, + valid_pairs: List[Tuple[float, float]], + outcome_wide: np.ndarray, + cohort_masks: Dict[float, np.ndarray], + never_treated_mask: np.ndarray, + period_to_col: Dict[float, int], + period_1_col: int, + cohort_fractions: Dict[float, float], + never_treated_val: float = np.inf, + unit_weights: Optional[np.ndarray] = None, +) -> np.ndarray: + """Gram/GEMM form of :func:`compute_omega_star_nocov` (v3.7 fast path). + + Same five-term structure and guard semantics, but the H(H+1)/2 pair loop + of :func:`_sample_cov` calls collapses to column-batched covariances over + the DISTINCT ``(t_pre_j, t_pre_k)`` combinations (Term 2/5 covariances + depend only on the t_pre columns, not the full pair), assembled by + lookup. Used by the estimator's default ridge path; ``omega_ridge=0`` + keeps the legacy loop. Agreement with the loop is at BLAS-reassociation + level (~1e-15 relative). + """ + H = len(valid_pairs) + if H == 0: + return np.empty((0, 0)) + + t_col = period_to_col[target_t] + y1_col = period_1_col + + g_mask = cohort_masks[target_g] + Y_g = outcome_wide[g_mask] + pi_g = cohort_fractions[target_g] + w_g = unit_weights[g_mask] if unit_weights is not None else None + + Y_inf = outcome_wide[never_treated_mask] + pi_inf = cohort_fractions.get(never_treated_val, 0.0) + w_inf = unit_weights[never_treated_mask] if unit_weights is not None else None + + pcols = sorted({period_to_col[tp] for _, tp in valid_pairs}) + pindex = {c: i for i, c in enumerate(pcols)} + n_p = len(pcols) + j_p = [pindex[period_to_col[tp]] for _, tp in valid_pairs] + j_gp = [gp for gp, _ in valid_pairs] + + iu, ju = np.triu_indices(n_p) + combo_idx = np.zeros((n_p, n_p), dtype=int) + combo_idx[iu, ju] = np.arange(len(iu)) + combo_idx[ju, iu] = combo_idx[iu, ju] + + a_g = Y_g[:, t_col] - Y_g[:, y1_col] + term1 = 0.0 + if pi_g > 0: + term1 = (1.0 / pi_g) * float(_sample_cov_cols(a_g[:, None], a_g[:, None], w_g)[0]) + + # Term 2 / Term 5 (never-treated) tables - _sample_cov_cols returns zeros + # for n_inf < 2, matching the legacy table-not-populated skip. + u_inf = Y_inf[:, [t_col]] - Y_inf[:, pcols] + v_inf = Y_inf[:, pcols] - Y_inf[:, [y1_col]] + t2 = _sample_cov_cols(u_inf[:, iu], u_inf[:, ju], w_inf) + t5_inf = _sample_cov_cols(v_inf[:, iu], v_inf[:, ju], w_inf) + + # Term 3/4 table (same-cohort pairs only) + t34 = None + if pi_g > 0 and any(gp == target_g for gp in j_gp): + b_g = Y_g[:, pcols] - Y_g[:, [y1_col]] + t34 = _sample_cov_cols(np.repeat(a_g[:, None], n_p, axis=1), b_g, w_g) + + # Term 5 tables per finite comparison cohort + t5: Dict[float, np.ndarray] = {} + for gp in set(j_gp): + if np.isinf(gp) or gp in t5: + continue + pi_gp = cohort_fractions.get(gp, 0.0) + if pi_gp <= 0 or gp not in cohort_masks: + continue + Y_c = outcome_wide[cohort_masks[gp]] + w_c = unit_weights[cohort_masks[gp]] if unit_weights is not None else None + v_c = Y_c[:, pcols] - Y_c[:, [y1_col]] + t5[gp] = _sample_cov_cols(v_c[:, iu], v_c[:, ju], w_c) + + omega = np.empty((H, H)) + for j in range(H): + pj = j_p[j] + gpj = j_gp[j] + for k in range(j, H): + pk = j_p[k] + gpk = j_gp[k] + val = term1 + if pi_inf > 0: + val += (1.0 / pi_inf) * t2[combo_idx[pj, pk]] + if gpj == target_g and t34 is not None: + val -= (1.0 / pi_g) * t34[pj] + if gpk == target_g and t34 is not None: + val -= (1.0 / pi_g) * t34[pk] + if gpj == gpk: + if np.isinf(gpj): + if pi_inf > 0: + val += (1.0 / pi_inf) * t5_inf[combo_idx[pj, pk]] + elif gpj in t5: + val += (1.0 / cohort_fractions[gpj]) * t5[gpj][combo_idx[pj, pk]] + omega[j, k] = val + if j != k: + omega[k, j] = val + + return omega + + def compute_efficient_weights( omega_star: np.ndarray, cond_threshold: float = 1e12, + omega_ridge: float = 0.0, ) -> Tuple[np.ndarray, bool, float]: """Compute efficient weights from Omega* inverse (Eq 3.13 / 4.3). ``w = ones @ inv(Omega*) / (ones @ inv(Omega*) @ ones)`` + With ``omega_ridge = 0`` (this helper's default), runs the legacy path: + exact inverse, pseudoinverse + per-cell warning when the condition number + exceeds ``cond_threshold``. With ``omega_ridge > 0``, solves the + ridge-regularized system ``(Omega* + lam * max(trace/H, 0) * I) x = 1`` + instead - numerically stable on the singular Omega* produced by PT-All's + telescoping moments (see ``OMEGA_RIDGE_DEFAULT`` in + ``efficient_did_covariates``) - and emits no per-cell warning (the + estimator consolidates ill-conditioning into one fit-level warning). The + condition number is still computed and returned for diagnostics either + way. + Parameters ---------- omega_star : ndarray, shape (H, H) Covariance matrix from :func:`compute_omega_star_nocov`. cond_threshold : float - If condition number exceeds this, use pseudoinverse + warning. + If condition number exceeds this: legacy path uses pseudoinverse + + warning; ridge path only reports it via the returned cond. + omega_ridge : float + Relative ridge scale; 0 = legacy inv/pinv path. Returns ------- weights : ndarray, shape (H,) Efficient combination weights (sum to 1). used_pinv : bool - True if pseudoinverse was used. + True if pseudoinverse was used (always False on the ridge path). cond_number : float Condition number of Omega* (avoids recomputation by caller). """ @@ -361,6 +514,24 @@ def compute_efficient_weights( return ones / H, False, np.inf cond = float(np.linalg.cond(omega_star)) + + if omega_ridge > 0: + scale = max(float(np.trace(omega_star)) / H, 0.0) + om_ridged = omega_star + omega_ridge * scale * np.eye(H) + try: + numerator = np.linalg.solve(om_ridged, ones) + except np.linalg.LinAlgError: + # Unreachable for PSD Omega* with lam > 0; minimum-norm backstop. + numerator = np.linalg.pinv(om_ridged) @ ones + denominator = float(numerator @ ones) + if abs(denominator) < 1e-15: + warnings.warn( + "Denominator of efficient weights is near zero; using uniform weights.", + UserWarning, + stacklevel=2, + ) + return ones / H, False, cond + return numerator / denominator, False, cond if cond > cond_threshold: warnings.warn( f"Omega* condition number ({cond:.2e}) exceeds threshold " diff --git a/diff_diff/guides/llms-full.txt b/diff_diff/guides/llms-full.txt index 9f921ec7b..3a09b2572 100644 --- a/diff_diff/guides/llms-full.txt +++ b/diff_diff/guides/llms-full.txt @@ -885,6 +885,7 @@ EfficientDiD( bootstrap_weights: str = "rademacher", # "rademacher", "mammen", or "webb" seed: int | None = None, anticipation: int = 0, + omega_ridge: float = 1e-6, # Ridge for the Omega* inversion behind the efficient weights (PT-All's overidentified moments make sample Omega* numerically singular); 0 = legacy exact-inverse/pseudoinverse path. See the Omega* ridge Note in the methodology registry. ) ``` diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 5adf49f7f..75eb1babc 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -1240,7 +1240,7 @@ where `q_{g,e} = pi_g / sum_{g' in G_{trt,e}} pi_{g'}`. *Edge cases:* - **Single pre-treatment period (g=2)**: `V*_{gt}(X)` is 1x1, efficient weights are trivially 1, estimator collapses to standard DiD with single baseline -- **Rank deficiency in `V*_{gt}(X)` or `Omega*_{gt}(X)`**: Inverse does not exist if outcome changes are linearly dependent conditional on covariates. Detect via matrix condition number; fall back to pseudoinverse or standard estimator +- **Rank deficiency in `V*_{gt}(X)` or `Omega*_{gt}(X)`**: Inverse does not exist if outcome changes are linearly dependent conditional on covariates. Under the default `omega_ridge > 0` the ridge-regularized solve (see the Omega* ridge Note below) handles this without a routing cliff; at `omega_ridge=0` the legacy behavior applies (detect via matrix condition number; fall back to pseudoinverse with a per-cell warning) - **Near-zero propensity scores**: Ratio `p_g(X)/p_{g'}(X)` explodes. Overlap assumption (O) rules this out in population; implement trimming or warn on finite-sample instability - **Note:** When no sieve degree K succeeds for ratio estimation (basis dimension exceeds comparison group size, or all linear systems are singular), the estimator falls back to a constant ratio of 1 for all units with a UserWarning. The outcome regression adjustment remains active, so the generated outcomes (Eq 4.4) still incorporate covariate information via the m_hat terms. The DR property ensures consistency as long as the outcome regression is correctly specified. - **Note:** When no sieve degree K succeeds for inverse propensity estimation (algorithm step 4), the estimator falls back to unconditional n/n_group scaling with a UserWarning, which reduces to the unconditional Omega* approximation for the affected group. @@ -1248,7 +1248,8 @@ where `q_{g,e} = pi_g / sum_{g' in G_{trt,e}} pi_{g'}`. - **Negative weights**: Explicitly stated as harmless for bias and beneficial for precision; arise from efficiency optimization under overidentification (Section 5.2) - **PT-Post regime (just-identified)**: Under PT-Post, EDiD automatically reduces to standard single-baseline estimator (Corollary 3.2). No downside to using EDiD -- it subsumes standard estimators - **Duplicate rows**: Duplicate `(unit, time)` entries are rejected with `ValueError`. The estimator requires exactly one observation per unit-period -- **Note:** PT-All index set includes g'=∞ (never-treated) as a candidate comparison group and excludes period_1 for all g'. When g'=∞, the second and third Eq 3.9 terms telescope so all (∞, t_pre) moments produce the same 2x2 DiD value; these redundant moments are handled by Omega*'s pseudoinverse. When t_pre = period_1, the third term degenerates to E[Y_1 - Y_1 | G=g'] = 0 for any g', adding no information. Valid pairs require only t_pre < g' (pre-treatment for comparison group), not t_pre < g. Same-group pairs (g'=g) are valid and contribute overidentifying moments (Equation 3.9). +- **Note:** PT-All index set includes g'=∞ (never-treated) as a candidate comparison group and excludes period_1 for all g'. When g'=∞, the second and third Eq 3.9 terms telescope so all (∞, t_pre) moments produce the same 2x2 DiD value; these redundant moments are damped by the default Omega* ridge (see the Omega* ridge Note below; at `omega_ridge=0`, by the legacy pseudoinverse). When t_pre = period_1, the third term degenerates to E[Y_1 - Y_1 | G=g'] = 0 for any g', adding no information. Valid pairs require only t_pre < g' (pre-treatment for comparison group), not t_pre < g. Same-group pairs (g'=g) are valid and contribute overidentifying moments (Equation 3.9) — except the degenerate PRE-treatment self-pair (g'=g, t_pre=t), which the default ridge path excludes (see the ridge Note). +- **Note:** Omega* ridge regularization (`omega_ridge`, default 1e-6; v3.7) — a documented refinement of the Omega* inversion in Eq 3.5/3.13/4.3, in the space the paper leaves open (it assumes Omega* is invertible and does not prescribe finite-sample handling of singularity). Under PT-All the overidentified moment set makes the SAMPLE Omega* numerically singular by construction (the telescoping (∞, t_pre) moments above plus near-duplicate cross-cohort moments): measured cond 1e17–1e22 for 100% of units on realistic panels, with a spectrum of one exact-null direction (the degenerate self-pair, below) plus a cluster of statistically-null directions at relative eigenvalue 1e-5–1e-8, far below the ~1e-2 sampling noise of the covariance entries — and no clean spectral gap. The prior pseudoinverse fallback therefore sat on the rcond-cutoff cliff: ANY floating-point-level change (BLAS reordering, platform change, a 1-ulp data perturbation) redrew per-cell weights and moved per-cell ATT(g,t) at up to ~1e-2 relative (measured 1.2e-4 per-cell rel for a 1-ulp outcome perturbation on the pre-v3.7 code; overall ATT stable at ~1e-9 because the redraw averages out across units and cells). The ridge solves `(Omega* + omega_ridge * max(trace/H, 0) * I) x = 1` (trace-scaled: scale-equivariant, O(H), no SVD) — a smooth regularization with sensitivity bounded by 1/omega_ridge, not a cutoff. **Why the deviation is safe:** every moment individually identifies ATT(g,t), so any fixed weights summing to 1 keep the estimator consistent; the ridge trades a numerically ill-defined efficiency optimum for a stable one, and the plug-in EIF treatment of estimated weights (Remark 4.2) is unaffected. **Calibration evidence (2026-07):** default 1e-6 chosen by 1-ulp stability (per-cell rel <= 3e-9 vs 1.2e-4 legacy; candidates 1e-4/1e-6/1e-8 all pass the <=1e-6 target) with the HRS Table 6 anchors as a hard gate (all anchors within the published-value tolerance at every candidate; worst deviation 0.0257 SE, unchanged from legacy; shift <= 0.0001 SE); Monte Carlo on covariate-confounded DGPs shows bias/RMSE/SE-calibration/coverage statistically identical to legacy (ridge marginally better point metrics). **One-time value shift:** per-cell and event-study values on the covariate path move within the pre-existing indeterminacy band when upgrading (worst observed post-treatment cell shift ~0.6 of its own SE at n=500, shrinking with n: overall-ATT shift 1.6e-2 → 4e-3 → 1.1e-3 rel at n=500/1k/2k); the no-covariates path is essentially unchanged (~1e-7). `omega_ridge=0` restores the ENTIRE legacy code path bit-for-bit (both the omega constructions and the inv/pinv weights, including per-cell condition-number warnings and the legacy O(n^2 H^2) runtime). **Degenerate self-pair:** for PRE-treatment cells (t < g), the pair (g'=g, t_pre=t) telescopes to the identically-zero moment 0=0 (the exact-null Omega* direction). The legacy pseudoinverse truncated it, spreading weight over noisy moments (spurious pre-treatment placebos of ~5% of the effect size, pure noise amplification); a naive ridge would instead load all weight on the zero-variance moment, making placebos deterministically zero and silently disabling the pre-trend diagnostic. The default ridge path therefore drops this zero-information pair (fit-level filter; post-treatment cells never contain it), restoring honest data-driven placebos; `omega_ridge=0` keeps the legacy pair set. **Warnings/diagnostics:** the no-covariates path still computes per-cell condition numbers (cheap at (H,H)) for `results.omega_condition_numbers` and consolidates cells with cond > 1e12 into ONE fit-level warning (count + max cond) instead of the legacy per-cell pseudoinverse warnings; the covariate path intentionally computes NO per-unit condition numbers — they would cost exactly the per-unit SVDs the v3.7 rewrite removes, near-singularity there is structural (~always true under PT-All, so a warning would be always-on noise rather than signal), and the ridge handles it by design (the legacy covariate path likewise had no per-unit diagnostics). The scalability warning threshold moved from n > 5000 (legacy per-cell warning) to n > 50,000 (one fit-level warning; the kernel stage is still intrinsically O(n^2) but with a ~100x lower constant and tile-bounded memory). - **Note:** Bootstrap aggregation uses fixed cohort-size weights for overall/event-study reaggregation, matching the CallawaySantAnna bootstrap pattern (staggered_bootstrap.py:281 computes `bootstrap_overall = bootstrap_atts_gt[:, post_indices] @ weights`; L297 uses the same fixed-weight pattern for event study). The analytical path includes a WIF correction; fixed-weight bootstrap captures the same sampling variability through per-cell EIF perturbation without re-estimating aggregation weights, consistent with both the library's CS implementation and the R `did` package. - **Overall ATT convention**: The library's `overall_att` uses cohort-size-weighted averaging of post-treatment (g,t) cells, matching the CallawaySantAnna simple aggregation. This differs from the paper's ES_avg (Eq 2.3), which uniformly averages over event-time horizons. ES_avg can be computed from event study output as `mean(event_study_effects[e]["effect"] for e >= 0)` diff --git a/docs/performance-plan.md b/docs/performance-plan.md index 202655e4c..578189baa 100644 --- a/docs/performance-plan.md +++ b/docs/performance-plan.md @@ -231,6 +231,58 @@ covariate counts (IRLS `lstsq`, rank-guard QR — Phase 3 candidate). --- +## EfficientDiD analytical path: omega_ridge + fused tiled GEMM (v3.7, 2026-07) + +Resolves the "analytical stage pathological scaling" TODO filed during the Phase 2 +bootstrap work. Instrumented stage split at n=500/1k/2k units (20 periods, 5 cohorts, +5 covariates, 95 cells, H=60 pairs/cell; `.bench-local/edid_analytical_*.json`): + +| stage | n=500 | n=1k | n=2k | mechanism | +|---|---|---|---|---| +| `compute_omega_star_conditional` | 23.9s | 71.3s | 289.8s | O(n^2 H^2) Python pair loop, ~4 (n x n_group) broadcast temporaries per pair | +| `compute_per_unit_weights` | 11.9s | 23.9s | 47.7s | per-unit SVD cond + SVD pinv Python loop | +| everything else | ~0.4s | ~0.5s | ~1.8s | sieves, gen_out, EIF | + +The spike also uncovered the numerical-stability story documented in REGISTRY.md +(Omega* ridge Note): Omega* is numerically singular by construction (cond 1e17-1e22, +100% of units), so the legacy pseudoinverse redrew per-cell values at ~1e-2 rel under +ANY floating-point change (1-ulp input perturbation on the shipped code: 1.2e-4). +An eigh-based pinv (2x faster) was measured and REJECTED — O(1) per-unit weight +changes through the same cutoff cliff. The ridge (default 1e-6, trace-scaled) is both +the stability fix and what unlocks the batched solve. + +Rewrite (one PR): `_kcov_batch` GEMM identity `KCov = W@(A0*B0) - (W@A0)*(W@B0)` +(rows of the kernel matrix sum to 1, survey-weighted included; validated 28-29x per +cell at 2e-15 rel on captured inputs), (t_pre_j, t_pre_k) dedup (1830 pairs -> ~190 +GEMM columns), fused unit-tiled pass 2 (per-group kernel matrices computed once per +tile, reused across ALL 95 cells; omega assembled per tile; batched +`np.linalg.solve` ridge weights; EIF centered after all tiles; +`_TARGET_OMEGA_TILE_BYTES` = 256MB, call-time resolved), nocov Gram twin. + +Measured arms (3-rep medians; BEFORE = pristine main ced49e89 via baseline-worktree +PYTHONPATH; `.bench-local/edid_arms_results.jsonl`): + +| scenario | BEFORE | AFTER | speedup | maxrss B->A | +|---|---|---|---|---| +| conditional n=500 | 35.9s | 1.83s | 19.7x | 0.18 -> 0.25 GB | +| conditional n=1k | 95.4s | 3.48s | 27.5x | 0.25 -> 0.37 GB | +| conditional n=2k | 337.3s | 7.81s | 43.2x | 0.37 -> 0.59 GB | +| conditional n=10k | ~2.3h (extrapolated O(n^2)) | 57.7s | ~146x | - -> 0.71 GB | +| conditional n=100k | memory-killed | 85.6 min (1 rep) | completes | - -> 4.0 GB | +| nocov n=2k | 1.22s | 0.47s | 2.6x | 0.15 -> 0.15 GB | +| survey n=1k | 95.3s | 3.34s | 28.6x | 0.25 -> 0.37 GB | + +Estimate deltas AFTER vs BEFORE (the documented one-time ridge redraw, shrinking +with n as the indeterminacy band tightens): overall ATT 1.6e-2 / 4.0e-3 / 1.1e-3 rel +at n=500/1k/2k; worst post-treatment cell <= 0.58 of its own SE; nocov path ~1e-7. +1-ulp per-cell stability after: <= 3e-9 rel (before: 1.2e-4). + +Remaining lever (TODO row): cross-cell kernel-table hoisting (Term 5 tables are +(g,t)-independent; Term 2 depends only on t) — ~10x on the intrinsic per-cell GEMM +factor at 100k-unit scale; not needed below ~50k units. + +--- + ## Memory scaling at the millions-of-units tail (v3.x, June 2026) At the scale where the dense working arrays - not compute - are the binding constraint diff --git a/tests/test_efficient_did.py b/tests/test_efficient_did.py index 6ac513462..2458aa93c 100644 --- a/tests/test_efficient_did.py +++ b/tests/test_efficient_did.py @@ -449,7 +449,13 @@ def test_weights_uniform_under_iid(self): assert w.std() > 0 def test_condition_number_warning(self): - """Near-singular Omega* should trigger a warning.""" + """Near-singular Omega* should trigger a warning (legacy path). + + Re-scoped to omega_ridge=0 with the v3.7 ridge default: this test + exercises the legacy inv/pinv path's per-cell warning contract, which + the default ridge path intentionally replaces with one aggregate + fit-level warning (see TestOmegaRidge). + """ # Use a perfectly collinear DGP to produce near-singular Omega* n_units, n_periods = 100, 5 units = np.repeat(np.arange(n_units), n_periods) @@ -469,7 +475,7 @@ def test_condition_number_warning(self): ) with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") - EfficientDiD().fit(df, "y", "unit", "time", "first_treat") + EfficientDiD(omega_ridge=0.0).fit(df, "y", "unit", "time", "first_treat") # Should get a warning about condition number or zero matrix warning_msgs = [str(x.message) for x in w] assert any( @@ -481,6 +487,343 @@ def test_condition_number_warning(self): ), f"Expected condition/zero warning, got: {warning_msgs}" +class TestOmegaRidge: + """omega_ridge parameter surface, warning contract, and engagement lock. + + The v3.7 default (OMEGA_RIDGE_DEFAULT) ridge-regularizes the Omega* + inversion behind the efficient weights; omega_ridge=0 restores the exact + legacy inv/pinv code path. See REGISTRY.md (EfficientDiD, Omega* ridge + regularization note). + """ + + @staticmethod + def _duplicated_period_panel(n_units=120, seed=11): + """Panel where period 3 duplicates period 2 exactly. + + The never-treated moments for t_pre=2 and t_pre=3 become identical, + so every overidentified cell's Omega* is exactly singular + (cond > 1e12) WITHOUT being the all-zero matrix. + """ + rng = np.random.default_rng(seed) + n_periods = 6 + ft = np.full(n_units, np.inf) + ft[: n_units // 2] = 5 + y_wide = rng.normal(0, 1.0, (n_units, n_periods)) + y_wide[:, 2] = y_wide[:, 1] # period 3 == period 2 + treated = (ft[:, None] < np.inf) & (np.arange(1, n_periods + 1)[None, :] >= ft[:, None]) + y_wide = y_wide + 2.0 * treated + return pd.DataFrame( + { + "unit": np.repeat(np.arange(n_units), n_periods), + "time": np.tile(np.arange(1, n_periods + 1), n_units), + "first_treat": np.repeat(ft, n_periods), + "y": y_wide.ravel(), + } + ) + + def test_param_surface(self): + from diff_diff.efficient_did_covariates import OMEGA_RIDGE_DEFAULT + + est = EfficientDiD() + assert est.get_params()["omega_ridge"] == OMEGA_RIDGE_DEFAULT + est.set_params(omega_ridge=0.0) + assert est.omega_ridge == 0.0 + est.set_params(omega_ridge=1e-4) + assert est.get_params()["omega_ridge"] == 1e-4 + + @pytest.mark.parametrize("bad", [-1e-6, np.nan, np.inf, -np.inf]) + def test_validation_rejects_bad_values(self, bad): + with pytest.raises(ValueError, match="omega_ridge"): + EfficientDiD(omega_ridge=bad) + + def test_set_params_transactional(self): + from diff_diff.efficient_did_covariates import OMEGA_RIDGE_DEFAULT + + est = EfficientDiD() + with pytest.raises(ValueError): + est.set_params(omega_ridge=-1.0) + assert est.omega_ridge == OMEGA_RIDGE_DEFAULT + + def test_results_echo(self): + df = _make_staggered_panel() + res = EfficientDiD(omega_ridge=1e-5).fit(df, "y", "unit", "time", "first_treat") + assert res.omega_ridge == 1e-5 + + def test_aggregate_warning_payload_and_legacy_per_cell(self): + """Default ridge: ONE fit-level warning whose cell count matches the + number of genuinely ill-conditioned cells; omega_ridge=0: the legacy + per-cell pseudoinverse warnings for the same cells.""" + import re + + df = self._duplicated_period_panel() + + with warnings.catch_warnings(record=True) as w_ridge: + warnings.simplefilter("always") + EfficientDiD().fit(df, "y", "unit", "time", "first_treat") + agg = [str(x.message) for x in w_ridge if "regularization handled" in str(x.message)] + assert len(agg) == 1, f"expected exactly one aggregate warning, got {agg}" + m = re.search(r"in (\d+) of (\d+) \(g, t\) cells", agg[0]) + assert m is not None, agg[0] + n_ill, n_cells = int(m.group(1)), int(m.group(2)) + assert 1 <= n_ill <= n_cells + + with warnings.catch_warnings(record=True) as w_legacy: + warnings.simplefilter("always") + EfficientDiD(omega_ridge=0.0).fit(df, "y", "unit", "time", "first_treat") + legacy_msgs = [str(x.message) for x in w_legacy] + per_cell = [m_ for m_ in legacy_msgs if "using pseudoinverse for weights" in m_] + # Same pathology surfaces per-cell on the legacy path. The legacy pair + # set additionally contains the degenerate (g'=g, t_pre=t) self-pair + # for pre-treatment cells (dropped on the ridge path), so legacy may + # warn on MORE cells - never fewer. + assert len(per_cell) >= n_ill + assert not any("regularization handled" in m_ for m_ in legacy_msgs) + + def test_ridge_engaged_and_overall_stable(self): + """Default vs omega_ridge=0 on an overidentified panel: overall ATT + agrees tightly while at least one per-cell effect differs (proves the + ridge path is actually engaged).""" + df = _make_staggered_panel(n_per_group=80, n_control=100) + res_r = EfficientDiD().fit(df, "y", "unit", "time", "first_treat") + res_l = EfficientDiD(omega_ridge=0.0).fit(df, "y", "unit", "time", "first_treat") + assert res_r.overall_att == pytest.approx(res_l.overall_att, rel=1e-4) + diffs = [ + abs(res_r.group_time_effects[k]["effect"] - res_l.group_time_effects[k]["effect"]) + for k in res_r.group_time_effects + ] + assert max(diffs) > 0.0 + + def test_degenerate_self_pair_dropped_only_on_ridge_path(self): + """Pre-treatment cells lose the identically-zero (g'=g, t_pre=t) + self-pair under ridge (stored nocov weights shrink by exactly one), + while post-treatment cells keep the same pair count.""" + df = _make_staggered_panel() + res_r = EfficientDiD().fit(df, "y", "unit", "time", "first_treat") + res_l = EfficientDiD(omega_ridge=0.0).fit(df, "y", "unit", "time", "first_treat") + assert res_r.efficient_weights is not None + assert res_l.efficient_weights is not None + period_1 = min(res_r.time_periods) + for gt, w_l in res_l.efficient_weights.items(): + g, t = gt + w_r = res_r.efficient_weights[gt] + if t < g and t != period_1: + assert len(w_r) == len(w_l) - 1, f"cell {gt}" + else: + assert len(w_r) == len(w_l), f"cell {gt}" + + def test_pretreatment_placebos_remain_data_driven(self): + """Under default ridge, pre-treatment placebo cells stay data-driven + (nonzero noise), NOT deterministically zero - the degenerate-pair + drop preserves the pre-trend diagnostic.""" + df = _make_staggered_panel(rho=0.3, seed=99) + res = EfficientDiD().fit(df, "y", "unit", "time", "first_treat") + pre = [ + abs(v["effect"]) + for (g, t), v in res.group_time_effects.items() + if t < g and np.isfinite(v["effect"]) + ] + assert pre, "expected pre-treatment cells" + assert max(pre) > 1e-8 + + +class TestFusedConditionalPath: + """The v3.7 fused unit-tiled GEMM conditional path (default ridge): + semantic equivalence to the dense construction, tile invariance, and the + 1-ulp stability contract that motivated the ridge.""" + + @staticmethod + def _cov_df(n_units=200, n_periods=9, seed=42): + return _make_covariate_panel(n_units=n_units, n_periods=n_periods, seed=seed) + + @staticmethod + def _surfaces(res): + gt = {k: (v["effect"], v["se"]) for k, v in res.group_time_effects.items()} + es = ( + {k: (v["effect"], v["se"]) for k, v in res.event_study_effects.items()} + if res.event_study_effects + else {} + ) + return gt, es, (res.overall_att, res.overall_se) + + @classmethod + def _assert_close_surfaces(cls, r1, r2, rtol, atol=1e-12): + gt1, es1, ov1 = cls._surfaces(r1) + gt2, es2, ov2 = cls._surfaces(r2) + np.testing.assert_allclose(ov1, ov2, rtol=rtol, atol=atol) + for k in gt2: + np.testing.assert_allclose(gt1[k], gt2[k], rtol=rtol, atol=atol, err_msg=str(k)) + for k in es2: + np.testing.assert_allclose(es1[k], es2[k], rtol=rtol, atol=atol, err_msg=str(k)) + + def test_fused_matches_dense_reference(self): + """Fused tiled GEMM cells match a dense reference built from the + legacy compute_omega_star_conditional + the same ridge solve. The + residual is the omega GEMM's ~1e-15 reassociation drift amplified by + the ridge's bounded 1/lambda sensitivity (~1e6) - i.e. <= ~1e-8, the + designed stability, vs ~1e-2 under the legacy pseudoinverse.""" + import diff_diff.efficient_did as ed_mod + from diff_diff.efficient_did_covariates import ( + _ridge_solve_weights, + compute_generated_outcomes_cov, + compute_omega_star_conditional, + ) + + def dense_reference( + cell_specs, + outcome_wide, + covariate_matrix, + cohort_masks, + never_treated_mask, + period_to_col, + cohort_fractions, + m_hat_cache, + r_hat_cache, + s_hat_cache, + bandwidth, + omega_ridge, + unit_weights=None, + never_treated_val=np.inf, + tile_bytes=None, + ): + out = {} + for spec in cell_specs: + g, t, pairs = spec["g"], spec["t"], spec["pairs"] + gen_out = compute_generated_outcomes_cov( + target_g=g, + target_t=t, + valid_pairs=pairs, + outcome_wide=outcome_wide, + cohort_masks=cohort_masks, + never_treated_mask=never_treated_mask, + period_to_col=period_to_col, + period_1_col=spec["y1_col"], + cohort_fractions=cohort_fractions, + m_hat_cache=m_hat_cache, + r_hat_cache=r_hat_cache, + ) + if len(pairs) == 1: + scores = gen_out[:, 0] + else: + omega = compute_omega_star_conditional( + target_g=g, + target_t=t, + valid_pairs=pairs, + outcome_wide=outcome_wide, + cohort_masks=cohort_masks, + never_treated_mask=never_treated_mask, + period_to_col=period_to_col, + period_1_col=spec["y1_col"], + cohort_fractions=cohort_fractions, + covariate_matrix=covariate_matrix, + s_hat_cache=s_hat_cache, + bandwidth=bandwidth, + unit_weights=unit_weights, + ) + w = _ridge_solve_weights(omega, omega_ridge) + scores = np.sum(w * gen_out, axis=1) + att = ( + float(np.average(scores, weights=unit_weights)) + if unit_weights is not None + else float(np.mean(scores)) + ) + out[(g, t)] = (att, scores - att) + return out + + df = self._cov_df() + fit_kwargs = dict(covariates=["x1", "x2"], aggregate="all") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r_fused = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", **fit_kwargs) + orig = ed_mod.compute_conditional_cells_tiled + ed_mod.compute_conditional_cells_tiled = dense_reference + try: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r_ref = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", **fit_kwargs) + finally: + ed_mod.compute_conditional_cells_tiled = orig + self._assert_close_surfaces(r_fused, r_ref, rtol=1e-6, atol=1e-8) + + def test_tile_forced_twin(self, monkeypatch): + """Forcing one-unit tiles must reproduce the single-tile fit + (rel 1e-10) - and must actually execute multi-tile.""" + import diff_diff.efficient_did_covariates as cov_mod + + df = self._cov_df(n_units=150) + fit_kwargs = dict(covariates=["x1", "x2"], aggregate="all") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r_single = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", **fit_kwargs) + + calls = {"n": 0} + orig_kwm = cov_mod._kernel_weights_matrix + + def counting_kwm(*args, **kwargs): + calls["n"] += 1 + return orig_kwm(*args, **kwargs) + + monkeypatch.setattr(cov_mod, "_TARGET_OMEGA_TILE_BYTES", 1) + monkeypatch.setattr(cov_mod, "_kernel_weights_matrix", counting_kwm) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r_tiled = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", **fit_kwargs) + # one-unit tiles -> at least n_units kernel-matrix builds (vs ~4 for + # a single tile), proving the multi-tile path executed + assert calls["n"] >= 150, calls["n"] + self._assert_close_surfaces(r_tiled, r_single, rtol=1e-10, atol=1e-12) + + def test_tile_forced_twin_survey_weights(self, monkeypatch): + """Tile invariance under survey weights (weighted kernels, weighted + ATT averaging).""" + import diff_diff.efficient_did_covariates as cov_mod + + df = self._cov_df(n_units=150) + rng = np.random.default_rng(5) + n_units = df["unit"].nunique() + pw = np.exp(rng.normal(0, 0.4, n_units)) + df = df.merge(pd.DataFrame({"unit": sorted(df["unit"].unique()), "pw": pw}), on="unit") + fit_kwargs = dict( + covariates=["x1", "x2"], aggregate="all", survey_design=SurveyDesign(weights="pw") + ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r_single = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", **fit_kwargs) + monkeypatch.setattr(cov_mod, "_TARGET_OMEGA_TILE_BYTES", 1) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r_tiled = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", **fit_kwargs) + self._assert_close_surfaces(r_tiled, r_single, rtol=1e-10, atol=1e-12) + + def test_one_ulp_stability_conditional(self): + """THE stability contract: a 1-ulp outcome perturbation moves + per-cell and event-study ATTs and SEs by <= 1e-6 relative under the + default ridge (the legacy pseudoinverse path moves ~1e-4).""" + df = self._cov_df() + fit_kwargs = dict(covariates=["x1", "x2"], aggregate="all") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r_base = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", **fit_kwargs) + df_ulp = df.copy() + df_ulp["y"] = np.nextafter(df_ulp["y"].to_numpy(), np.inf) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r_ulp = EfficientDiD().fit(df_ulp, "y", "unit", "time", "first_treat", **fit_kwargs) + self._assert_close_surfaces(r_ulp, r_base, rtol=1e-6, atol=1e-9) + + def test_one_ulp_stability_nocov(self): + """Same stability contract on the no-covariates PT-All path.""" + df = _make_staggered_panel(n_per_group=80, n_control=100) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r_base = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", aggregate="all") + df_ulp = df.copy() + df_ulp["y"] = np.nextafter(df_ulp["y"].to_numpy(), np.inf) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r_ulp = EfficientDiD().fit(df_ulp, "y", "unit", "time", "first_treat", aggregate="all") + self._assert_close_surfaces(r_ulp, r_base, rtol=1e-6, atol=1e-9) + + class TestValidTriples: """Test enumerate_valid_triples with hand-worked examples.""" diff --git a/tests/test_methodology_efficient_did.py b/tests/test_methodology_efficient_did.py index d3650b858..8c6ee9b0b 100644 --- a/tests/test_methodology_efficient_did.py +++ b/tests/test_methodology_efficient_did.py @@ -25,7 +25,11 @@ - ``TestEfficientWeights`` — inverse-covariance optimal weights ``w = 1' Omega*^-1 / (1' Omega*^-1 1)`` (Eq 3.5 single-date / Eq 3.13 staggered), - the min-variance property, and the singular-``Omega*`` pseudoinverse path. + the min-variance property, and the legacy singular-``Omega*`` pseudoinverse + path (``omega_ridge=0``; the helpers default to it). +- ``TestOmegaRidgeWeights`` / ``TestNoCovOmegaGram`` / ``TestKernelCovBatch`` — + the v3.7 ridge-regularized inversion and the GEMM/Gram twins of the omega + constructions (see the Omega* ridge Note in REGISTRY.md). - ``TestGeneratedOutcomeNoCov`` — the no-covariate generated outcome (Eq 3.9): the ``g'=g`` same-cohort pair telescopes to the per-baseline DiD ``E[Y_t-Y_tpre|g] - E[Y_t-Y_tpre|inf]`` (Eq 3.3); the ``g'=inf`` pairs telescope @@ -200,6 +204,275 @@ def test_singular_omega_uses_pseudoinverse(self): assert abs(w.sum() - 1.0) < 1e-8 +class TestOmegaRidgeWeights: + """Ridge-regularized weights ``(Omega + lam*max(trace/H,0)*I)^-1 1`` + (normalized) - the v3.7 stabilization of the Eq 3.5/3.13 inversion on + numerically singular Omega*. omega_ridge=0 preserves the legacy path + exactly; see REGISTRY.md.""" + + def test_global_weights_hand_computed(self): + omega = np.array([[2.0, 0.5, 0.0], [0.5, 1.0, 0.2], [0.0, 0.2, 3.0]]) + lam = 1e-3 + w, used_pinv, _ = compute_efficient_weights(omega, omega_ridge=lam) + scale = np.trace(omega) / 3.0 + num = np.linalg.solve(omega + lam * scale * np.eye(3), np.ones(3)) + expected = num / num.sum() + assert not used_pinv + np.testing.assert_allclose(w, expected, atol=1e-12) + assert abs(w.sum() - 1.0) < 1e-12 + + def test_per_unit_weights_hand_computed(self): + from diff_diff.efficient_did_covariates import compute_per_unit_weights + + base = np.array([[2.0, 0.5, 0.0], [0.5, 1.0, 0.2], [0.0, 0.2, 3.0]]) + stack = np.stack([base, 2.0 * base]) + lam = 1e-3 + w = compute_per_unit_weights(stack, omega_ridge=lam) + for i in range(2): + scale = np.trace(stack[i]) / 3.0 + num = np.linalg.solve(stack[i] + lam * scale * np.eye(3), np.ones(3)) + np.testing.assert_allclose(w[i], num / num.sum(), atol=1e-12) + # trace-relative ridge is scale-equivariant: same weights for c*Omega + np.testing.assert_allclose(w[0], w[1], atol=1e-12) + + def test_zero_matrix_and_indefinite_zero_trace_fall_back_to_uniform(self): + from diff_diff.efficient_did_covariates import compute_per_unit_weights + + zero = np.zeros((2, 2)) + # Indefinite with zero trace: scale clamps to 0 (no negative ridge) + # and 1'A^-1 1 = 0 exactly -> den guard -> uniform. + indefinite = np.array([[1.0, 0.0], [0.0, -1.0]]) + spd = np.array([[1.0, 0.2], [0.2, 2.0]]) + stack = np.stack([zero, indefinite, spd]) + w = compute_per_unit_weights(stack, omega_ridge=1e-6) + np.testing.assert_allclose(w[0], [0.5, 0.5], atol=1e-15) + np.testing.assert_allclose(w[1], [0.5, 0.5], atol=1e-15) + assert abs(w[2].sum() - 1.0) < 1e-12 + assert not np.allclose(w[2], [0.5, 0.5]) + + def test_dose_response_larger_lambda_pushes_toward_uniform(self): + omega = np.array([[1.0, 0.0, 0.0], [0.0, 5.0, 0.0], [0.0, 0.0, 25.0]]) + uniform = np.ones(3) / 3.0 + dist = [] + for lam in (1e-8, 1e-2, 1e2, 1e6): + w, _, _ = compute_efficient_weights(omega, omega_ridge=lam) + dist.append(float(np.linalg.norm(w - uniform))) + assert dist == sorted(dist, reverse=True), dist + assert dist[-1] < 1e-4 # huge ridge ~ uniform + + def test_lam_zero_is_exact_legacy(self): + omega = np.array([[2.0, 0.5], [0.5, 1.0]]) + w_legacy, p1, c1 = compute_efficient_weights(omega) + w_zero, p2, c2 = compute_efficient_weights(omega, omega_ridge=0.0) + assert np.array_equal(w_legacy, w_zero) + assert (p1, c1) == (p2, c2) + + def test_ridge_weights_remain_minimum_variance_up_to_lambda(self): + """For small lam on a well-conditioned Omega, ridge weights stay + within O(lam) of the exact minimum-variance solution.""" + omega = np.array([[1.0, 0.3], [0.3, 2.0]]) + w_exact, _, _ = compute_efficient_weights(omega) + w_ridge, _, _ = compute_efficient_weights(omega, omega_ridge=1e-8) + np.testing.assert_allclose(w_ridge, w_exact, atol=1e-6) + + +class TestNoCovOmegaGram: + """Gram/GEMM twin of compute_omega_star_nocov (v3.7 ridge path): exact + five-term + guard-semantics parity with the _sample_cov loop, including + the n<2 group, sum(w)<=1 group, and zero-fraction-cohort guards.""" + + @staticmethod + def _fixture( + n_per_cohort=(30, 25, 1), + n_control=20, + n_periods=6, + seed=8, + weighted=False, + control_weight_total=None, + ): + """Wide-outcome fixture with cohorts at g=3,4 and a 1-UNIT cohort at + g=5 (exercises the n<2 covariance guard).""" + from diff_diff.efficient_did_weights import enumerate_valid_triples + + rng = np.random.default_rng(seed) + groups = [3.0, 4.0, 5.0] + sizes = list(n_per_cohort) + [n_control] + n_units = sum(sizes) + wide = rng.normal(0, 1.0, (n_units, n_periods)) + cohorts = np.concatenate( + [np.full(s, g) for s, g in zip(n_per_cohort, groups)] + [np.full(n_control, np.inf)] + ) + cohort_masks = {g: cohorts == g for g in groups} + never_mask = np.isinf(cohorts) + cohort_masks[np.inf] = never_mask + time_periods = [float(p) for p in range(1, n_periods + 1)] + period_to_col = {p: i for i, p in enumerate(time_periods)} + weights = None + if weighted: + weights = np.exp(rng.normal(0, 0.5, n_units)) + if control_weight_total is not None: + w_ctrl = weights[never_mask] + weights[never_mask] = w_ctrl * (control_weight_total / w_ctrl.sum()) + if weights is not None: + total = weights.sum() + fractions = {g: float(weights[cohort_masks[g]].sum()) / total for g in groups} + fractions[np.inf] = float(weights[never_mask].sum()) / total + else: + fractions = {g: float(cohort_masks[g].sum()) / n_units for g in groups} + fractions[np.inf] = float(never_mask.sum()) / n_units + pairs = enumerate_valid_triples(3.0, groups, time_periods, 1.0, "all") + return dict( + wide=wide, + cohort_masks=cohort_masks, + never_mask=never_mask, + period_to_col=period_to_col, + fractions=fractions, + pairs=pairs, + weights=weights, + ) + + def _compare(self, fx, target_g=3.0, target_t=4.0, rtol=1e-12, atol=1e-13): + from diff_diff.efficient_did_weights import ( + _omega_star_nocov_gram, + compute_omega_star_nocov, + ) + + kwargs = dict( + target_g=target_g, + target_t=target_t, + valid_pairs=fx["pairs"], + outcome_wide=fx["wide"], + cohort_masks=fx["cohort_masks"], + never_treated_mask=fx["never_mask"], + period_to_col=fx["period_to_col"], + period_1_col=0, + cohort_fractions=fx["fractions"], + unit_weights=fx["weights"], + ) + legacy = compute_omega_star_nocov(**kwargs) + gram = _omega_star_nocov_gram(**kwargs) + assert gram.shape == legacy.shape + np.testing.assert_allclose(gram, legacy, rtol=rtol, atol=atol) + + @pytest.mark.parametrize("weighted", [False, True]) + def test_matches_loop_incl_one_unit_cohort(self, weighted): + # includes a 1-unit cohort (g=5) in the pair set -> n<2 guard parity + fx = self._fixture(weighted=weighted) + assert any(gp == 5.0 for gp, _ in fx["pairs"]) + self._compare(fx) + + def test_sub_unit_total_weight_group_guard(self): + # never-treated total survey weight <= 1 -> weighted Bessel guard + # (sum_w - 1 <= 0) returns 0.0 in _sample_cov; the batch must agree + fx = self._fixture(weighted=True, control_weight_total=0.9) + self._compare(fx) + + def test_empty_pairs_shape_contract(self): + from diff_diff.efficient_did_weights import _omega_star_nocov_gram + + fx = self._fixture() + out = _omega_star_nocov_gram( + target_g=3.0, + target_t=4.0, + valid_pairs=[], + outcome_wide=fx["wide"], + cohort_masks=fx["cohort_masks"], + never_treated_mask=fx["never_mask"], + period_to_col=fx["period_to_col"], + period_1_col=0, + cohort_fractions=fx["fractions"], + ) + assert out.shape == (0, 0) + + def test_fit_level_twin(self): + """Default-ridge nocov fit routed through the Gram omega matches the + same fit forced onto the legacy loop omega at rel 1e-10.""" + import diff_diff.efficient_did as ed_mod + from diff_diff.efficient_did_weights import compute_omega_star_nocov + + rng = np.random.default_rng(21) + n_units, n_periods = 240, 7 + ft = np.full(n_units, np.inf) + ft[:80] = 3 + ft[80:150] = 5 + y = rng.normal(0, 1, (n_units, n_periods)) + y += 2.0 * ((ft[:, None] < np.inf) & (np.arange(1, n_periods + 1)[None] >= ft[:, None])) + df = pd.DataFrame( + { + "unit": np.repeat(np.arange(n_units), n_periods), + "time": np.tile(np.arange(1, n_periods + 1), n_units), + "first_treat": np.repeat(ft, n_periods), + "y": y.ravel(), + } + ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r_gram = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", aggregate="all") + orig = ed_mod._omega_star_nocov_gram + ed_mod._omega_star_nocov_gram = compute_omega_star_nocov + try: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r_loop = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", aggregate="all") + finally: + ed_mod._omega_star_nocov_gram = orig + assert r_gram.overall_att == pytest.approx(r_loop.overall_att, rel=1e-10) + for k in r_loop.group_time_effects: + np.testing.assert_allclose( + ( + r_gram.group_time_effects[k]["effect"], + r_gram.group_time_effects[k]["se"], + ), + ( + r_loop.group_time_effects[k]["effect"], + r_loop.group_time_effects[k]["se"], + ), + rtol=1e-8, + atol=1e-10, + err_msg=str(k), + ) + + +class TestKernelCovBatch: + """GEMM-form kernel covariance (v3.7 fused path) vs the dense + _kernel_weighted_cov loop: KCov_i(A,B) = (W@(A0*B0))_i - (W@A0)_i(W@B0)_i + for row-normalized W - exact identity, verified with and without survey + weights and on a row-tile of W.""" + + @staticmethod + def _setup(seed, weighted): + from diff_diff.efficient_did_covariates import _kernel_weights_matrix + + rng = np.random.default_rng(seed) + n_all, n_group, m = 60, 25, 7 + x_all = rng.normal(size=(n_all, 3)) + x_group = rng.normal(size=(n_group, 3)) + gw = np.exp(rng.normal(0, 0.5, n_group)) if weighted else None + W = _kernel_weights_matrix(x_all, x_group, bandwidth=0.8, group_weights=gw) + # offset columns exercise the global-centering cancellation guard + A = rng.normal(size=(n_group, m)) + 1e6 + B = rng.normal(size=(n_group, m)) + 1e6 + return W, A, B + + @pytest.mark.parametrize("weighted", [False, True]) + def test_matches_dense_loop(self, weighted): + from diff_diff.efficient_did_covariates import _kcov_batch, _kernel_weighted_cov + + W, A, B = self._setup(3, weighted) + got = _kcov_batch(W, A, B) + for j in range(A.shape[1]): + expected = _kernel_weighted_cov(A[:, j], B[:, j], W) + np.testing.assert_allclose(got[:, j], expected, rtol=1e-12, atol=1e-9) + + def test_row_tile_equals_full(self): + from diff_diff.efficient_did_covariates import _kcov_batch + + W, A, B = self._setup(4, True) + full = _kcov_batch(W, A, B) + tiled = np.vstack([_kcov_batch(W[i : i + 7], A, B) for i in range(0, W.shape[0], 7)]) + np.testing.assert_allclose(tiled, full, rtol=1e-13, atol=1e-12) + + # ============================================================================= # Eq 3.9 — no-covariate generated outcome (telescoping) # ============================================================================= @@ -285,7 +558,11 @@ def test_efficient_att_equals_weighted_generated_outcomes(self): ) with warnings.catch_warnings(): warnings.simplefilter("ignore") - w, _, _ = compute_efficient_weights(omega) + # Pass the shipped default ridge so this closed-form lock keeps + # guarding the DEFAULT fit path (the fit above uses it too). + from diff_diff.efficient_did_covariates import OMEGA_RIDGE_DEFAULT + + w, _, _ = compute_efficient_weights(omega, omega_ridge=OMEGA_RIDGE_DEFAULT) y_hat = compute_generated_outcomes_nocov( g, t,