diff --git a/CHANGELOG.md b/CHANGELOG.md index 40ead62d..d00f6239 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -222,6 +222,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 the same dispatch globals the fit consumes (`bootstrap_utils` / `linalg`), with a coherence assert between the two (defense in depth). +### Fixed +- **Rust backend: saturated full-rank designs (`n == k`) no longer leak `Inf` + variance-covariance silently.** `solve_ols`'s Rust HC1/CR1 path computed the + `n/(n-k)` (clustered `(n-1)/(n-k)`) adjustment with zero residual degrees of + freedom, producing all-`Inf` vcov that bypassed the dispatcher's Python-fallback + check (which keyed on NaN only) and reached users without a warning. The Rust + kernel now returns the documented non-finite-inference contract — an all-NaN + vcov, matching the canonical numpy saturated guard — with the `G >= 2` cluster + error keeping precedence, and the dispatcher's fallback check now rejects any + non-finite vcov (NaN or Inf). Point estimates and residuals are unchanged; + only the undefined-inference sentinel changes (`Inf` → `NaN` + warning). + Found while building the opt-in Cholesky fast path below, which honors the + same contract. + ### Added - **`SyntheticControl` conformal extensions: one-sided alternatives + covariates in the proxy (Chernozhukov-Wüthrich-Zhu 2021).** `conformal_test` / `conformal_confidence_intervals` @@ -236,6 +250,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 T*-block structure. Locked by a hand-rolled signed-statistic permutation oracle, directional rejection/half-line/composition tests, a proxy-weight mechanism check, and a permutation-floor invariant; all 31 pre-existing conformal tests pass unchanged. +- **Opt-in `DIFF_DIFF_SOLVE_OLS_FASTPATH` normal-equations Cholesky fast path for + `solve_ols` (both backends).** Set the env var to `1` (a positive integer; boolean + spellings like `true`/`on` are silently ignored per the resolver convention, mirroring + `DIFF_DIFF_DEMEAN_CHUNK_COLS`) to route certified-well-conditioned full-rank OLS solves + through an equilibrated normal-equations Cholesky instead of the default SVD/gelsd: + the Python twin reuses the stage-0 rank-certification Gram and gates on LAPACK `dpocon` + reciprocal condition > 1e-6; the Rust side is a new self-certifying faer `Llt` kernel + (`solve_ols_chol`, faer-only heavy ops — identical behavior across the + accelerate/openblas/no-BLAS wheels) whose exact 1-norm rcond inverse doubles as the + sandwich bread. Default OFF = byte-identical legacy behavior on both backends (locked + by dispatch spies, exact-equality tests, and a pristine-base-tree identity gate; sole + exception: the saturated n == k vcov contract fix above); any + certification decline falls back verbatim to the SVD path, so a knob-on decline equals + knob-off exactly. Within the opt-in path, parity vs the default is tol-bounded (fitted + ~1e-8 abs / SE ~1e-6 rel; certified forward-error budget ~eps·cond ≤ 2e-10). Measured + (M4 Max, medians): SunAbraham 1.13→0.63 s and 16.8→9.9 s (1.7-1.8x), CallawaySantAnna + dr 40-covariate 4.40→3.57 s (1.23x), `skip_rank_check` micro-solve 1.56x; rust-side + allocator high-water on a 2.4M×130 clustered solve 7.27→2.66 GB (no thin-SVD U + transient). See docs/performance-plan.md § "Opt-in solve_ols normal-equations Cholesky + fast path". - **Opt-in `df_convention="cluster"` inference-df knob (DiD / TWFE / MultiPeriodDiD + `LinearRegression`).** Clustered analytical fits historically compute t-statistics, p-values, and CIs at the fitted **residual df** (`n − K_full`), while `fixest`/Stata use diff --git a/TODO.md b/TODO.md index 831b9792..0d5ab6dd 100644 --- a/TODO.md +++ b/TODO.md @@ -42,7 +42,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| | `EfficientDiD` conditional path: the largest remaining O(n) stage is the sieve/nuisance construction outside the tiled pass (~9s at 10k). (The `_ridge_solve_weights` Python-prep shave landed 2026-07-07 — the `omega_stack[rest]` fancy-index copy and tail scatter are skipped when no row is zero-masked, byte-identical outputs; the `zero_mask` abs scan itself remains, needed for correctness.) | `efficient_did_covariates.py` | CS-scaling | Mid | Low | -| Rust `solve_ols` always runs a full thin SVD (equilibrated, gelsd-parity); at 40 covariates it is the top per-cell solver item on a CS dr fit (1.04s of 3.1s at 2M rows, 95 calls). A Cholesky/QR fast path with SVD fallback — mirroring the Phase 3 Python-side pattern (certify well-conditioned, fall back verbatim) — is the natural lever, but `solve_ols` is the universal OLS entry point (every estimator), so the blast radius needs the full backend-parity treatment (`TestSolveOLSSkipRankCheckParity` posture: fitted-values parity, not beta). | `rust/src/linalg.rs::solve_ols` | CS-scaling | Heavy | Low | +| Evaluate flipping `DIFF_DIFF_SOLVE_OLS_FASTPATH` default-ON after an opt-in soak (the 2026-07 certified normal-equations Cholesky fast path, both backends). A flip needs: golden/parity-suite recapture at the tol-bounded posture (fitted ~1e-8 abs / SE ~1e-6 rel — the default today is byte-pinned in several benchmark conventions), certification-rate telemetry across real workloads (any decline is silent-correct but forfeits the speedup), and the staged default-flip protocol used for `df_convention` (v4-class change). | `diff_diff/linalg.py::_resolve_solve_ols_fastpath`, `rust/src/linalg.rs::solve_ols_chol` | CS-scaling | Mid | Low | ### Testing / docs @@ -114,8 +114,8 @@ Doable in principle, but no current caller and/or explicitly out of paper scope. | CallawaySantAnna **unbalanced-panel R parity — LANDED** via `allow_unbalanced_panel=True` (matches R `did::att_gt(allow_unbalanced_panel=TRUE)` / `DRDID::reg_did_rc`: ATT bit-exact on cells AND dynamic aggregation via fixed unit-cohort-mass `pg` + a per-unit WIF; SE up to the documented CR1 `sqrt(G/(G-1))` factor). The earlier "weighting" framing was a mis-diagnosis — on unbalanced panels the dominant divergence from R is the *estimator* (within-cell differencing vs RC-on-pooled-obs), not only the weighting; both are resolved by the flag. The DEFAULT path keeps within-cell differencing as a documented design choice and now emits a `UserWarning` on unbalanced input (no-silent-failures). **Remaining deferred:** `survey_design=` × `allow_unbalanced_panel=` (per-obs vs per-unit weight resolution — currently fail-closed `NotImplementedError`); and covariate / ipw / dr × the flag R-parity verification (the RC path supports them; the committed golden covers `reg` no-cov). | `staggered.py`, `staggered_aggregation.py` | SE-audit D3 | Low | | CallawaySantAnna event-study bucket/weight construction is duplicated between the analytical aggregator (`staggered_aggregation.py::_aggregate_event_study`) and the multiplier bootstrap (`staggered_bootstrap.py`): both group (g,t) by `e = t - g`, apply the finite/NaN/reference masks, and read cohort weights. Both already consume the same source-materialized universal reference cells (so they agree), but the bucket logic is copy-pasted. Extract one shared helper returning per-event-time buckets (finite cells, NaN cells, reference flags, cohort weights, combined-IF inputs) used by both. Pure refactor; gate on byte-identical analytical + bootstrap output. | `staggered_aggregation.py`, `staggered_bootstrap.py` | SE-audit D3 | Low | | `StackedDiD` survey re-resolution intra-file dedup (raw-weight extraction ×3, compose-normalize ×3, resolve-on-stacked ×2). The cross-estimator ContinuousDiD/EfficientDiD panel-to-unit collapse consolidation LANDED (#226 shared helpers `ResolvedSurveyDesign.subset_to_units_by_row_idx` / `build_unit_first_row_index`); StackedDiD is deliberately NOT on that path (control units are duplicated across sub-experiments, so it re-resolves at stacked granularity rather than collapsing to one row per unit). The residual is stacked-specific, low value, and touches the numerically-sensitive composed-weight renormalization. Post-filter re-resolution / metadata-recompute unification across the three estimators was assessed and is not warranted — they use genuinely different mechanisms and already delegate to shared `_resolve_survey_for_fit` / `compute_survey_metadata`. | `stacked_did.py` | #226 | Low | -| `solve_ols` residual memory floor is inherent to SVD-based lstsq after the PR-E marshalling slim (rust-side allocator high-water 15.32 -> 7.81 GB on the 2.4M x 130 clustered solve; remaining = one fused equilibrated input copy + thin-SVD U transient + vcov scores block + LAPACK gelsd internals on the python path). Further reduction requires the tall-skinny-QR algorithm change - same trade-off as the parked QR-reuse row below (library-wide last-digit perturbation + loses the second independent collinearity detector); the two land together if ever reopened. Diagnosis + measurements in docs/performance-plan.md. | `rust/src/linalg.rs::solve_ols`, `diff_diff/linalg.py` | PR-E | Low | -| SunAbraham-family fits are SOLVER-bound after the v3.6.x demeaning speedups: on the county-class shape (3.1k units x 60 months, ~100 interaction columns) `solve_ols` = ~60% of fit time (faer SVD 1.24s + pivoted-QR rank detection 0.70s of 3.27s), and the share grows further under the Rust demean kernel; the pivoted QR alone is ~20% of the whole fit (measured attribution in docs/performance-plan.md). Candidate fix is QR-reuse (one factorization for rank detection + solve). **Parked per the 2026-07 correctness-first decision**: the change perturbs every coefficient in the library at the last digits (golden/parity/bit-identity recapture) and removes the SVD's independent singular-value truncation - a second, different near-collinearity detector that is deliberate defense-in-depth for the no-silent-failures contract (the FE-span guard episode showed borderline cases are real). Re-open only on practitioner demand for wide event studies; any implementation gates on fixest/R end-to-end parity, not just internal consistency. | `linalg.py::solve_ols`, `linalg.py::_detect_rank_deficiency` | PR-C attribution | Medium | +| `solve_ols` residual memory floor is inherent to SVD-based lstsq after the PR-E marshalling slim (rust-side allocator high-water 15.32 -> 7.81 GB on the 2.4M x 130 clustered solve; remaining = one fused equilibrated input copy + thin-SVD U transient + vcov scores block + LAPACK gelsd internals on the python path). Further reduction requires the tall-skinny-QR algorithm change - same trade-off as the parked QR-reuse row below (library-wide last-digit perturbation + loses the second independent collinearity detector); the two land together if ever reopened. Diagnosis + measurements in docs/performance-plan.md. (The 2026-07 opt-in `DIFF_DIFF_SOLVE_OLS_FASTPATH` Cholesky path sidesteps the floor when enabled — no thin-SVD U transient, scores reuse the one equilibrated buffer; this row tracks the DEFAULT path only.) | `rust/src/linalg.rs::solve_ols`, `diff_diff/linalg.py` | PR-E | Low | +| SunAbraham-family fits are SOLVER-bound after the v3.6.x demeaning speedups: on the county-class shape (3.1k units x 60 months, ~100 interaction columns) `solve_ols` = ~60% of fit time (faer SVD 1.24s + pivoted-QR rank detection 0.70s of 3.27s), and the share grows further under the Rust demean kernel; the pivoted QR alone is ~20% of the whole fit (measured attribution in docs/performance-plan.md). Candidate fix is QR-reuse (one factorization for rank detection + solve). **Parked per the 2026-07 correctness-first decision**: the change perturbs every coefficient in the library at the last digits (golden/parity/bit-identity recapture) and removes the SVD's independent singular-value truncation - a second, different near-collinearity detector that is deliberate defense-in-depth for the no-silent-failures contract (the FE-span guard episode showed borderline cases are real). Re-open only on practitioner demand for wide event studies; any implementation gates on fixest/R end-to-end parity, not just internal consistency. (2026-07 updates: the pivoted-QR share quoted above predates #634's stage-0 Gram certification — rank detection is now ~0.03s on the county shape — and the opt-in `DIFF_DIFF_SOLVE_OLS_FASTPATH` Cholesky path is the certification-gated alternative that captures the solver win WITHOUT the library-wide perturbation: the default stays byte-identical and keeps the SVD truncation as the second collinearity detector.) | `linalg.py::solve_ols`, `linalg.py::_detect_rank_deficiency` | PR-C attribution | Medium | | dCDH parity-test SE/CI assertions only cover pure-direction scenarios; mixed-direction SE comparison is structurally apples-to-oranges (cell-count vs obs-count weighting). | `test_chaisemartin_dhaultfoeuille_parity.py` | #294 | Low | | `HeterogeneousAdoptionDiD` joint cross-horizon covariance / sup-t bands: per-horizon SEs use independent sandwiches (paper-faithful pointwise CIs per Pierce-Schott Fig 2). Follow-ups (low demand): IF-based stacking for joint cross-horizon inference; analytical H×H covariance on the weighted ES path. (A sup-t band on the unweighted ES path shipped via the clustered band — `cluster=` fires the simultaneous band on the unweighted path.) | `had.py::_fit_event_study` | Phase 2b / 4.5 B | Low | | `HeterogeneousAdoptionDiD` event-study staggered timing beyond the last cohort: Phase 2b auto-filters to the last cohort (paper App B.2); earlier-cohort effects aren't HAD-identified (redirect to dCDH). Full staggered HAD needs a different identification path (out of paper scope). | `had.py::_validate_had_panel_event_study` | Phase 2b | Low | diff --git a/benchmarks/speed_review/baselines/solve_ols_fastpath_after.json b/benchmarks/speed_review/baselines/solve_ols_fastpath_after.json new file mode 100644 index 00000000..d3183d1e --- /dev/null +++ b/benchmarks/speed_review/baselines/solve_ols_fastpath_after.json @@ -0,0 +1,135 @@ +{ + "metadata": { + "platform": "macOS-26.5.2-arm64-arm-64bit-Mach-O", + "arms": [ + "off", + "on" + ] + }, + "county_policy": { + "off": { + "times": [ + 1.127356459001021, + 1.1081956670022919, + 1.1512349999975413 + ], + "median": 1.127356459001021, + "att": 0.28037854139903057, + "se": 0.031244864610215983 + }, + "on": { + "times": [ + 0.6245858750044135, + 0.6251335830020253, + 0.6231642499988084 + ], + "median": 0.6245858750044135, + "att": 0.28037854139903007, + "se": 0.031244864610213023 + } + }, + "firm_churn": { + "off": { + "times": [ + 17.31286900000123, + 16.22713416699844 + ], + "median": 16.770001583499834, + "att": 0.314432718449715, + "se": 0.005837281701253922 + }, + "on": { + "times": [ + 9.93284683299862, + 9.850229166004283 + ], + "median": 9.891537999501452, + "att": 0.3144327184497242, + "se": 0.005837281701238088 + } + }, + "scanner_twfe": { + "off": { + "times": [ + 0.5774591250010417, + 0.5748795420004171, + 0.5841960840043612 + ], + "median": 0.5774591250010417, + "att": 0.25116413149499167, + "se": 0.002414938742160611 + }, + "on": { + "times": [ + 0.5765144159959164, + 0.5650162499950966, + 0.5698209999973187 + ], + "median": 0.5698209999973187, + "att": 0.25116413149499867, + "se": 0.0024149387421607276 + } + }, + "cs_cov40": { + "off": { + "times": [ + 4.3890298749975045, + 4.401376791000075, + 4.414817625001888 + ], + "median": 4.401376791000075, + "att": 2.0008128010664223, + "se": 0.005610568442727097 + }, + "on": { + "times": [ + 3.700359333997767, + 3.5647381669987226, + 3.5419092910015024 + ], + "median": 3.5647381669987226, + "att": 2.0008128010664223, + "se": 0.005610568442727097 + } + }, + "survey_absorb": { + "off": { + "times": [ + 2.7668485419999342, + 2.7689070839987835 + ], + "median": 2.767877812999359, + "att": 0.24554568144130481, + "se": 0.007662698712656388 + }, + "on": { + "times": [ + 2.5810511670060805, + 2.595715333998669 + ], + "median": 2.5883832505023747, + "att": 0.24554568144130481, + "se": 0.00766269871265628 + } + }, + "skiprank_micro": { + "off": { + "times": [ + 6.819884582997474, + 6.934963792002236 + ], + "median": 6.877424187499855, + "att": 2.6266509653021792, + "se": 0.0006455844619141583 + }, + "on": { + "times": [ + 4.4150471670000115, + 4.423518374998821 + ], + "median": 4.419282770999416, + "att": 2.626650965303164, + "se": 0.0006455844619141326 + } + } +} \ No newline at end of file diff --git a/benchmarks/speed_review/baselines/solve_ols_fastpath_before.json b/benchmarks/speed_review/baselines/solve_ols_fastpath_before.json new file mode 100644 index 00000000..f5a7d197 --- /dev/null +++ b/benchmarks/speed_review/baselines/solve_ols_fastpath_before.json @@ -0,0 +1,77 @@ +{ + "metadata": { + "platform": "macOS-26.5.2-arm64-arm-64bit-Mach-O", + "arms": [ + "off" + ] + }, + "county_policy": { + "off": { + "times": [ + 1.0828709579946008, + 1.1249838750009076, + 1.0858678750009858 + ], + "median": 1.0858678750009858, + "att": 0.28037854139903057, + "se": 0.031244864610215983 + } + }, + "firm_churn": { + "off": { + "times": [ + 17.221773457997188, + 16.04917508300423 + ], + "median": 16.63547427050071, + "att": 0.314432718449715, + "se": 0.005837281701253922 + } + }, + "scanner_twfe": { + "off": { + "times": [ + 0.575249708002957, + 0.5752674159957678, + 0.5818307080044178 + ], + "median": 0.5752674159957678, + "att": 0.25116413149499167, + "se": 0.002414938742160611 + } + }, + "cs_cov40": { + "off": { + "times": [ + 4.368334792001406, + 4.373289832998125, + 4.364916250000533 + ], + "median": 4.368334792001406, + "att": 2.0008128010664223, + "se": 0.005610568442727097 + } + }, + "survey_absorb": { + "off": { + "times": [ + 2.8271069999973406, + 2.837882542000443 + ], + "median": 2.832494770998892, + "att": 0.24554568144130481, + "se": 0.007662698712656388 + } + }, + "skiprank_micro": { + "off": { + "times": [ + 7.0087420839990955, + 7.03886975000205 + ], + "median": 7.023805917000573, + "att": 2.6266509653021792, + "se": 0.0006455844619141583 + } + } +} \ No newline at end of file diff --git a/benchmarks/speed_review/bench_solve_ols_fastpath.py b/benchmarks/speed_review/bench_solve_ols_fastpath.py new file mode 100644 index 00000000..cf4f3daa --- /dev/null +++ b/benchmarks/speed_review/bench_solve_ols_fastpath.py @@ -0,0 +1,293 @@ +"""solve_ols opt-in Cholesky fast-path A/B benchmark (DIFF_DIFF_SOLVE_OLS_FASTPATH). + +Each scenario runs in a fresh subprocess (the speed_review noise protocol); +within a scenario the knob-off and knob-on arms share that process — the +knob is resolved per call, so same-process A/B is exact. Scenarios are the +solver-bound shapes from the 2026-07 attribution: + +- ``county_policy`` / ``firm_churn`` - SunAbraham saturated fits (fe_absorption shapes) +- ``scanner_twfe`` - TwoWayFixedEffects (demean-bound control arm) +- ``cs_cov40`` - CallawaySantAnna dr, 40 covariates, 2M rows + (per-cell solver floor) +- ``survey_absorb`` - DiD absorb + BRR replicates (weighted lane -> + numpy twin coverage) +- ``skiprank_micro`` - direct solve_ols(skip_rank_check=True) on a + firm-shape matrix (self-certification lane) + +The knob-off arm doubles as the BYTE-IDENTITY gate: run this script once on +the pristine base tree (``PYTHONPATH= --arms off --out before.json``) +and once on the branch (``--arms off,on``), then ``--check-identity before.json`` +asserts the knob-off ATT/SE are exactly equal to the base tree's (the default +path must be untouched) and reports the knob-on deltas. + +Usage:: + + python benchmarks/speed_review/bench_solve_ols_fastpath.py --repeats 3 \ + --out benchmarks/speed_review/baselines/solve_ols_fastpath_after.json \ + --check-identity benchmarks/speed_review/baselines/solve_ols_fastpath_before.json +""" + +import argparse +import json +import os +import platform +import statistics +import sys +import time +from pathlib import Path + +import numpy as np +import pandas as pd + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +import fe_absorption_datagen as datagen # noqa: E402 + +KNOB = "DIFF_DIFF_SOLVE_OLS_FASTPATH" + +# (scenario, timed fits per arm) +SCENARIOS = { + "county_policy": 3, + "firm_churn": 2, + "scanner_twfe": 3, + "cs_cov40": 3, + "survey_absorb": 2, + "skiprank_micro": 2, +} + + +def make_cs_panel(n_units=100_000, n_periods=20, n_cov=40, seed=0): + """Staggered panel, 5 treated cohorts + never-treated, covariate-selected + cohorts (same DGP as the 2026-07 attribution profiling).""" + rng = np.random.default_rng(seed) + X = rng.standard_normal((n_units, n_cov)) + score = 0.4 * X[:, 0] - 0.3 * X[:, 1] + 0.2 * X[:, 2] + rng.standard_normal(n_units) + qs = np.quantile(score, [0.35, 0.48, 0.61, 0.74, 0.87]) + cohort_vals = np.array([0, 4, 7, 10, 13, 16]) + g_unit = cohort_vals[np.searchsorted(qs, score)] + unit = np.repeat(np.arange(n_units, dtype=np.int64), n_periods) + t = np.tile(np.arange(n_periods, dtype=np.int64), n_units) + g = np.repeat(g_unit, n_periods) + alpha = np.repeat(rng.normal(0, 1, n_units), n_periods) + beta = rng.normal(0.1, 0.05, n_cov) + xb = np.repeat(X @ beta, n_periods) + treated = (g > 0) & (t >= g) + y = alpha + 0.05 * t + xb + 2.0 * treated + rng.normal(0, 1, n_units * n_periods) + cols = {"unit": unit, "time": t, "first_treat": g, "y": y} + for j in range(n_cov): + cols[f"x{j}"] = np.repeat(X[:, j], n_periods) + return pd.DataFrame(cols) + + +def make_skiprank_matrix(n=2_400_000, k=130, seed=1): + rng = np.random.default_rng(seed) + X = np.column_stack([np.ones(n), rng.standard_normal((n, k - 1))]) + y = X @ rng.normal(0, 0.5, k) + rng.standard_normal(n) + return X, y + + +def build(scenario): + """Return (payload, fit) where fit(payload) -> (att-like, se-like).""" + if scenario in ("county_policy", "firm_churn"): + df, _ = datagen.build(scenario) + from diff_diff import SunAbraham + + def fit(d): + res = SunAbraham().fit( + d, outcome="y", unit="unit", time="time", first_treat="first_treat" + ) + return float(res.att), float(res.se) + + return df, fit + if scenario == "scanner_twfe": + df, _ = datagen.build(scenario) + from diff_diff import TwoWayFixedEffects + + def fit(d): + res = TwoWayFixedEffects().fit( + d, outcome="y", treatment="treated", time="post", unit="unit" + ) + return float(res.att), float(res.se) + + return df, fit + if scenario == "cs_cov40": + df = make_cs_panel() + from diff_diff import CallawaySantAnna + + covs = [f"x{j}" for j in range(40)] + + def fit(d): + res = CallawaySantAnna(estimation_method="dr").fit( + d, + outcome="y", + unit="unit", + time="time", + first_treat="first_treat", + covariates=covs, + aggregate="all", + ) + return float(res.overall_att), float(res.overall_se) + + return df, fit + if scenario == "survey_absorb": + df, _ = datagen.build(scenario) + from diff_diff import DifferenceInDifferences + from diff_diff.survey import SurveyDesign + + rep_cols = [c for c in df.columns if c.startswith("rw")] + + def fit(d): + design = SurveyDesign(weights="w", replicate_weights=rep_cols, replicate_method="BRR") + res = DifferenceInDifferences().fit( + d, + outcome="y", + treatment="treated", + time="post", + absorb=["state", "month"], + survey_design=design, + ) + return float(res.att), float(res.se) + + return df, fit + if scenario == "skiprank_micro": + X, y = make_skiprank_matrix() + from diff_diff.linalg import solve_ols + + def fit(payload): + Xm, ym = payload + coef, _, vcov = solve_ols(Xm, ym, skip_rank_check=True) + return float(np.sum(coef)), float(np.sqrt(vcov[1, 1])) + + return (X, y), fit + raise ValueError(scenario) + + +def run_scenario(scenario, arms, repeats_scale): + import warnings + + warnings.filterwarnings("ignore") + payload, fit = build(scenario) + n_fits = max(1, round(SCENARIOS[scenario] * repeats_scale)) + + # Warmup once through the same front door (imports, BLAS/rayon init). + os.environ.pop(KNOB, None) + fit(payload) + + out = {} + for arm in arms: + if arm == "on": + os.environ[KNOB] = "1" + else: + os.environ.pop(KNOB, None) + times = [] + att = se = None + for _ in range(n_fits): + t0 = time.perf_counter() + att, se = fit(payload) + times.append(time.perf_counter() - t0) + out[arm] = { + "times": times, + "median": statistics.median(times), + "att": att, + "se": se, + } + os.environ.pop(KNOB, None) + return out + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--arms", default="off,on", help="comma list: off,on") + ap.add_argument("--only", default=None, help="comma list of scenario ids") + ap.add_argument("--repeats", type=float, default=1.0, help="scale timed-fit counts") + ap.add_argument("--out", default=None) + ap.add_argument("--check-identity", default=None, metavar="BEFORE_JSON") + ap.add_argument( + "--in-process", + action="store_true", + help="run scenarios in this process (default: one fresh subprocess per " + "scenario, the speed_review noise protocol — in-process runs showed " + "cross-scenario state contaminating arm timings by ~25%%)", + ) + args = ap.parse_args() + + arms = [a.strip() for a in args.arms.split(",") if a.strip()] + ids = list(SCENARIOS) if args.only is None else args.only.split(",") + + results = {"metadata": {"platform": platform.platform(), "arms": arms}} + if not args.in_process and len(ids) > 1: + # Fresh subprocess per scenario, strictly sequential (matches + # bench_fe_absorption's noise protocol). + import subprocess + import tempfile + + for scenario in ids: + with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as tf: + tmp = tf.name + cmd = [ + sys.executable, + str(Path(__file__).resolve()), + "--only", + scenario, + "--arms", + args.arms, + "--repeats", + str(args.repeats), + "--out", + tmp, + "--in-process", + ] + subprocess.run(cmd, check=True) + with open(tmp) as f: + results[scenario] = json.load(f)[scenario] + os.unlink(tmp) + else: + for scenario in ids: + print(f"=== {scenario}", flush=True) + results[scenario] = run_scenario(scenario, arms, args.repeats) + for arm, r in results[scenario].items(): + print( + f" {arm:>3}: median {r['median']:.3f}s att={r['att']:.6f} se={r['se']:.6g}", + flush=True, + ) + + if args.out: + Path(args.out).parent.mkdir(parents=True, exist_ok=True) + with open(args.out, "w") as f: + json.dump(results, f, indent=2) + print(f"wrote {args.out}") + + if args.check_identity: + with open(args.check_identity) as f: + before = json.load(f) + failures = [] + for scenario in ids: + if scenario not in before or scenario not in results: + continue + b = before[scenario]["off"] + a_off = results[scenario]["off"] + # Byte-identity gate: knob-off on the branch == the base tree. + if not (b["att"] == a_off["att"] and b["se"] == a_off["se"]): + failures.append( + f"{scenario}: knob-off (att={a_off['att']!r}, se={a_off['se']!r}) " + f"!= base tree (att={b['att']!r}, se={b['se']!r})" + ) + if "on" in results[scenario]: + a_on = results[scenario]["on"] + d_att = abs(a_on["att"] - a_off["att"]) + d_se = abs(a_on["se"] - a_off["se"]) / max(abs(a_off["se"]), 1e-300) + speedup = a_off["median"] / a_on["median"] + print( + f"{scenario}: speedup {speedup:.2f}x |dATT|={d_att:.2e} " + f"|dSE|/SE={d_se:.2e}" + ) + if failures: + print("IDENTITY GATE FAILED:") + for f_ in failures: + print(" " + f_) + sys.exit(1) + print("identity gate: knob-off byte-identical to base tree ✓") + + +if __name__ == "__main__": + main() diff --git a/diff_diff/_backend.py b/diff_diff/_backend.py index 9175fcbb..2d379d8c 100644 --- a/diff_diff/_backend.py +++ b/diff_diff/_backend.py @@ -93,6 +93,17 @@ except ImportError: _rust_compute_robust_vcov_hc2 = None +# Opt-in normal-equations Cholesky OLS fast path: imported independently +# for the same mixed-version reason as demean_map. A stale extension +# missing only this symbol keeps every older Rust acceleration: Rust-eligible +# fits fall back to the legacy SVD solve_ols kernel (the knob simply has no +# Rust acceleration there), while numpy-lane fits (weighted, non-hc1, +# forced-python) still use the numpy Cholesky twin. +try: + from diff_diff._rust_backend import solve_ols_chol as _rust_solve_ols_chol +except ImportError: + _rust_solve_ols_chol = None + # Determine final backend based on environment variable and availability if _backend_env == "python": # Force pure Python mode - disable Rust even if available @@ -107,6 +118,8 @@ _rust_batched_ridge_chol_solve = None # HC2 robust vcov _rust_compute_robust_vcov_hc2 = None + # Opt-in normal-equations Cholesky OLS fast path + _rust_solve_ols_chol = None # TROP estimator acceleration (local method) _rust_unit_distance_matrix = None _rust_loocv_grid_search = None @@ -158,6 +171,8 @@ def rust_backend_info(): "_rust_project_simplex", "_rust_solve_ols", "_rust_compute_robust_vcov", + # Opt-in normal-equations Cholesky OLS fast path + "_rust_solve_ols_chol", # FE-absorption MAP demeaning kernel "_rust_demean_map", # Batched ridge-regularized SPD solve (EfficientDiD per-unit weights) diff --git a/diff_diff/linalg.py b/diff_diff/linalg.py index 28b75a00..9aeb7ab9 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -32,6 +32,7 @@ - "silent": No warning, but still set NA for dropped coefficients """ +import os import warnings from dataclasses import dataclass from typing import Any, Dict, List, Literal, Optional, Tuple, Union, overload @@ -49,6 +50,7 @@ _rust_compute_robust_vcov, _rust_compute_robust_vcov_hc2, _rust_solve_ols, + _rust_solve_ols_chol, ) # Conley (1999) spatial HAC helpers live in diff_diff.conley to keep this @@ -96,6 +98,8 @@ def _factorize_cluster_ids(cluster_ids: np.ndarray) -> np.ndarray: def _detect_rank_deficiency( X: np.ndarray, rcond: Optional[float] = None, + *, + _cert_out: Optional[dict] = None, ) -> Tuple[int, np.ndarray, np.ndarray]: """ Detect rank deficiency using pivoted QR decomposition. @@ -137,6 +141,16 @@ def _detect_rank_deficiency( Column permutation from QR decomposition. For a full-rank result the pivot carries no information (no caller consumes it; the stage-0 certification below returns a trivial ``arange`` pivot). + + Other Parameters + ---------------- + _cert_out : dict, optional + Private out-parameter for the opt-in solve_ols Cholesky fast path. + When a dict is passed and stage-0 certification runs, it is populated + with the stage-0 artifacts (``gram``, ``scales``, ``gram_eq``, + ``eig_min``, ``eig_max`` and, on the certified branch, ``certified``) + so the caller can reuse the Gram work instead of rebuilding it. Pure + out-parameter: the return value and every rank decision are unchanged. """ n, k = X.shape if k == 0: @@ -174,12 +188,22 @@ def _detect_rank_deficiency( gram_eq = gram / scales[:, None] / scales[None, :] eigvals = np.linalg.eigvalsh(gram_eq) eig_min, eig_max = eigvals[0], eigvals[-1] + if _cert_out is not None: + _cert_out.update( + gram=gram, + scales=scales, + gram_eq=gram_eq, + eig_min=eig_min, + eig_max=eig_max, + ) if ( np.isfinite(eig_min) and np.isfinite(eig_max) and eig_max > 0.0 and eig_min > 1e-10 * eig_max ): + if _cert_out is not None: + _cert_out["certified"] = True return k, np.array([], dtype=int), np.arange(k, dtype=int) def _rank_and_pivot(M: np.ndarray) -> Tuple[int, np.ndarray]: @@ -410,6 +434,109 @@ def _equilibrated_lstsq(X: np.ndarray, y: np.ndarray) -> np.ndarray: return coef_scaled / safe_norms +# Reciprocal-condition guard for the opt-in solve_ols normal-equations Cholesky +# fast path. Same 1e-6 bound and rationale as _IRLS_CHOL_RCOND_GUARD (see the +# solve_logit IRLS inner solve): cond(G_eq) <= 1e6 bounds the Cholesky forward +# error at ~eps*cond ~ 2e-10 relative in the equilibrated basis. Kept as a +# separate constant so the two paths can be tuned independently. +_SOLVE_OLS_CHOL_RCOND_GUARD = 1e-6 + +# Module default for the opt-in fast path (OFF = byte-identical legacy +# behavior). The env var is read PER CALL (see resolver below) so benchmarks +# and tests can A/B within one process; this constant is the monkeypatch seam +# and the fallback for unset/invalid env values. +_SOLVE_OLS_FASTPATH: bool = False + + +def _resolve_solve_ols_fastpath() -> bool: + """Resolve the DIFF_DIFF_SOLVE_OLS_FASTPATH opt-in knob. + + Set to a positive integer (``1``) to enable the certification-gated + normal-equations Cholesky fast path in ``solve_ols``. Read PER CALL, not + at import, so a single process can A/B the two paths. Unset or invalid + values (non-integer strings, zero, negatives) fall back silently to the + module default, mirroring ``_resolve_demean_chunk_cols``'s convention. + """ + raw = os.environ.get("DIFF_DIFF_SOLVE_OLS_FASTPATH") + if raw is None: + return _SOLVE_OLS_FASTPATH + try: + value = int(raw) + except ValueError: + return _SOLVE_OLS_FASTPATH + return True if value > 0 else _SOLVE_OLS_FASTPATH + + +def _solve_ols_chol_numpy( + X: np.ndarray, + y: np.ndarray, + *, + cert_info: Optional[dict] = None, +) -> Optional[Tuple[np.ndarray, np.ndarray]]: + """Certified equilibrated normal-equations Cholesky solve (opt-in path). + + Returns ``(coefficients, gram_raw)`` on success, or ``None`` when + certification declines — the caller then falls back VERBATIM to the gelsd + path, so a decline is always output-identical to the fast path being off. + ``gram_raw`` is exactly the ``X.T @ X`` expression on the same array, so + the caller may reuse it as the robust-vcov bread matrix bit-for-bit. + + Guard chain mirrors the solve_logit IRLS inner solve: symmetric + equilibration of the Gram (== column equilibration of X), ``cho_factor``, + then an explicit ``dpocon`` reciprocal-condition estimate gated at + ``_SOLVE_OLS_CHOL_RCOND_GUARD`` — factorization success alone is NOT a + certificate (cho_factor can succeed with a garbage solution at + cond ~ 1e10+). Stage-0 artifacts from ``_detect_rank_deficiency`` are + reused when supplied (``cert_info``); the self-build branch exists for the + ``skip_rank_check`` route where no stage-0 certification ran. Artifacts + present but uncertified mean the design already failed the (looser) + stage-0 rank gate, so the solve gate cannot pass: decline immediately. + """ + n, k = X.shape + if k == 0 or n < k: + return None + + if cert_info is not None: + if "gram_eq" not in cert_info or not cert_info.get("certified", False): + return None + gram = cert_info["gram"] + scales = cert_info["scales"] + gram_eq = cert_info["gram_eq"] + else: + # Self-build (skip_rank_check route). Same expressions and guards as + # stage-0: decline unless diag is finite AND strictly positive (zero + # diag = zero/underflowed column; non-finite diag = NaN/Inf in X — + # any non-finite entry poisons its own column's diagonal). + with np.errstate(divide="ignore", over="ignore", invalid="ignore"): + gram = X.T @ X + diag = np.diag(gram) + if not (np.all(np.isfinite(diag)) and np.all(diag > 0)): + return None + scales = np.sqrt(diag) + gram_eq = gram / scales[:, None] / scales[None, :] + + # 1-norm BEFORE factorization (dpocon contract). + anorm = float(np.max(np.sum(np.abs(gram_eq), axis=0))) + if not np.isfinite(anorm): + return None + try: + chol = cho_factor(gram_eq) + except (np.linalg.LinAlgError, ValueError): + return None + rcond_gram, pocon_info = dpocon(chol[0], anorm) + if not ( + pocon_info == 0 and np.isfinite(rcond_gram) and rcond_gram > _SOLVE_OLS_CHOL_RCOND_GUARD + ): + return None + + with np.errstate(divide="ignore", over="ignore", invalid="ignore"): + xty = X.T @ y + coefficients = cho_solve(chol, xty / scales) / scales + if not np.all(np.isfinite(coefficients)): + return None + return coefficients, gram + + @overload def _rank_guarded_inv( A: np.ndarray, @@ -660,6 +787,88 @@ def _solve_ols_rust( return coefficients, residuals, vcov +def _solve_ols_chol_rust( + X: np.ndarray, + y: np.ndarray, + *, + cluster_ids: Optional[np.ndarray] = None, + return_vcov: bool = True, + return_fitted: bool = False, +) -> Optional[ + Union[ + Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]], + Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]], + ] +]: + """Rust certified normal-equations Cholesky solve (opt-in fast path). + + Mirrors ``_solve_ols_rust``'s wrapper contract. The kernel is + self-certifying and returns Python ``None`` on a certification decline + (zero/non-finite column norm, non-PD Gram, exact 1-norm rcond <= 1e-6, + n < k, non-finite coefficient); this wrapper propagates the ``None`` so + the dispatcher falls through VERBATIM to the SVD kernel and then the + numpy path. The "Need at least 2 clusters" ValueError from clustered + vcov re-raises exactly as on the SVD path (its message names neither + instability nor singularity). + """ + if cluster_ids is not None: + _validate_cluster_ids(cluster_ids) + cluster_ids = _factorize_cluster_ids(cluster_ids) + + try: + result = _rust_solve_ols_chol(X, y, cluster_ids=cluster_ids, return_vcov=return_vcov) + except ValueError as e: + error_msg = str(e).lower() + if "numerically unstable" in error_msg or "singular" in error_msg: + warnings.warn( + f"Rust backend detected numerical instability: {e}. " + "Falling back to Python backend.", + UserWarning, + stacklevel=3, + ) + return None + raise + + if result is None: + return None # certification declined — caller falls through verbatim + + coefficients, residuals, vcov = result + coefficients = np.asarray(coefficients) + residuals = np.asarray(residuals) + if vcov is not None: + vcov = np.asarray(vcov) + + if return_fitted: + fitted = np.dot(X, coefficients) + return coefficients, residuals, fitted, vcov + return coefficients, residuals, vcov + + +def _nonfinite_vcov_needs_python_rerun( + vcov: Optional[np.ndarray], *, nan_is_sentinel: bool +) -> bool: + """Shared guard: should a Rust-backend vcov be rejected in favor of the + canonical numpy re-run? + + On the rank-checked route (``nan_is_sentinel=False``) any non-finite + entry (NaN or Inf) reroutes: the design was certified full rank, so the + numpy re-run safely applies the canonical contracts (rank handling, + saturated all-NaN guard). + + On the ``skip_rank_check`` route (``nan_is_sentinel=True``) an all-NaN + vcov is the DOCUMENTED sentinel for what the caller asserted away + (rank deficiency / saturation) and passes through unchanged — rerouting + it to a numpy path that assumes full rank could raise on a singular + bread where users previously received the safe NaN answer. Only Inf + (e.g. numerical overflow, the silent-corruption class) reroutes there. + """ + if vcov is None: + return False + if nan_is_sentinel: + return bool(np.any(np.isinf(vcov))) + return not bool(np.all(np.isfinite(vcov))) + + @overload def solve_ols( X: np.ndarray, @@ -682,6 +891,7 @@ def solve_ols( conley_time: Optional[np.ndarray] = ..., conley_unit: Optional[np.ndarray] = ..., conley_lag_cutoff: Optional[int] = ..., + diagnostics_out: Optional[dict] = ..., ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]: ... @@ -707,6 +917,7 @@ def solve_ols( conley_time: Optional[np.ndarray] = ..., conley_unit: Optional[np.ndarray] = ..., conley_lag_cutoff: Optional[int] = ..., + diagnostics_out: Optional[dict] = ..., ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]: ... @@ -732,6 +943,7 @@ def solve_ols( conley_time: Optional[np.ndarray] = ..., conley_unit: Optional[np.ndarray] = ..., conley_lag_cutoff: Optional[int] = ..., + diagnostics_out: Optional[dict] = ..., ) -> Union[ Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]], Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]], @@ -793,6 +1005,7 @@ def solve_ols( conley_time: Optional[np.ndarray] = None, conley_unit: Optional[np.ndarray] = None, conley_lag_cutoff: Optional[int] = None, + diagnostics_out: Optional[dict] = None, ) -> Union[ Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]], Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]], @@ -891,6 +1104,17 @@ def solve_ols( a ``UserWarning`` if the resulting meat is materially indefinite; the radial 1-D Bartlett (matching R ``conleyreg``) is not formally PSD-guaranteed — see :func:`compute_robust_vcov`. + diagnostics_out : dict, optional + Observability sink (zero-cost when None). Populated with + ``"solve_ols_fastpath"`` ∈ {``"off"``, ``"chol_numpy"``, + ``"chol_rust"``, ``"fallback_declined"``} recording which solve + branch produced the coefficients under the opt-in + ``DIFF_DIFF_SOLVE_OLS_FASTPATH`` normal-equations Cholesky fast + path (``"off"`` when the knob is unset; ``"fallback_declined"`` + when the knob is on but the fast path did not produce the result — + certification declined, the design was rank-deficient/structurally + ineligible, or the Rust fast-path symbol was unavailable — and the + verbatim legacy chain ran instead). Returns ------- @@ -1014,12 +1238,42 @@ def solve_ols( _weighted_vcov_external = weights is not None _backend_return_vcov = return_vcov and not _weighted_vcov_external + # Opt-in normal-equations Cholesky fast path (DIFF_DIFF_SOLVE_OLS_FASTPATH). + # Resolved per call so a single process can A/B; OFF (the default) leaves + # every line below byte-identical to the legacy path. On a certification + # decline the fast path returns None and execution falls through the + # UNCHANGED legacy chain (Rust SVD, then gelsd), so knob-on-decline output + # equals knob-off output exactly. + fastpath = _resolve_solve_ols_fastpath() + if diagnostics_out is not None: + # Pessimistic default; overwritten with the branch actually taken. + diagnostics_out["solve_ols_fastpath"] = "fallback_declined" if fastpath else "off" + # Fast path: skip rank check and use Rust directly when requested # This saves O(nk²) QR overhead but won't detect rank-deficient matrices result = None # Will hold the tuple from backend functions if skip_rank_check: if ( + fastpath + and HAS_RUST_BACKEND + and _rust_solve_ols_chol is not None + and weights is None + and vcov_type == "hc1" + ): + # Self-certifying Rust Cholesky kernel (no stage-0 cert exists + # on this route). None = certification declined; fall through + # to the UNCHANGED legacy chain (Rust SVD, then numpy). + result = _solve_ols_chol_rust( + X, + y, + cluster_ids=cluster_ids, + return_vcov=_backend_return_vcov, + return_fitted=return_fitted, + ) + if result is not None and diagnostics_out is not None: + diagnostics_out["solve_ols_fastpath"] = "chol_rust" + if result is None and ( HAS_RUST_BACKEND and _rust_solve_ols is not None and weights is None @@ -1033,6 +1287,18 @@ def solve_ols( return_fitted=return_fitted, ) # result is None on numerical instability → fall through + if result is not None and _nonfinite_vcov_needs_python_rerun( + result[-1] if _backend_return_vcov else None, + nan_is_sentinel=True, + ): + warnings.warn( + "Rust backend detected ill-conditioned matrix (non-finite " + "variance-covariance). Re-running with Python backend for " + "proper rank detection.", + UserWarning, + stacklevel=2, + ) + result = None # Force Python fallback below if result is None: result = _solve_ols_numpy( X, @@ -1043,6 +1309,8 @@ def solve_ols( rank_deficient_action=rank_deficient_action, column_names=column_names, _skip_rank_check=True, + _fastpath=fastpath, + _diagnostics_out=diagnostics_out, vcov_type=vcov_type, conley_coords=conley_coords, conley_cutoff_km=conley_cutoff_km, @@ -1055,8 +1323,13 @@ def solve_ols( else: # Check for rank deficiency using fast pivoted QR decomposition. # Rank detection operates on (possibly weighted) X since collinearity - # depends on the weighted column space. - rank, dropped_cols, pivot = _detect_rank_deficiency(X) + # depends on the weighted column space. When the fast path is on, + # collect the stage-0 certification artifacts (Gram, scales, + # eigenvalues) so the Cholesky solve can reuse them instead of + # rebuilding the Gram — the artifacts are of the (possibly weighted) + # X, which is exactly the matrix being solved. + cert_out: Optional[dict] = {} if fastpath else None + rank, dropped_cols, pivot = _detect_rank_deficiency(X, _cert_out=cert_out) is_rank_deficient = len(dropped_cols) > 0 # Routing strategy: @@ -1064,6 +1337,27 @@ def solve_ols( # - Weighted or rank-deficient or non-HC1 vcov_type → Python backend # - Rust numerical instability → Python fallback (via None return) if ( + fastpath + and HAS_RUST_BACKEND + and _rust_solve_ols_chol is not None + and not is_rank_deficient + and weights is None + and vcov_type == "hc1" + ): + # Self-certifying Rust Cholesky kernel (rebuilds its own + # equilibrated Gram; the stage-0 artifacts stay with the numpy + # twin). None = certification declined; fall through to the + # UNCHANGED legacy chain below. + result = _solve_ols_chol_rust( + X, + y, + cluster_ids=cluster_ids, + return_vcov=_backend_return_vcov, + return_fitted=return_fitted, + ) + if result is not None and diagnostics_out is not None: + diagnostics_out["solve_ols_fastpath"] = "chol_rust" + if result is None and ( HAS_RUST_BACKEND and _rust_solve_ols is not None and not is_rank_deficient @@ -1079,11 +1373,15 @@ def solve_ols( ) if result is not None: - vcov_check = result[-1] - if _backend_return_vcov and vcov_check is not None and np.any(np.isnan(vcov_check)): + vcov_check = result[-1] if _backend_return_vcov else None + if _nonfinite_vcov_needs_python_rerun(vcov_check, nan_is_sentinel=False): + # Non-finite covers both the NaN rank-deficiency sentinel + # and Inf leakage (e.g. numerical overflow); the Python + # backend re-run applies the canonical numpy contracts. warnings.warn( - "Rust backend detected ill-conditioned matrix (NaN in variance-covariance). " - "Re-running with Python backend for proper rank detection.", + "Rust backend detected ill-conditioned matrix (non-finite " + "variance-covariance). Re-running with Python backend for " + "proper rank detection.", UserWarning, stacklevel=2, ) @@ -1099,6 +1397,9 @@ def solve_ols( rank_deficient_action=rank_deficient_action, column_names=column_names, _precomputed_rank_info=(rank, dropped_cols, pivot), + _fastpath=fastpath, + _cert_info=cert_out, + _diagnostics_out=diagnostics_out, vcov_type=vcov_type, conley_coords=conley_coords, conley_cutoff_km=conley_cutoff_km, @@ -1187,6 +1488,9 @@ def _solve_ols_numpy( column_names: Optional[List[str]] = ..., _precomputed_rank_info: Optional[Tuple[int, np.ndarray, np.ndarray]] = ..., _skip_rank_check: bool = ..., + _fastpath: bool = ..., + _cert_info: Optional[dict] = ..., + _diagnostics_out: Optional[dict] = ..., vcov_type: str = ..., conley_coords: Optional[np.ndarray] = ..., conley_cutoff_km: Optional[float] = ..., @@ -1210,6 +1514,9 @@ def _solve_ols_numpy( column_names: Optional[List[str]] = ..., _precomputed_rank_info: Optional[Tuple[int, np.ndarray, np.ndarray]] = ..., _skip_rank_check: bool = ..., + _fastpath: bool = ..., + _cert_info: Optional[dict] = ..., + _diagnostics_out: Optional[dict] = ..., vcov_type: str = ..., conley_coords: Optional[np.ndarray] = ..., conley_cutoff_km: Optional[float] = ..., @@ -1233,6 +1540,9 @@ def _solve_ols_numpy( column_names: Optional[List[str]] = ..., _precomputed_rank_info: Optional[Tuple[int, np.ndarray, np.ndarray]] = ..., _skip_rank_check: bool = ..., + _fastpath: bool = ..., + _cert_info: Optional[dict] = ..., + _diagnostics_out: Optional[dict] = ..., vcov_type: str = ..., conley_coords: Optional[np.ndarray] = ..., conley_cutoff_km: Optional[float] = ..., @@ -1258,6 +1568,9 @@ def _solve_ols_numpy( column_names: Optional[List[str]] = None, _precomputed_rank_info: Optional[Tuple[int, np.ndarray, np.ndarray]] = None, _skip_rank_check: bool = False, + _fastpath: bool = False, + _cert_info: Optional[dict] = None, + _diagnostics_out: Optional[dict] = None, vcov_type: str = "hc1", conley_coords: Optional[np.ndarray] = None, conley_cutoff_km: Optional[float] = None, @@ -1299,6 +1612,17 @@ def _solve_ols_numpy( _skip_rank_check : bool, default False If True, skip rank detection entirely and assume full rank. Used when caller has already determined matrix is full rank. + _fastpath : bool, default False + Opt-in normal-equations Cholesky fast path (resolved from + DIFF_DIFF_SOLVE_OLS_FASTPATH by solve_ols). Only the full-rank + branch is affected; a certification decline falls back verbatim + to the gelsd solve. + _cert_info : dict, optional + Stage-0 certification artifacts from _detect_rank_deficiency + (Gram reuse). None on the _skip_rank_check route — the fast path + then self-builds and self-certifies its Gram. + _diagnostics_out : dict, optional + solve_ols's diagnostics sink; records which solve branch ran. Returns ------- @@ -1400,7 +1724,21 @@ def _solve_ols_numpy( # Full-rank case: proceed normally. Equilibrate columns before the lstsq # so a large-scale column cannot truncate the genuine small-scale # direction (cond=1e-07 is relative to the largest singular value). - coefficients = _equilibrated_lstsq(X, y) + # + # Opt-in fast path (DIFF_DIFF_SOLVE_OLS_FASTPATH): try the certified + # normal-equations Cholesky solve first; a None return (certification + # declined) falls through to the verbatim gelsd line below, so the + # decline case is output-identical to the fast path being off. + coefficients = None + gram_bread = None + if _fastpath: + fast = _solve_ols_chol_numpy(X, y, cert_info=_cert_info) + if fast is not None: + coefficients, gram_bread = fast + if _diagnostics_out is not None: + _diagnostics_out["solve_ols_fastpath"] = "chol_numpy" + if coefficients is None: + coefficients = _equilibrated_lstsq(X, y) # Compute residuals and fitted values fitted = np.dot(X, coefficients) @@ -1421,6 +1759,7 @@ def _solve_ols_numpy( conley_time=conley_time, conley_unit=conley_unit, conley_lag_cutoff=conley_lag_cutoff, + _bread_matrix=gram_bread, ) if return_fitted: @@ -2882,11 +3221,18 @@ def _compute_robust_vcov_numpy( conley_time: Optional[np.ndarray] = None, conley_unit: Optional[np.ndarray] = None, conley_lag_cutoff: Optional[int] = None, + _bread_matrix: Optional[np.ndarray] = None, ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: """ NumPy fallback implementation of compute_robust_vcov. See :func:`compute_robust_vcov` for parameter and return semantics. + + ``_bread_matrix`` is a private reuse seam for the opt-in solve_ols + Cholesky fast path: the caller passes its already-built ``X.T @ X`` + (the exact same expression on the same array, so the bread is + byte-identical to building it here). Honored only for unweighted calls; + every vcov formula downstream is unchanged. """ # Re-run the shared validation here too. The public wrapper validates # before dispatch, but solve_ols / _solve_ols_numpy call this function @@ -2905,6 +3251,8 @@ def _compute_robust_vcov_numpy( if weights is not None: XtWX = X.T @ (X * weights[:, np.newaxis]) bread_matrix = XtWX + elif _bread_matrix is not None: + bread_matrix = _bread_matrix else: bread_matrix = X.T @ X diff --git a/docs/doc-deps.yaml b/docs/doc-deps.yaml index b6dba61f..51b8b854 100644 --- a/docs/doc-deps.yaml +++ b/docs/doc-deps.yaml @@ -880,6 +880,7 @@ sources: type: methodology - path: docs/performance-plan.md type: performance + note: "Opt-in DIFF_DIFF_SOLVE_OLS_FASTPATH normal-equations Cholesky fast path (certified twin + stage-0 Gram reuse)" - path: docs/benchmarks.rst type: performance - path: docs/api/utils.rst @@ -1021,4 +1022,4 @@ sources: docs: - path: docs/performance-plan.md type: performance - note: "Rust backend detection" + note: "Rust backend detection; independently-imported solve_ols_chol symbol (stale-extension degradation)" diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index f9ec0121..dd45ced7 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -686,6 +686,7 @@ The multiplier bootstrap uses random weights w_i with E[w]=0 and Var(w)=1: - Handling: Warns and drops linearly dependent columns, sets NA for dropped coefficients (R-style, matches `lm()`) - Parameter: `rank_deficient_action` controls behavior: "warn" (default), "error", or "silent" - **Note:** Rank detection and the least-squares solve are invariant to per-column scaling in BOTH the Python and Rust backends. The rank threshold is anchored to the largest pivot/singular value, so a column on a large raw scale (e.g. an unstandardized covariate in raw population or currency units) previously inflated the threshold and false-dropped the intercept/treatment/interaction to NaN on an otherwise full-rank design — or truncated the small-scale direction in the solve, returning finite-but-wrong coefficients. Detection now runs a raw pivoted QR first and only re-checks on column-equilibrated (unit 2-norm) columns when the raw pass reports a deficiency. If the equilibrated rank is higher (a scale-induced false-drop), the equilibrated rank AND its pivot selection are adopted — its first `rank` columns are independent under the scale-corrected criterion, so the retained design is guaranteed identified; otherwise (genuine collinearity, no scale disparity) the raw rank and pivot selection are kept unchanged. The solve equilibrates columns and unscales the coefficients. The Python (LAPACK `gelsd`) and Rust (faer thin-SVD) backends implement this same contract and agree at the parity-suite tolerances (coefficients ~1e-8, vcov ~1e-5), never bit-identically. This repairs the scale bug while leaving everything else unchanged: it is a no-op for full-rank well-conditioned designs (R-parity unaffected) and does **not** change which column is dropped in a *well-scaled* collinear design (the established raw pivot selection is preserved). A scale-induced under-count instead adopts the scale-corrected equilibrated selection — which may differ from the raw choice in a mixed scale+collinearity design but is guaranteed to retain an identified (full-rank) subset. This shared `diff_diff/linalg.py` behavior covers every covariate outcome-regression fit routed through `solve_ols` — DiD, TwoWayFixedEffects, MultiPeriodDiD, ImputationDiD, TwoStageDiD, and TripleDifference. **Scope (now scale-equilibrated):** CallawaySantAnna's covariate outcome-regression (`_compute_all_att_gt_covariate_reg`) and doubly-robust (`_doubly_robust`) nuisance solves — and StaggeredTripleDifference's per-cohort OR solve (`_compute_or`) — now route through `solve_ols` (column-equilibrated SVD/gelsd), matching TripleDifference's already-`solve_ols`-routed OR fit (`triple_diff.py:1438/1444`) and R's `lm()`/QR. The prior estimator-local `cho_solve(X'X)` / `scipy.linalg.lstsq(cond=1e-7)` fast paths were not scale-equilibrated: a covariate **correlated with another regressor at a very large scale** (e.g. a large constant offset, near-collinear with the intercept) could perturb the point-estimate ATT — and the IF SE that follows it — because the normal-equations Cholesky squares the condition number (pure orthogonal ill-scaling was already safe). The equilibrated SVD is offset-invariant to ~1e-11 where the prior solve drifted; scale-invariance is now pinned by tests (covariate + 1e6 offset → ATT(g,t) unchanged). The change is not bit-identical (cho/normal-equations → SVD) but well-scaled designs move only ~1e-12. The separate DR/OR influence-function SE rank-guard — which previously returned garbage SEs (~1e13) when these local Gram matrices were near-singular — is also **implemented** for CS / TripleDifference / StaggeredTripleDifference via `_rank_guarded_inv` (see the "rank-guarded IF standard errors" Note above). + - **Note (opt-in solve_ols Cholesky fast path, 2026-07):** setting `DIFF_DIFF_SOLVE_OLS_FASTPATH=1` (positive integer; read per call) routes full-rank `solve_ols` fits through an equilibrated normal-equations Cholesky solve, certification-gated at reciprocal condition 1e-6 on the equilibrated Gram (Python: LAPACK `dpocon`, the same guard shape as the IRLS inner-solver Note under Propensity score estimation; Rust: the exact 1-norm rcond, whose inverse byproduct is reused as the sandwich bread). Any certification decline falls back VERBATIM to the SVD/gelsd solve, so the knob-off default path - and every rank-deficiency decision - stays byte-identical to the legacy behavior on both backends. The fast path never runs on a design the QR detector would call deficient: the 1e-6 solve guard is strictly inside the 1e-10 stage-0 rank certification, which is itself two orders stricter than the 1e-7 QR boundary. Within the opt-in path, parity vs the default solve is tol-bounded, not bit-level (fitted values ~1e-8 abs, SEs ~1e-6 rel; certified forward-error budget ~eps*cond(G_eq) <= 2e-10 in the equilibrated basis). The default path keeps the SVD's independent singular-value truncation as defense-in-depth; the knob is opt-in per the 2026-07 correctness-first decision. - Non-finite inference values: - Analytic SE: Returns NaN to signal invalid inference (not biased via zeroing) - Bootstrap: Drops non-finite samples, warns, and adjusts p-value floor accordingly. SE, CI, and p-value are all NaN if the original point estimate is non-finite, SE is non-finite or zero (e.g., n_valid=1 with ddof=1, or identical samples) diff --git a/docs/performance-plan.md b/docs/performance-plan.md index 43a11a40..a870c0c0 100644 --- a/docs/performance-plan.md +++ b/docs/performance-plan.md @@ -4,6 +4,76 @@ This document outlines the strategy for improving diff-diff's performance on lar --- +## Opt-in solve_ols normal-equations Cholesky fast path (v3.7, 2026-07) + +`solve_ols` is the universal OLS entry point; both backends run an equilibrated +thin SVD (gelsd-parity) by default — deliberate defense-in-depth (the SVD's +singular-value truncation is a second, independent near-collinearity detector). +After the Phase 3 stage-0 Gram certification removed the pivoted-QR cost, the +SVD solve itself became the dominant stage on solver-bound fits (2026-07-10 +attribution, M4 Max, rust+Accelerate): + +| Scenario | Fit (median) | rust `solve_ols` | share | +|---|---|---|---| +| county (SunAbraham, 177k rows, ~100 cols) | 1.12 s | 0.65 s | 58% | +| firm (SunAbraham, 2.4M rows, ~130 cols) | 17.6 s | 9.07 s | 54% | +| scanner (TWFE, 3.25M rows) | 0.52 s | 0.12 s | 25% | +| cs_cov40 (CS dr, 2M rows, 40 cov, 95 cells) | 4.28 s | 1.66 s | 40% | + +**The lever (opt-in only, per the 2026-07 correctness-first decision):** +`DIFF_DIFF_SOLVE_OLS_FASTPATH=1` (set to a positive integer; read per call so a +process can A/B) routes certified-well-conditioned full-rank solves through an +equilibrated normal-equations Cholesky instead of the SVD: + +- **Python twin**: reuses the stage-0 rank-certification artifacts (the Gram is + already built to certify rank — the marginal solve cost is X'y + a k×k + `cho_factor`/`cho_solve`), gated by LAPACK `dpocon` at reciprocal condition + 1e-6 on the equilibrated Gram (the `_IRLS_CHOL_RCOND_GUARD` budget: forward + error ~eps·cond ≈ 2e-10). Self-builds and self-certifies on the + `skip_rank_check` route (no stage-0 artifacts exist there — factorization + success alone is NOT a certificate). +- **Rust kernel** (`solve_ols_chol`, separate pyfunction so a stale extension + degrades gracefully): self-certifying faer `Llt` with the exact 1-norm rcond, + whose inverse byproduct doubles as the sandwich bread; one owned n×k buffer + (the equilibrated copy, reused in place for the vcov scores — no thin-SVD U + transient); faer-only heavy ops, so the accelerate/openblas/no-BLAS wheel + variants behave identically. +- **Fallback contract**: any certification decline falls through the UNCHANGED + legacy chain (Rust SVD, then gelsd), so knob-on-decline output equals + knob-off exactly; the knob-off default is byte-identical to the pre-change + tree (pinned by a pristine-base-tree identity gate in the benchmark lane and + spy + exact-equality tests) with one deliberate exception: the saturated + n == k silent-Inf vcov leak on the legacy Rust path was fixed alongside + (CHANGELOG Fixed entry; the sentinel changes Inf -> NaN + warning — no + benchmark scenario is saturated, so the identity gate is unaffected). + +**Measured** (`benchmarks/speed_review/bench_solve_ols_fastpath.py`, +subprocess-per-scenario, medians; deltas are knob-on vs knob-off): + +| Scenario | off | on | speedup | max deltas | +|---|---|---|---|---| +| county_policy (SunAbraham) | 1.13 s | 0.63 s | **1.80x** | dATT 5e-16, dSE/SE 9e-14 | +| firm_churn (SunAbraham) | 16.77 s | 9.89 s | **1.70x** | dATT 9e-15, dSE/SE 3e-12 | +| scanner_twfe (TWFE) | 0.58 s | 0.57 s | 1.01x | demean-bound | +| cs_cov40 (CS dr per-cell) | 4.40 s | 3.57 s | **1.23x** | dATT = 0 (bit-stable) | +| survey_absorb (WLS -> numpy twin) | 2.77 s | 2.59 s | 1.07x | dATT = 0 | +| skiprank_micro (2.4M×130 direct) | 6.88 s | 4.42 s | **1.56x** | dATT 1e-12 | + +Memory: rust-side allocator high-water on the 2.4M×130 clustered solve +(alloc-profile build) drops **7.27 GB -> 2.66 GB** knob-on — the fast path's +working set is essentially the single equilibrated buffer, sidestepping the +SVD residual memory floor documented below (that floor still governs the +default path). + +**Why opt-in and not default**: the default path's byte-stability is load-bearing +(benchmark identity conventions, golden pins at 1e-8) and the SVD truncation is +deliberate defense-in-depth. Within the opt-in path, parity is tol-bounded +(fitted ~1e-8 abs, SE ~1e-6 rel), never bit-level. A default-on flip is tracked +as a TODO follow-up (needs golden recapture + certification-rate telemetry + +the staged default-flip protocol). + +--- + ## FE-absorption baseline: the MAP demeaning hot path (v3.6.x, July 2026) A measured head-to-head against pyfixest 0.60 (Rust demeaner core), prompted by diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 10f25dbd..7ad95f9a 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -52,6 +52,7 @@ fn _rust_backend(m: &Bound<'_, PyModule>) -> PyResult<()> { // Linear algebra operations m.add_function(wrap_pyfunction!(linalg::solve_ols, m)?)?; + m.add_function(wrap_pyfunction!(linalg::solve_ols_chol, m)?)?; m.add_function(wrap_pyfunction!(linalg::compute_robust_vcov, m)?)?; m.add_function(wrap_pyfunction!(linalg::compute_robust_vcov_hc2, m)?)?; diff --git a/rust/src/linalg.rs b/rust/src/linalg.rs index 246b2ab2..660c6af7 100644 --- a/rust/src/linalg.rs +++ b/rust/src/linalg.rs @@ -12,6 +12,7 @@ use std::collections::HashMap; // faer for pure Rust linear algebra (no external BLAS/LAPACK dependencies) use faer::linalg::solvers::{PartialPivLu, Solve}; +use faer::Side; /// Solve OLS regression: β = (X'X)^{-1} X'y /// @@ -312,6 +313,32 @@ fn compute_robust_vcov_internal( n: usize, k: usize, ) -> PyResult> { + // Saturated design (n <= k): zero residual degrees of freedom make the + // HC1/CR1 (n - k) adjustment undefined. The documented non-finite + // inference contract — and the numpy backend, which is canonical (its + // saturated guard fires before any bread inversion) — is an all-NaN + // vcov; the previous behavior silently leaked Inf through the n/(n-k) + // division. Placed BEFORE the bread inversion to mirror the numpy + // guard order, so direct compute_robust_vcov calls with a singular + // saturated Gram get the NaN sentinel rather than an inversion error. + // The G >= 2 cluster contract keeps precedence (evaluated here only in + // the saturated case, where n is tiny; the general path's in-arm check + // below is unchanged). + if n <= k { + if let Some(clusters) = cluster_ids { + let (_, n_clusters) = cluster_first_appearance_index(clusters, n); + if n_clusters < 2 { + return Err(PyErr::new::(format!( + "Need at least 2 clusters for cluster-robust SEs, got {}", + n_clusters + ))); + } + } + let mut nan_vcov = Array2::::zeros((k, k)); + nan_vcov.fill(f64::NAN); + return Ok(nan_vcov); + } + // Compute X'X let xtx = x.t().dot(x); @@ -358,13 +385,10 @@ fn compute_robust_vcov_internal( // is bit-stable). Rows accumulate in first-appearance order // instead: for the factorized 0..G-1 ids the Python dispatcher // passes this is ascending id order, matching NumPy's groupby. - let mut cluster_index: HashMap = HashMap::new(); - for i in 0..n_obs { - let next = cluster_index.len(); - cluster_index.entry(clusters[i]).or_insert(next); - } - - let n_clusters = cluster_index.len(); + // Index construction is shared with the Cholesky fast-path + // kernel (cluster_first_appearance_index) so the ordering + // contract cannot drift between the two vcov paths. + let (cluster_index, n_clusters) = cluster_first_appearance_index(clusters, n_obs); if n_clusters < 2 { return Err(PyErr::new::(format!( @@ -510,6 +534,321 @@ fn invert_symmetric(a: &Array2) -> PyResult> { Ok(result) } +/// Build the deterministic first-appearance-order cluster index shared by +/// the SVD-path vcov (`compute_robust_vcov_internal`) and the Cholesky +/// fast-path kernel. Returns the id -> row map and the cluster count; the +/// caller enforces its own G >= 2 contract so the two paths raise the same +/// error through their own channels. +fn cluster_first_appearance_index( + clusters: &Array1, + n_obs: usize, +) -> (HashMap, usize) { + let mut cluster_index: HashMap = HashMap::new(); + for i in 0..n_obs { + let next = cluster_index.len(); + cluster_index.entry(clusters[i]).or_insert(next); + } + let n_clusters = cluster_index.len(); + (cluster_index, n_clusters) +} + +/// Kernel error that must surface as a Python exception, as opposed to a +/// certification decline (which returns `None` so the dispatcher falls +/// through to the SVD path). +pub(crate) enum CholKernelError { + /// Cluster-robust vcov requested with fewer than 2 clusters. Message + /// parity with `compute_robust_vcov_internal` matters: the Python + /// dispatcher re-raises unless the text mentions instability. + TooFewClusters(usize), +} + +/// Reciprocal-1-norm-condition guard for the Cholesky fast path. Same 1e-6 +/// bound as the Python twin's `_SOLVE_OLS_CHOL_RCOND_GUARD` (dpocon): +/// cond(G_eq) <= 1e6 bounds the Cholesky forward error at ~eps*cond +/// ~ 2e-10 relative in the equilibrated basis. The exact 1-norm computed +/// here is >= dpocon's lower-bound estimate, so this rcond is <= the +/// twin's and the guard is marginally STRICTER — a disagreement near the +/// boundary only flips to the verbatim-correct SVD fallback. +const CHOL_RCOND_GUARD: f64 = 1e-6; + +/// Output of the certified normal-equations Cholesky solve. +pub(crate) struct CholSolveOutput { + pub(crate) coefficients: Array1, + pub(crate) residuals: Array1, + pub(crate) vcov: Option>, +} + +/// Certified equilibrated normal-equations Cholesky OLS (opt-in fast path). +/// +/// Self-certifying: returns `Ok(None)` (the dispatcher then falls through +/// VERBATIM to the SVD `solve_ols`) on any of: k == 0, n < k, zero or +/// non-finite column norm, Llt factorization failure, exact 1-norm +/// rcond(G_eq) <= 1e-6, or a non-finite coefficient. Every heavy operation +/// on this path is faer matmul or a hand-rolled sequential loop — never a +/// BLAS-feature-gated ndarray `.dot()` — so results and performance are +/// uniform across the accelerate/openblas/no-BLAS wheel variants. +/// +/// Memory discipline: ONE owned n x k buffer (the equilibrated copy), +/// reused in place for the vcov scores; unlike the SVD path there is no +/// U-factor transient, so the allocator high-water is strictly lower. +/// +/// vcov (hc1 / CR1 only, matching the SVD kernel's scope) fuses the bread +/// from the certification byproduct: the rcond computation solves +/// G_eq * Z = I, and (X'X)^{-1} = D^{-1} Z D^{-1} with D = diag(col norms). +/// Adjustment factors, the G >= 2 contract, and the saturated n == k +/// all-NaN vcov contract replicate `compute_robust_vcov_internal` +/// (both paths honor the documented non-finite-inference contract; +/// the SVD path's former silent-Inf leak was fixed alongside this +/// kernel). +pub(crate) fn solve_ols_chol_core( + x: &ArrayView2, + y: &ArrayView1, + cluster_ids: Option<&Array1>, + return_vcov: bool, +) -> Result, CholKernelError> { + let n = x.nrows(); + let k = x.ncols(); + if k == 0 || n < k { + return Ok(None); + } + + // Column 2-norms in the legacy j-outer/i-inner order. Unlike the SVD + // path (which substitutes 1.0 and lets truncation absorb zero columns), + // a zero norm here means a structurally deficient design and a + // non-finite norm means NaN/Inf in X (any non-finite entry poisons its + // own column's norm) — both are certification declines. + let mut norms = Array1::::zeros(k); + for j in 0..k { + let mut acc = 0.0_f64; + for i in 0..n { + let v = x[[i, j]]; + acc += v * v; + } + let norm = acc.sqrt(); + if !norm.is_finite() || !(norm > 0.0) { + return Ok(None); + } + norms[j] = norm; + } + + // The single owned n x k buffer: equilibrated copy (unit 2-norm + // columns, entries bounded by 1, so the Gram below is finite by + // construction given the finite-norm gate above). + let mut x_eq = faer::Mat::from_fn(n, k, |i, j| x[[i, j]] / norms[j]); + + // Equilibrated Gram and its Cholesky factor. Factorization failure + // (not numerically PD) declines. + let gram_eq = x_eq.as_ref().transpose() * x_eq.as_ref(); + let llt = match gram_eq.llt(Side::Lower) { + Ok(f) => f, + Err(_) => return Ok(None), + }; + + // Certification: exact 1-norm reciprocal condition number via the full + // inverse Z = G_eq^{-1} (a k x k solve against I — same O(k^3) order as + // the factorization, microseconds at the k this library sees). Z + // doubles as the vcov bread below, so certification is not wasted work. + let identity = faer::Mat::::from_fn(k, k, |i, j| if i == j { 1.0 } else { 0.0 }); + let z = llt.solve(&identity); + let mut anorm = 0.0_f64; // ||G_eq||_1 + let mut znorm = 0.0_f64; // ||G_eq^{-1}||_1 + for j in 0..k { + let mut col_a = 0.0_f64; + let mut col_z = 0.0_f64; + for i in 0..k { + col_a += gram_eq[(i, j)].abs(); + col_z += z[(i, j)].abs(); + } + anorm = anorm.max(col_a); + znorm = znorm.max(col_z); + } + let denom = anorm * znorm; + let rcond = if denom > 0.0 && denom.is_finite() { + 1.0 / denom + } else { + 0.0 + }; + if !(rcond > CHOL_RCOND_GUARD) { + return Ok(None); + } + + // beta_eq = G_eq^{-1} (X_eq' y); X_eq'y via sequential column-order + // accumulation off the col-major faer buffer (deterministic bits, like + // the SVD path's uty loop). + let mut xty_eq = faer::Mat::::zeros(k, 1); + for j in 0..k { + let mut acc = 0.0_f64; + for i in 0..n { + acc += x_eq[(i, j)] * y[i]; + } + xty_eq[(j, 0)] = acc; + } + let beta_eq = llt.solve(&xty_eq); + let mut coefficients = Array1::::zeros(k); + for j in 0..k { + let b = beta_eq[(j, 0)] / norms[j]; + if !b.is_finite() { + return Ok(None); + } + coefficients[j] = b; + } + + // fitted = X beta = X_eq beta_eq exactly (equilibration is a column + // reparameterization), so the matvec reuses the owned equilibrated + // buffer through faer instead of a BLAS-gated dot on the raw view. + let fitted = x_eq.as_ref() * beta_eq.as_ref(); + let mut residuals = Array1::::zeros(n); + for i in 0..n { + residuals[i] = y[i] - fitted[(i, 0)]; + } + + let vcov = if return_vcov { + // Saturated design (n == k; n < k already declined above): zero + // residual degrees of freedom, so the sandwich's (n - k) adjustment + // is undefined. The documented non-finite-inference contract is an + // all-NaN vcov (matching the numpy saturated guard and the SVD + // path's guard in compute_robust_vcov_internal). + // Cluster-contract precedence: the G >= 2 check fires first. + if let Some(clusters) = cluster_ids { + let (_, n_clusters) = cluster_first_appearance_index(clusters, n); + if n_clusters < 2 { + return Err(CholKernelError::TooFewClusters(n_clusters)); + } + } + if n == k { + let mut nan_vcov = Array2::::zeros((k, k)); + nan_vcov.fill(f64::NAN); + return Ok(Some(CholSolveOutput { + coefficients, + residuals, + vcov: Some(nan_vcov), + })); + } + + // Overwrite the equilibrated buffer in place with the RAW scores + // x_ij * u_i (undo the column scaling while folding in the + // residual) — the memory-floor trick: no second n x k allocation. + for j in 0..k { + let nj = norms[j]; + for i in 0..n { + x_eq[(i, j)] *= nj * residuals[i]; + } + } + + let (meat, adjustment) = match cluster_ids { + None => { + // HC1 meat: S'S with S = diag(u) X. + let m = x_eq.as_ref().transpose() * x_eq.as_ref(); + (m, n as f64 / (n - k) as f64) + } + Some(clusters) => { + let (cluster_index, n_clusters) = cluster_first_appearance_index(clusters, n); + // Unreachable when n == k (early return above), but the + // G >= 2 contract is kept here too for the general path. + if n_clusters < 2 { + return Err(CholKernelError::TooFewClusters(n_clusters)); + } + let mut cluster_scores = faer::Mat::::zeros(n_clusters, k); + for i in 0..n { + let idx = cluster_index[&clusters[i]]; + for j in 0..k { + cluster_scores[(idx, j)] += x_eq[(i, j)]; + } + } + let g = n_clusters as f64; + let adjustment = (g / (g - 1.0)) * ((n - 1) as f64 / (n - k) as f64); + let m = cluster_scores.as_ref().transpose() * cluster_scores.as_ref(); + (m, adjustment) + } + }; + + // Bread on the raw scale from the certification byproduct: + // (X'X)^{-1} = D^{-1} G_eq^{-1} D^{-1}. + let bread = faer::Mat::::from_fn(k, k, |i, j| z[(i, j)] / (norms[i] * norms[j])); + let half = bread.as_ref() * meat.as_ref(); + let sandwich = half.as_ref() * bread.as_ref(); + let mut vcov_arr = Array2::::zeros((k, k)); + for i in 0..k { + for j in 0..k { + vcov_arr[[i, j]] = sandwich[(i, j)] * adjustment; + } + } + Some(vcov_arr) + } else { + None + }; + + Ok(Some(CholSolveOutput { + coefficients, + residuals, + vcov, + })) +} + +/// Opt-in certified normal-equations Cholesky OLS (see +/// `solve_ols_chol_core`). Returns Python `None` when certification +/// declines — the dispatcher then falls through verbatim to the SVD +/// `solve_ols` — and the same "Need at least 2 clusters" ValueError as the +/// SVD path's vcov when clustered vcov is requested with G < 2. +/// +/// Registered as a SEPARATE pyfunction (not a new parameter on +/// `solve_ols`) so a stale installed extension missing this symbol +/// degrades gracefully via the independent-import pattern in +/// `_backend.py` instead of raising TypeError at call time. +#[pyfunction] +#[pyo3(signature = (x, y, cluster_ids=None, return_vcov=true))] +#[allow(clippy::type_complexity)] +pub fn solve_ols_chol<'py>( + py: Python<'py>, + x: PyReadonlyArray2<'py, f64>, + y: PyReadonlyArray1<'py, f64>, + cluster_ids: Option>, + return_vcov: bool, +) -> PyResult< + Option<( + Bound<'py, PyArray1>, + Bound<'py, PyArray1>, + Option>>, + )>, +> { + let x_arr = x.as_array(); + let y_arr = y.as_array(); + // Shape contract up front: the Python dispatcher always passes matched + // shapes, but direct pyfunction misuse should raise a clean ValueError + // instead of a low-level index panic from the core. + if y_arr.len() != x_arr.nrows() { + return Err(PyErr::new::(format!( + "y length ({}) must match X rows ({})", + y_arr.len(), + x_arr.nrows() + ))); + } + let cluster_arr = cluster_ids.as_ref().map(|c| c.as_array().to_owned()); + if let Some(cl) = cluster_arr.as_ref() { + if cl.len() != x_arr.nrows() { + return Err(PyErr::new::(format!( + "cluster_ids length ({}) must match X rows ({})", + cl.len(), + x_arr.nrows() + ))); + } + } + match solve_ols_chol_core(&x_arr, &y_arr, cluster_arr.as_ref(), return_vcov) { + Ok(Some(out)) => Ok(Some(( + out.coefficients.to_pyarray(py), + out.residuals.to_pyarray(py), + out.vcov.map(|v| v.to_pyarray(py)), + ))), + Ok(None) => Ok(None), + Err(CholKernelError::TooFewClusters(g)) => { + Err(PyErr::new::(format!( + "Need at least 2 clusters for cluster-robust SEs, got {}", + g + ))) + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -570,6 +909,221 @@ mod tests { // not k, otherwise vt.t().dot(&s_inv_uty) will have mismatched dimensions } + /// Deterministic pseudo-random design for the chol-kernel tests (no + /// rand dependency in tests; simple LCG, values in roughly [-1, 1]). + fn lcg_design(n: usize, k: usize, seed: u64) -> (Array2, Array1) { + let mut state = seed + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + let mut next = || { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((state >> 33) as f64) / ((1u64 << 31) as f64) - 1.0 + }; + let mut x = Array2::::zeros((n, k)); + for i in 0..n { + x[[i, 0]] = 1.0; // intercept + for j in 1..k { + x[[i, j]] = next(); + } + } + let mut y = Array1::::zeros(n); + for i in 0..n { + let mut acc = 0.0; + for j in 0..k { + acc += x[[i, j]] * (j as f64 + 0.5); + } + y[i] = acc + 0.1 * next(); + } + (x, y) + } + + #[test] + fn test_chol_core_matches_lu_oracle() { + // Well-conditioned design: chol beta must match the LU + // normal-equations oracle to ~1e-12 (both exact to ~eps*cond). + let (x, y) = lcg_design(200, 5, 7); + let out = solve_ols_chol_core(&x.view(), &y.view(), None, true) + .ok() + .flatten() + .expect("well-conditioned design must certify"); + + let xtx = x.t().dot(&x); + let xty = x.t().dot(&y); + let xtx_inv = invert_symmetric(&xtx).expect("oracle inverse"); + let beta_oracle = xtx_inv.dot(&xty); + for j in 0..5 { + assert!( + (out.coefficients[j] - beta_oracle[j]).abs() < 1e-12, + "beta[{}]: chol {} vs oracle {}", + j, + out.coefficients[j], + beta_oracle[j] + ); + } + + // Bread-from-factor must match the LU-inverse HC1 sandwich to 1e-12. + let u2: Array1 = out.residuals.mapv(|r| r * r); + let xw = &x * &u2.insert_axis(ndarray::Axis(1)); + let meat = x.t().dot(&xw); + let adj = 200.0 / (200.0 - 5.0); + let vcov_oracle = xtx_inv.dot(&meat).dot(&xtx_inv) * adj; + let vcov = out.vcov.expect("vcov requested"); + for i in 0..5 { + for j in 0..5 { + assert!( + (vcov[[i, j]] - vcov_oracle[[i, j]]).abs() < 1e-12, + "vcov[{},{}]: {} vs {}", + i, + j, + vcov[[i, j]], + vcov_oracle[[i, j]] + ); + } + } + } + + #[test] + fn test_chol_core_declines_near_singular() { + // x2 = x1 + 1e-8 * noise: cond(G_eq) far beyond the 1e-6 guard. + let (x0, y) = lcg_design(300, 2, 11); + let mut x = Array2::::zeros((300, 3)); + let mut state = 12345u64; + for i in 0..300 { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + let noise = ((state >> 33) as f64) / ((1u64 << 31) as f64) - 1.0; + x[[i, 0]] = x0[[i, 0]]; + x[[i, 1]] = x0[[i, 1]]; + x[[i, 2]] = x0[[i, 1]] + 1e-8 * noise; + } + let result = solve_ols_chol_core(&x.view(), &y.view(), None, true) + .ok() + .flatten(); + assert!(result.is_none(), "near-singular design must decline"); + } + + #[test] + fn test_chol_core_declines_on_nan_and_zero_column() { + let (mut x, y) = lcg_design(50, 3, 13); + x[[7, 1]] = f64::NAN; + assert!( + solve_ols_chol_core(&x.view(), &y.view(), None, false) + .ok() + .flatten() + .is_none(), + "NaN in X poisons its column norm and must decline" + ); + + let (mut x2, y2) = lcg_design(50, 3, 17); + for i in 0..50 { + x2[[i, 2]] = 0.0; + } + assert!( + solve_ols_chol_core(&x2.view(), &y2.view(), None, false) + .ok() + .flatten() + .is_none(), + "zero column must decline (structural deficiency)" + ); + } + + #[test] + fn test_chol_core_declines_underdetermined() { + let (x, y) = lcg_design(3, 5, 19); + assert!( + solve_ols_chol_core(&x.view(), &y.view(), None, false) + .ok() + .flatten() + .is_none(), + "n < k must decline" + ); + } + + #[test] + fn test_chol_core_saturated_design_nan_vcov() { + // n == k full rank: zero residual df makes the sandwich adjustment + // undefined; the contract is all-NaN vcov (never Inf). Coefficients + // and residuals stay finite. Cluster precedence: G < 2 still errors + // BEFORE the saturated guard. + let (x, y) = lcg_design(4, 4, 31); + let out = solve_ols_chol_core(&x.view(), &y.view(), None, true) + .ok() + .flatten() + .expect("saturated full-rank design must certify"); + assert!(out.coefficients.iter().all(|v| v.is_finite())); + let vcov = out.vcov.expect("vcov requested"); + assert!( + vcov.iter().all(|v| v.is_nan()), + "saturated vcov must be all-NaN, got {:?}", + vcov + ); + + let mut clusters = Array1::::zeros(4); + for i in 0..4 { + clusters[i] = (i % 2) as i64; + } + let out2 = solve_ols_chol_core(&x.view(), &y.view(), Some(&clusters), true) + .ok() + .flatten() + .expect("clustered saturated design must certify"); + assert!(out2.vcov.expect("vcov").iter().all(|v| v.is_nan())); + + let one_cluster = Array1::::zeros(4); + let res = solve_ols_chol_core(&x.view(), &y.view(), Some(&one_cluster), true); + assert!( + matches!(res, Err(CholKernelError::TooFewClusters(1))), + "G < 2 must take precedence over the saturated guard" + ); + } + + #[test] + fn test_chol_core_too_few_clusters_errors() { + let (x, y) = lcg_design(60, 3, 23); + let clusters = Array1::::zeros(60); // single cluster + let result = solve_ols_chol_core(&x.view(), &y.view(), Some(&clusters), true); + assert!( + matches!(result, Err(CholKernelError::TooFewClusters(1))), + "G < 2 with vcov must error, not decline" + ); + // ...but with return_vcov=false the cluster contract never fires, + // matching the SVD path (its check lives inside the vcov build). + let no_vcov = solve_ols_chol_core(&x.view(), &y.view(), Some(&clusters), false); + assert!(matches!(no_vcov, Ok(Some(_)))); + } + + #[test] + fn test_chol_core_clustered_matches_lu_oracle() { + let (x, y) = lcg_design(120, 4, 29); + let mut clusters = Array1::::zeros(120); + for i in 0..120 { + clusters[i] = (i % 8) as i64; + } + let out = solve_ols_chol_core(&x.view(), &y.view(), Some(&clusters), true) + .ok() + .flatten() + .expect("well-conditioned design must certify"); + let vcov = out.vcov.expect("vcov requested"); + + let oracle = + compute_robust_vcov_internal(&x.view(), &out.residuals.view(), Some(&clusters), 120, 4) + .expect("oracle vcov"); + for i in 0..4 { + for j in 0..4 { + assert!( + (vcov[[i, j]] - oracle[[i, j]]).abs() < 1e-12, + "clustered vcov[{},{}]: {} vs {}", + i, + j, + vcov[[i, j]], + oracle[[i, j]] + ); + } + } + } + // Note: Singular and near-singular matrix tests removed because: // 1. invert_symmetric() returns PyResult, which requires Python initialization // to create PyErr - `cargo test` without Python causes panic diff --git a/tests/test_linalg.py b/tests/test_linalg.py index f91738ac..68a16ef5 100644 --- a/tests/test_linalg.py +++ b/tests/test_linalg.py @@ -3026,3 +3026,340 @@ def test_contrast_chunking_invariant(self, monkeypatch): monkeypatch.setattr(la, "_CR2_BM_CONTRAST_CHUNK_BYTES", G * k * 8) # 1 contrast/chunk _, dof_many = la._compute_cr2_bm(X, resid, cl, bread) np.testing.assert_allclose(dof_many, dof_one, rtol=1e-13) + + +class TestSolveOLSFastpathResolver: + """DIFF_DIFF_SOLVE_OLS_FASTPATH env resolver conventions. + + Mirrors TestDemeanChunkResolver: unset -> module default (OFF), positive + integer enables, invalid values fall back silently to the module default. + """ + + def test_default_when_unset_is_off(self, monkeypatch): + import diff_diff.linalg as lmod + + monkeypatch.delenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", raising=False) + assert lmod._SOLVE_OLS_FASTPATH is False + assert lmod._resolve_solve_ols_fastpath() is False + + @pytest.mark.parametrize("val", ["1", "2", "10"]) + def test_valid_override_honored(self, monkeypatch, val): + import diff_diff.linalg as lmod + + monkeypatch.setenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", val) + assert lmod._resolve_solve_ols_fastpath() is True + + @pytest.mark.parametrize("bad", ["abc", "0", "-4", "", "3.5"]) + def test_invalid_values_fall_back_to_default(self, monkeypatch, bad): + import diff_diff.linalg as lmod + + monkeypatch.setenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", bad) + assert lmod._resolve_solve_ols_fastpath() is False + + def test_module_default_is_monkeypatch_seam(self, monkeypatch): + import diff_diff.linalg as lmod + + monkeypatch.delenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", raising=False) + monkeypatch.setattr(lmod, "_SOLVE_OLS_FASTPATH", True) + assert lmod._resolve_solve_ols_fastpath() is True + + +class TestSolveOLSCholFastPath: + """Opt-in normal-equations Cholesky fast path (numpy twin). + + Parity posture is TOL-BOUNDED, not bit-level: the fast path swaps the + solve algorithm (certified Cholesky on the equilibrated Gram vs gelsd + SVD), so outputs differ at the numerical-error level. The dpocon guard + rcond > _SOLVE_OLS_CHOL_RCOND_GUARD (1e-6) bounds cond(G_eq) < 1e6 and + hence the Cholesky forward error at ~eps*cond ~ 2e-10 relative in the + equilibrated basis (same budget as _IRLS_CHOL_RCOND_GUARD); the gelsd + reference carries its own ~eps*cond(X_eq) error. Parity fixtures use + unit-scale y and designs a decade or more inside the guard, so the + fitted rtol=1e-6/atol=1e-8 and SE rtol=1e-6 gates have >= 2 orders of + headroom. SEs are compared RELATIVELY via sqrt(diag(vcov)) — never + absolute-decimal — because coefficient scales vary across fixtures. + + Byte-identity IS asserted wherever the gelsd path actually ran (forced + or natural certification decline, knob off, rank-deficient designs): + a decline falls through the verbatim legacy lines. + """ + + @staticmethod + def _force_python(monkeypatch): + import diff_diff.linalg as lmod + + monkeypatch.setattr(lmod, "HAS_RUST_BACKEND", False) + monkeypatch.setattr(lmod, "_rust_solve_ols", None) + + @staticmethod + def _make_design(seed, n=400, k=5, scale_col=False, add_cluster=False): + rng = np.random.default_rng(seed) + X = np.column_stack([np.ones(n), rng.standard_normal((n, k - 1))]) + if scale_col: + X[:, k - 1] *= 1e8 # equilibration absorbs raw scale + beta = rng.standard_normal(k) / np.maximum(1.0, np.abs(X).max(axis=0) / 10.0) + y = X @ beta + rng.standard_normal(n) + cluster_ids = rng.integers(0, 25, n) if add_cluster else None + return X, y, cluster_ids + + def _run_pair(self, monkeypatch, X, y, diag=None, **kwargs): + """Run solve_ols knob-on then knob-off (both forced-python).""" + self._force_python(monkeypatch) + monkeypatch.setenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", "1") + res_on = solve_ols(X, y, diagnostics_out=diag, **kwargs) + monkeypatch.delenv("DIFF_DIFF_SOLVE_OLS_FASTPATH") + res_off = solve_ols(X, y, **kwargs) + return res_on, res_off + + def test_property_parity_vs_default(self, monkeypatch): + cases = [] + for seed in range(10): + cases.append(dict(seed=seed, n=200 + 40 * seed, k=3 + seed % 5)) + for seed in range(10, 15): + cases.append(dict(seed=seed, n=500, k=6, scale_col=True)) + for seed in range(15, 20): + cases.append(dict(seed=seed, n=600, k=4, add_cluster=True)) + + for case in cases: + add_cluster = case.pop("add_cluster", False) + X, y, cl = self._make_design(**case, add_cluster=add_cluster) + diag = {} + (c_on, r_on, v_on), (c_off, r_off, v_off) = self._run_pair( + monkeypatch, X, y, diag=diag, cluster_ids=cl + ) + assert diag["solve_ols_fastpath"] == "chol_numpy", case + np.testing.assert_allclose(c_on, c_off, rtol=0, atol=1e-8, err_msg=str(case)) + np.testing.assert_allclose(X @ c_on, X @ c_off, rtol=1e-6, atol=1e-8, err_msg=str(case)) + np.testing.assert_allclose( + np.sqrt(np.diag(v_on)), np.sqrt(np.diag(v_off)), rtol=1e-6, err_msg=str(case) + ) + + def test_return_fitted_and_no_vcov_shapes(self, monkeypatch): + X, y, _ = self._make_design(31) + on4, off4 = self._run_pair(monkeypatch, X, y, return_fitted=True, return_vcov=False) + assert len(on4) == 4 and len(off4) == 4 + assert on4[3] is None and off4[3] is None + np.testing.assert_allclose(on4[2], off4[2], rtol=1e-6, atol=1e-8) + + def test_forced_fallback_byte_identical(self, monkeypatch): + import diff_diff.linalg as lmod + + X, y, _ = self._make_design(7) + + def _always_raise(*args, **kwargs): + raise np.linalg.LinAlgError("forced") + + self._force_python(monkeypatch) + monkeypatch.setattr(lmod, "cho_factor", _always_raise) + monkeypatch.setenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", "1") + diag = {} + c_on, r_on, v_on = solve_ols(X, y, diagnostics_out=diag) + assert diag["solve_ols_fastpath"] == "fallback_declined" + # cho_factor stays patched for the off run: the gelsd path never + # calls it, which is itself part of the byte-identity claim. + monkeypatch.delenv("DIFF_DIFF_SOLVE_OLS_FASTPATH") + c_off, r_off, v_off = solve_ols(X, y) + np.testing.assert_array_equal(c_on, c_off) + np.testing.assert_array_equal(r_on, r_off) + np.testing.assert_array_equal(v_on, v_off) + + def test_natural_guard_trip_between_fixture(self, monkeypatch): + """Design that PASSES stage-0 rank cert (1e-10) but FAILS the solve + guard (1e-6): the twin must decline and fall through verbatim.""" + import diff_diff.linalg as lmod + + rng = np.random.default_rng(11) + n = 500 + x0 = rng.standard_normal(n) + X = np.column_stack([np.ones(n), x0, x0 + 1e-4 * rng.standard_normal(n)]) + y = X @ np.array([1.0, 2.0, 3.0]) + rng.standard_normal(n) + + # Fixture self-check (regenerate if either bound drifts): eig ratio of + # the equilibrated Gram must sit BETWEEN the two thresholds with a + # decade of margin on each side. + gram = X.T @ X + scales = np.sqrt(np.diag(gram)) + gram_eq = gram / scales[:, None] / scales[None, :] + eigvals = np.linalg.eigvalsh(gram_eq) + ratio = eigvals[0] / eigvals[-1] + assert 1e-9 < ratio < 1e-7, ( + f"BETWEEN fixture drifted (eig ratio {ratio:.2e}); regenerate so it " + "sits between the 1e-10 stage-0 cert and the 1e-6 solve guard" + ) + + # Certified full rank by stage 0 (no QR, no warning)... + rank, dropped, _ = lmod._detect_rank_deficiency(X) + assert rank == X.shape[1] and dropped.size == 0 + + diag = {} + (c_on, r_on, v_on), (c_off, r_off, v_off) = self._run_pair(monkeypatch, X, y, diag=diag) + # ...but the solve guard declines and the gelsd line runs verbatim. + assert diag["solve_ols_fastpath"] == "fallback_declined" + np.testing.assert_array_equal(c_on, c_off) + np.testing.assert_array_equal(r_on, r_off) + np.testing.assert_array_equal(v_on, v_off) + + def test_default_path_byte_identity_and_no_twin(self, monkeypatch): + import diff_diff.linalg as lmod + + X, y, _ = self._make_design(13) + self._force_python(monkeypatch) + monkeypatch.delenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", raising=False) + + calls = {"twin": 0, "gelsd": 0} + real_twin = lmod._solve_ols_chol_numpy + real_lstsq = lmod._equilibrated_lstsq + + def _twin(*a, **kw): + calls["twin"] += 1 + return real_twin(*a, **kw) + + def _lstsq(*a, **kw): + calls["gelsd"] += 1 + return real_lstsq(*a, **kw) + + monkeypatch.setattr(lmod, "_solve_ols_chol_numpy", _twin) + monkeypatch.setattr(lmod, "_equilibrated_lstsq", _lstsq) + + diag = {} + c1, r1, v1 = solve_ols(X, y, diagnostics_out=diag) + assert calls == {"twin": 0, "gelsd": 1} + assert diag["solve_ols_fastpath"] == "off" + + # env set to an invalid value (0) is also OFF and byte-identical + monkeypatch.setenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", "0") + c2, r2, v2 = solve_ols(X, y) + assert calls == {"twin": 0, "gelsd": 2} + np.testing.assert_array_equal(c1, c2) + np.testing.assert_array_equal(r1, r2) + np.testing.assert_array_equal(v1, v2) + + def test_cert_out_is_pure_out_param(self): + import diff_diff.linalg as lmod + + X, _, _ = self._make_design(17) + r1 = lmod._detect_rank_deficiency(X) + cert = {} + r2 = lmod._detect_rank_deficiency(X, _cert_out=cert) + assert r1[0] == r2[0] + np.testing.assert_array_equal(r1[1], r2[1]) + np.testing.assert_array_equal(r1[2], r2[2]) + assert cert.get("certified") is True + np.testing.assert_array_equal(cert["gram"], X.T @ X) + + def test_twin_reuse_equals_self_build(self): + import diff_diff.linalg as lmod + + X, y, _ = self._make_design(19) + cert = {} + lmod._detect_rank_deficiency(X, _cert_out=cert) + reused = lmod._solve_ols_chol_numpy(X, y, cert_info=cert) + built = lmod._solve_ols_chol_numpy(X, y, cert_info=None) + assert reused is not None and built is not None + np.testing.assert_array_equal(reused[0], built[0]) + np.testing.assert_array_equal(reused[1], built[1]) + + def test_twin_declines_on_uncertified_artifacts(self): + import diff_diff.linalg as lmod + + rng = np.random.default_rng(23) + n = 300 + x0 = rng.standard_normal(n) + X = np.column_stack([np.ones(n), x0, 2.0 * x0]) # exactly collinear + y = rng.standard_normal(n) + cert = {} + lmod._detect_rank_deficiency(X, _cert_out=cert) + assert "gram_eq" in cert and cert.get("certified") is not True + assert lmod._solve_ols_chol_numpy(X, y, cert_info=cert) is None + + def test_skip_rank_check_self_cert_declines_near_singular(self, monkeypatch): + """skip_rank_check has no stage-0 cert; the twin must SELF-certify. + cho_factor succeeding is not a certificate: this fixture has + cond(G_eq) ~ 1e12 where Cholesky succeeds with a garbage solution.""" + rng = np.random.default_rng(29) + n = 400 + x1 = rng.standard_normal(n) + X = np.column_stack([np.ones(n), x1, x1 + 1e-6 * rng.standard_normal(n)]) + y = X @ np.array([1.0, 2.0, 3.0]) + rng.standard_normal(n) + + diag = {} + (c_on, r_on, v_on), (c_off, r_off, v_off) = self._run_pair( + monkeypatch, X, y, diag=diag, skip_rank_check=True + ) + assert diag["solve_ols_fastpath"] == "fallback_declined" + np.testing.assert_array_equal(c_on, c_off) + np.testing.assert_array_equal(r_on, r_off) + np.testing.assert_array_equal(v_on, v_off) + + def test_skip_rank_check_well_conditioned_takes_twin(self, monkeypatch): + X, y, _ = self._make_design(37) + diag = {} + (c_on, _, _), (c_off, _, _) = self._run_pair( + monkeypatch, X, y, diag=diag, skip_rank_check=True + ) + assert diag["solve_ols_fastpath"] == "chol_numpy" + np.testing.assert_allclose(c_on, c_off, rtol=0, atol=1e-8) + + @pytest.mark.parametrize("zero_weights", [False, True]) + def test_wls_parity(self, monkeypatch, zero_weights): + rng = np.random.default_rng(41) + X, y, _ = self._make_design(41, n=500) + w = rng.uniform(0.5, 2.0, 500) + if zero_weights: + w[::17] = 0.0 + diag = {} + (c_on, _, v_on), (c_off, _, v_off) = self._run_pair(monkeypatch, X, y, diag=diag, weights=w) + assert diag["solve_ols_fastpath"] == "chol_numpy" + np.testing.assert_allclose(c_on, c_off, rtol=0, atol=1e-8) + np.testing.assert_allclose(np.sqrt(np.diag(v_on)), np.sqrt(np.diag(v_off)), rtol=1e-6) + + def test_rank_deficient_interplay(self, monkeypatch): + """Fast path is structurally skipped on rank-deficient designs: + warning, NaN pattern, and outputs are byte-identical to knob-off.""" + rng = np.random.default_rng(43) + n = 300 + x0 = rng.standard_normal(n) + X = np.column_stack([np.ones(n), x0, 2.0 * x0, rng.standard_normal(n)]) + y = rng.standard_normal(n) + + self._force_python(monkeypatch) + monkeypatch.setenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", "1") + with pytest.warns(UserWarning, match="Rank-deficient"): + c_on, r_on, v_on = solve_ols(X, y) + monkeypatch.delenv("DIFF_DIFF_SOLVE_OLS_FASTPATH") + with pytest.warns(UserWarning, match="Rank-deficient"): + c_off, r_off, v_off = solve_ols(X, y) + np.testing.assert_array_equal(np.isnan(c_on), np.isnan(c_off)) + np.testing.assert_array_equal(c_on, c_off) + np.testing.assert_array_equal(v_on, v_off) + + @pytest.mark.parametrize("vcov_type", ["classical", "hc2"]) + def test_non_hc1_vcov_spot_checks(self, monkeypatch, vcov_type): + """The twin changes beta for ALL vcov types reaching the full-rank + branch; spot-check the parity claim beyond the hc1 gate.""" + X, y, _ = self._make_design(47) + diag = {} + (c_on, _, v_on), (c_off, _, v_off) = self._run_pair( + monkeypatch, X, y, diag=diag, vcov_type=vcov_type + ) + assert diag["solve_ols_fastpath"] == "chol_numpy" + np.testing.assert_allclose(c_on, c_off, rtol=0, atol=1e-8) + np.testing.assert_allclose(np.sqrt(np.diag(v_on)), np.sqrt(np.diag(v_off)), rtol=1e-6) + + def test_conley_vcov_spot_check(self, monkeypatch): + rng = np.random.default_rng(53) + n = 120 + X = np.column_stack([np.ones(n), rng.standard_normal((n, 2))]) + y = X @ np.array([1.0, 0.5, -0.5]) + rng.standard_normal(n) + coords = rng.uniform(0.0, 10.0, (n, 2)) + kwargs = dict( + vcov_type="conley", + conley_coords=coords, + conley_cutoff_km=2.0, + conley_metric="euclidean", + ) + diag = {} + (c_on, _, v_on), (c_off, _, v_off) = self._run_pair(monkeypatch, X, y, diag=diag, **kwargs) + assert diag["solve_ols_fastpath"] == "chol_numpy" + np.testing.assert_allclose(c_on, c_off, rtol=0, atol=1e-8) + np.testing.assert_allclose(np.sqrt(np.diag(v_on)), np.sqrt(np.diag(v_off)), rtol=1e-6) diff --git a/tests/test_rust_backend.py b/tests/test_rust_backend.py index 485f4f6f..f0ec5295 100644 --- a/tests/test_rust_backend.py +++ b/tests/test_rust_backend.py @@ -3674,3 +3674,506 @@ def test_raw_kernel_noncontiguous_ids_bit_identical(self): baseline = raw_vcov(X, resid, cl) for _ in range(20): np.testing.assert_array_equal(raw_vcov(X, resid, cl), baseline) + + +@pytest.mark.skipif(not HAS_RUST_BACKEND, reason="Rust backend not available") +class TestSolveOLSCholKernel: + """Direct tests of the opt-in Rust solve_ols_chol pyfunction. + + Parity vs the numpy Cholesky twin is TOL-BOUNDED at the established + cross-backend bars (coef/resid decimal=8, hc1 vcov decimal=8, clustered + decimal=6) on unit-scale fixtures: both sides run the same certified + algorithm, but Gram/matvec accumulation orders differ (faer vs BLAS). + Never byte-identity across the twins. + """ + + @staticmethod + def _design(seed=42, n=800, k=6): + rng = np.random.default_rng(seed) + X = np.column_stack([np.ones(n), rng.standard_normal((n, k - 1))]) + y = X @ rng.standard_normal(k) + rng.standard_normal(n) + return X, y + + def test_parity_vs_numpy_twin_hc1(self): + from diff_diff._rust_backend import solve_ols_chol + + from diff_diff.linalg import ( + _compute_robust_vcov_numpy, + _solve_ols_chol_numpy, + ) + + X, y = self._design() + out = solve_ols_chol(X, y) + assert out is not None, "well-conditioned design must certify" + coef_r, resid_r, vcov_r = out + + twin = _solve_ols_chol_numpy(X, y) + assert twin is not None + coef_p, gram = twin + resid_p = y - X @ coef_p + vcov_p = _compute_robust_vcov_numpy(X, resid_p, None, _bread_matrix=gram) + + np.testing.assert_array_almost_equal(coef_r, coef_p, decimal=8) + np.testing.assert_array_almost_equal(resid_r, resid_p, decimal=8) + np.testing.assert_array_almost_equal(vcov_r, vcov_p, decimal=8) + + def test_parity_vs_numpy_twin_clustered(self): + from diff_diff._rust_backend import solve_ols_chol + + from diff_diff.linalg import ( + _compute_robust_vcov_numpy, + _solve_ols_chol_numpy, + ) + + X, y = self._design(seed=7) + rng = np.random.default_rng(7) + cl = rng.integers(0, 25, X.shape[0]).astype(np.int64) + + out = solve_ols_chol(X, y, cluster_ids=cl) + assert out is not None + coef_r, resid_r, vcov_r = out + + twin = _solve_ols_chol_numpy(X, y) + assert twin is not None + coef_p, gram = twin + resid_p = y - X @ coef_p + vcov_p = _compute_robust_vcov_numpy(X, resid_p, cl, _bread_matrix=gram) + + np.testing.assert_array_almost_equal(coef_r, coef_p, decimal=8) + np.testing.assert_array_almost_equal(vcov_r, vcov_p, decimal=6) + + def test_returns_none_on_near_singular_and_nan(self): + from diff_diff._rust_backend import solve_ols_chol + + rng = np.random.default_rng(11) + n = 400 + x1 = rng.standard_normal(n) + X = np.column_stack([np.ones(n), x1, x1 + 1e-8 * rng.standard_normal(n)]) + y = rng.standard_normal(n) + assert solve_ols_chol(X, y) is None + + Xn, yn = self._design(seed=13, n=100, k=3) + Xn = Xn.copy() + Xn[5, 1] = np.nan + assert solve_ols_chol(Xn, yn) is None + + def test_returns_none_underdetermined(self): + from diff_diff._rust_backend import solve_ols_chol + + rng = np.random.default_rng(17) + X = rng.standard_normal((3, 5)) + y = rng.standard_normal(3) + assert solve_ols_chol(X, y) is None + + def test_return_vcov_false_shape(self): + from diff_diff._rust_backend import solve_ols_chol + + X, y = self._design(seed=19) + out = solve_ols_chol(X, y, return_vcov=False) + assert out is not None + coef, resid, vcov = out + assert vcov is None + assert len(coef) == X.shape[1] and len(resid) == X.shape[0] + + def test_too_few_clusters_raises_same_message(self): + from diff_diff._rust_backend import solve_ols_chol + + X, y = self._design(seed=23, n=60, k=3) + with pytest.raises(ValueError, match="Need at least 2 clusters"): + solve_ols_chol(X, y, cluster_ids=np.zeros(60, dtype=np.int64)) + + def test_repeated_call_bit_determinism(self): + """Same determinism lock as the SVD path's clustered vcov: 20 + repeated calls must be bit-identical (first-appearance cluster + order; sequential/faer accumulation).""" + from diff_diff._rust_backend import solve_ols_chol + + X, y = self._design(seed=29) + rng = np.random.default_rng(29) + cl = rng.integers(0, 12, X.shape[0]).astype(np.int64) + base = solve_ols_chol(X, y, cluster_ids=cl) + assert base is not None + for _ in range(20): + rep = solve_ols_chol(X, y, cluster_ids=cl) + np.testing.assert_array_equal(rep[0], base[0]) + np.testing.assert_array_equal(rep[1], base[1]) + np.testing.assert_array_equal(rep[2], base[2]) + + +@pytest.mark.skipif(not HAS_RUST_BACKEND, reason="Rust backend not available") +class TestSolveOLSFastpathDispatch: + """Dispatch-level locks for DIFF_DIFF_SOLVE_OLS_FASTPATH on the Rust + backend: default path unchanged (SVD kernel runs, chol never), opt-in + routes to the chol kernel, stale-symbol degradation to the numpy twin, + Rust-decline fallthrough verbatim, and estimator-level opt-in parity. + + Estimator-level parity is TOL-BOUNDED (ATT abs 1e-8 / SE rel 1e-6), NOT + the demean-chunk suite's rtol=0/atol=1e-12: that suite's exactness is + by-construction (identical partitioned arithmetic), while this knob + swaps the solve algorithm; the certified error budget is ~eps*cond(G_eq) + <= 2e-10 relative in the equilibrated basis, leaving >= 2 orders of + headroom at these gates. + """ + + @staticmethod + def _spies(monkeypatch): + import diff_diff.linalg as lmod + + calls = {"svd": 0, "chol": 0} + real_svd = lmod._solve_ols_rust + real_chol = lmod._solve_ols_chol_rust + + def svd_spy(*a, **kw): + calls["svd"] += 1 + return real_svd(*a, **kw) + + def chol_spy(*a, **kw): + calls["chol"] += 1 + return real_chol(*a, **kw) + + monkeypatch.setattr(lmod, "_solve_ols_rust", svd_spy) + monkeypatch.setattr(lmod, "_solve_ols_chol_rust", chol_spy) + return calls + + @staticmethod + def _staggered_panel(n_units=120, n_periods=8, seed=3): + rng = np.random.default_rng(seed) + unit = np.repeat(np.arange(n_units), n_periods) + time = np.tile(np.arange(n_periods), n_units) + first_treat = np.repeat(rng.choice([0, 3, 5], n_units), n_periods) + treated = (first_treat > 0) & (time >= first_treat) + y = ( + np.repeat(rng.normal(0, 1, n_units), n_periods) + + 0.1 * time + + 1.5 * treated + + rng.normal(0, 0.5, n_units * n_periods) + ) + return pd.DataFrame({"unit": unit, "time": time, "first_treat": first_treat, "y": y}) + + def test_default_dispatch_unchanged(self, monkeypatch): + monkeypatch.delenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", raising=False) + calls = self._spies(monkeypatch) + from diff_diff import SunAbraham + + SunAbraham().fit( + self._staggered_panel(), + outcome="y", + unit="unit", + time="time", + first_treat="first_treat", + ) + assert calls["chol"] == 0, "chol wrapper must never run with the knob off" + assert calls["svd"] >= 1, "the legacy SVD kernel must have run" + + def test_opt_in_dispatches_chol(self, monkeypatch): + monkeypatch.setenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", "1") + calls = self._spies(monkeypatch) + from diff_diff import SunAbraham + + SunAbraham().fit( + self._staggered_panel(), + outcome="y", + unit="unit", + time="time", + first_treat="first_treat", + ) + assert calls["chol"] >= 1, "knob on must route through the chol wrapper" + assert calls["svd"] == 0, "certified fits must not fall through to SVD" + + def test_stale_symbol_falls_back_to_legacy_svd(self, monkeypatch): + """A stale extension missing only solve_ols_chol keeps the legacy + Rust SVD kernel for Rust-eligible fits (the knob has no Rust + acceleration there); the numpy twin serves only the numpy lane.""" + import diff_diff.linalg as lmod + + monkeypatch.setenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", "1") + monkeypatch.setattr(lmod, "_rust_solve_ols_chol", None) + + rng = np.random.default_rng(31) + X = np.column_stack([np.ones(300), rng.standard_normal((300, 3))]) + y = X @ rng.standard_normal(4) + rng.standard_normal(300) + diag = {} + coef, _, _ = lmod.solve_ols(X, y, diagnostics_out=diag) + # Rust chol unavailable -> the (unweighted, hc1, full-rank) fit + # dispatches to the legacy Rust SVD kernel, NOT the numpy twin; + # the knob simply has nothing to accelerate on this route. + assert diag["solve_ols_fastpath"] == "fallback_declined" + assert np.all(np.isfinite(coef)) + + def test_rust_decline_falls_through_verbatim(self, monkeypatch): + """Near-singular fixture: chol kernel declines in-kernel; output + must equal the knob-off SVD result exactly (NaN-aware).""" + rng = np.random.default_rng(37) + n = 500 + x1 = rng.standard_normal(n) + X = np.column_stack([np.ones(n), x1, x1 + 1e-8 * rng.standard_normal(n)]) + y = X @ np.array([1.0, 2.0, 3.0]) + rng.standard_normal(n) + + import diff_diff.linalg as lmod + + monkeypatch.setenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", "1") + diag = {} + c_on, r_on, v_on = lmod.solve_ols(X, y, skip_rank_check=True, diagnostics_out=diag) + assert diag["solve_ols_fastpath"] == "fallback_declined" + monkeypatch.delenv("DIFF_DIFF_SOLVE_OLS_FASTPATH") + c_off, r_off, v_off = lmod.solve_ols(X, y, skip_rank_check=True) + np.testing.assert_array_equal(c_on, c_off) + np.testing.assert_array_equal(r_on, r_off) + # This fixture is SVD-truncated on the skip route -> NaN vcov on + # BOTH runs (pre-existing behavior); equal_nan makes that exact. + assert np.array_equal(v_on, v_off, equal_nan=True) + + def test_estimator_level_opt_in_parity_sun_abraham(self, monkeypatch): + from diff_diff import SunAbraham + + df = self._staggered_panel(n_units=200, n_periods=10, seed=41) + monkeypatch.delenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", raising=False) + res_off = SunAbraham().fit( + df, outcome="y", unit="unit", time="time", first_treat="first_treat" + ) + monkeypatch.setenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", "1") + res_on = SunAbraham().fit( + df, outcome="y", unit="unit", time="time", first_treat="first_treat" + ) + assert res_on.att == pytest.approx(res_off.att, abs=1e-8) + assert res_on.se == pytest.approx(res_off.se, rel=1e-6) + + def test_estimator_level_opt_in_parity_did_absorb_cluster(self, monkeypatch): + from diff_diff import DifferenceInDifferences + + rng = np.random.default_rng(43) + n_units, n_periods = 150, 6 + unit = np.repeat(np.arange(n_units), n_periods) + time = np.tile(np.arange(n_periods), n_units) + treated = (unit < n_units // 2).astype(int) + post = (time >= n_periods // 2).astype(int) + y = ( + np.repeat(rng.normal(0, 1, n_units), n_periods) + + 0.2 * time + + 2.0 * treated * post + + rng.normal(0, 0.5, n_units * n_periods) + ) + df = pd.DataFrame({"unit": unit, "time": time, "treated": treated, "post": post, "y": y}) + + def fit(): + return DifferenceInDifferences(cluster="unit").fit( + df, + outcome="y", + treatment="treated", + time="post", + absorb=["unit", "time"], + ) + + monkeypatch.delenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", raising=False) + res_off = fit() + monkeypatch.setenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", "1") + res_on = fit() + assert res_on.att == pytest.approx(res_off.att, abs=1e-8) + assert res_on.se == pytest.approx(res_off.se, rel=1e-6) + + def test_cross_backend_fitted_parity_rust_vs_twin(self, monkeypatch): + from unittest.mock import patch + + import diff_diff.linalg as lmod + + rng = np.random.default_rng(47) + X = np.column_stack([np.ones(600), rng.standard_normal((600, 5))]) + y = X @ rng.standard_normal(6) + rng.standard_normal(600) + + monkeypatch.setenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", "1") + d_rust = {} + c_rust, _, v_rust = lmod.solve_ols(X, y, diagnostics_out=d_rust) + assert d_rust["solve_ols_fastpath"] == "chol_rust" + + d_twin = {} + with ( + patch.object(lmod, "HAS_RUST_BACKEND", False), + patch.object(lmod, "_rust_solve_ols", None), + patch.object(lmod, "_rust_solve_ols_chol", None), + ): + c_twin, _, v_twin = lmod.solve_ols(X, y, diagnostics_out=d_twin) + assert d_twin["solve_ols_fastpath"] == "chol_numpy" + + np.testing.assert_allclose(X @ c_rust, X @ c_twin, rtol=1e-6, atol=1e-8) + np.testing.assert_array_almost_equal(v_rust, v_twin, decimal=6) + + +@pytest.mark.skipif(not HAS_RUST_BACKEND, reason="Rust backend not available") +class TestSolveOLSCholSaturatedContract: + """Saturated n == k designs must honor the documented + non-finite-inference contract on EVERY path: all-NaN vcov, never Inf + (the numpy saturated guard is the canonical behavior; the legacy Rust + SVD kernel's former silent-Inf leak was fixed alongside the opt-in + Cholesky path).""" + + @staticmethod + def _saturated_design(seed=61, k=6): + rng = np.random.default_rng(seed) + X = np.column_stack([np.ones(k), rng.standard_normal((k, k - 1))]) + y = rng.standard_normal(k) + return X, y + + def test_direct_kernel_saturated_nan_vcov(self): + from diff_diff._rust_backend import solve_ols_chol + + X, y = self._saturated_design() + two_clusters = (np.arange(X.shape[0]) % 2).astype(np.int64) + out = solve_ols_chol(X, y, cluster_ids=two_clusters) + assert out is not None, "saturated full-rank design must certify" + coef, resid, vcov = out + assert np.all(np.isfinite(coef)) + assert np.all(np.isnan(vcov)), "saturated vcov must be all-NaN, never Inf" + + out_hc1 = solve_ols_chol(X, y) + assert np.all(np.isnan(out_hc1[2])) + + def test_dispatch_saturated_nan_vcov(self, monkeypatch): + import diff_diff.linalg as lmod + + X, y = self._saturated_design(seed=67) + monkeypatch.setenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", "1") + diag = {} + coef, resid, vcov = lmod.solve_ols(X, y, diagnostics_out=diag) + assert diag["solve_ols_fastpath"] == "chol_rust" + assert np.all(np.isfinite(coef)) + assert np.all(np.isnan(vcov)) + + # The numpy twin lane honors the same contract via the shared + # saturated guard in _compute_robust_vcov_numpy. + from unittest.mock import patch + + diag2 = {} + with ( + patch.object(lmod, "HAS_RUST_BACKEND", False), + patch.object(lmod, "_rust_solve_ols", None), + patch.object(lmod, "_rust_solve_ols_chol", None), + ): + _, _, v_twin = lmod.solve_ols(X, y, diagnostics_out=diag2) + assert diag2["solve_ols_fastpath"] == "chol_numpy" + assert np.all(np.isnan(v_twin)) + + def test_saturated_one_cluster_still_raises(self, monkeypatch): + from diff_diff._rust_backend import solve_ols_chol + + X, y = self._saturated_design(seed=71) + with pytest.raises(ValueError, match="Need at least 2 clusters"): + solve_ols_chol(X, y, cluster_ids=np.zeros(X.shape[0], dtype=np.int64)) + + +@pytest.mark.skipif(not HAS_RUST_BACKEND, reason="Rust backend not available") +class TestSolveOLSSaturatedDefaultPath: + """Default (knob-off) saturated n == k contract: the legacy Rust SVD + vcov path formerly leaked all-Inf vcov silently (its (n-k) adjustment + divides by zero and the dispatcher's fallback check keyed on NaN only). + It now returns the canonical all-NaN contract, and the dispatcher + rejects ANY non-finite Rust vcov.""" + + @staticmethod + def _saturated_design(seed=73, k=5): + rng = np.random.default_rng(seed) + X = np.column_stack([np.ones(k), rng.standard_normal((k, k - 1))]) + y = rng.standard_normal(k) + return X, y + + def test_raw_svd_kernel_saturated_nan_vcov(self): + from diff_diff._rust_backend import solve_ols as rust_svd + + X, y = self._saturated_design() + coef, resid, vcov = rust_svd(X, y, None, True) + assert np.all(np.isfinite(np.asarray(coef))) + assert np.all(np.isnan(np.asarray(vcov))), "saturated vcov must be all-NaN, never Inf" + + two_clusters = (np.arange(X.shape[0]) % 2).astype(np.int64) + _, _, vcov_cl = rust_svd(X, y, two_clusters, True) + assert np.all(np.isnan(np.asarray(vcov_cl))) + + def test_public_solve_ols_saturated_nan_vcov(self, monkeypatch): + import diff_diff.linalg as lmod + + monkeypatch.delenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", raising=False) + X, y = self._saturated_design(seed=79) + with pytest.warns(UserWarning, match="non-finite"): + coef, resid, vcov = lmod.solve_ols(X, y) + assert np.all(np.isfinite(coef)) + assert np.all(np.isnan(vcov)) + + two_clusters = (np.arange(X.shape[0]) % 2).astype(np.int64) + with pytest.warns(UserWarning, match="non-finite"): + _, _, vcov_cl = lmod.solve_ols(X, y, cluster_ids=two_clusters) + assert np.all(np.isnan(vcov_cl)) + + def test_saturated_one_cluster_precedence(self): + from diff_diff._rust_backend import solve_ols as rust_svd + + X, y = self._saturated_design(seed=83) + with pytest.raises(ValueError, match="Need at least 2 clusters"): + rust_svd(X, y, np.zeros(X.shape[0], dtype=np.int64), True) + + def test_dispatcher_rejects_inf_vcov(self, monkeypatch): + """Any non-finite Rust vcov (not just NaN) must trigger the Python + re-run, which applies the canonical numpy contracts.""" + import diff_diff.linalg as lmod + + monkeypatch.delenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", raising=False) + rng = np.random.default_rng(89) + X = np.column_stack([np.ones(50), rng.standard_normal((50, 2))]) + y = X @ np.array([1.0, 2.0, 3.0]) + rng.standard_normal(50) + + def inf_vcov_rust(Xa, ya, *, cluster_ids=None, return_vcov=True, return_fitted=False): + coef = np.linalg.lstsq(Xa, ya, rcond=None)[0] + resid = ya - Xa @ coef + vcov = np.full((Xa.shape[1], Xa.shape[1]), np.inf) + if return_fitted: + return coef, resid, Xa @ coef, vcov + return coef, resid, vcov + + monkeypatch.setattr(lmod, "_solve_ols_rust", inf_vcov_rust) + with pytest.warns(UserWarning, match="non-finite"): + coef, resid, vcov = lmod.solve_ols(X, y) + assert np.all(np.isfinite(vcov)), "numpy re-run must replace the Inf vcov" + + def test_dispatcher_rejects_inf_vcov_skip_rank_check(self, monkeypatch): + """The Inf-rejection guard applies on the skip_rank_check route too + (same silent-corruption class as the rank-checked route).""" + import diff_diff.linalg as lmod + + monkeypatch.delenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", raising=False) + rng = np.random.default_rng(97) + X = np.column_stack([np.ones(50), rng.standard_normal((50, 2))]) + y = X @ np.array([1.0, 2.0, 3.0]) + rng.standard_normal(50) + + def inf_vcov_rust(Xa, ya, *, cluster_ids=None, return_vcov=True, return_fitted=False): + coef = np.linalg.lstsq(Xa, ya, rcond=None)[0] + resid = ya - Xa @ coef + vcov = np.full((Xa.shape[1], Xa.shape[1]), np.inf) + if return_fitted: + return coef, resid, Xa @ coef, vcov + return coef, resid, vcov + + monkeypatch.setattr(lmod, "_solve_ols_rust", inf_vcov_rust) + with pytest.warns(UserWarning, match="non-finite"): + coef, resid, vcov = lmod.solve_ols(X, y, skip_rank_check=True) + assert np.all(np.isfinite(vcov)), "numpy re-run must replace the Inf vcov" + + def test_skip_rank_check_nan_sentinel_passes_through(self, monkeypatch): + """On the skip route an all-NaN vcov remains the DOCUMENTED sentinel + for what the caller asserted away (rank deficiency / saturation) and + must pass through unchanged: rerouting it to a numpy path that + assumes full rank could raise on a singular bread where users + previously received the safe NaN answer.""" + import warnings as _warnings + + import diff_diff.linalg as lmod + + monkeypatch.delenv("DIFF_DIFF_SOLVE_OLS_FASTPATH", raising=False) + rng = np.random.default_rng(101) + x1 = rng.standard_normal(60) + # Exactly collinear: the rust SVD kernel truncates (rank < k) and + # returns its all-NaN vcov sentinel. + X = np.column_stack([np.ones(60), x1, 2.0 * x1]) + y = x1 + rng.standard_normal(60) + + with _warnings.catch_warnings(): + _warnings.simplefilter("error") # no warning may fire + coef, resid, vcov = lmod.solve_ols(X, y, skip_rank_check=True) + assert np.all(np.isnan(vcov)), "NaN sentinel must pass through on the skip route"