diff --git a/CHANGELOG.md b/CHANGELOG.md index 1c0422ca..2805f125 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,6 +33,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 attribution in `docs/performance-plan.md`). ### Changed +- **`solve_ols` marshalling slimmed on both backends** (memory + speed; outputs preserved). + Rust: column norms computed from the borrowed numpy view, equilibration fused into the + single owned faer copy, `U'y` computed directly off the faer SVD factor (the old path + materialized a full ndarray copy of U for one k-vector dot), and factors dropped before + the fitted/residual/vcov stage - rust-side allocator high-water on a 2.4M x 130 clustered + solve falls from 15.32 GB to 7.81 GB (measured with the new feature-gated `alloc-profile` + counting allocator; never compiled into wheels). Python: rank detection uses + `qr(mode="r")` (same dgeqp3 R and pivot, skips forming the discarded Q) and + `_equilibrated_lstsq` materializes its scaled temporary F-order with `overwrite_a=True` + so gelsd consumes it in place - python-backend estimates are bit-identical (benchmark + identity deltas exactly 0.0). The removed copies also cost time: frozen-code benchmark + arms measured CV-clear improvements incl. firm-panel rust fits -21% (25.8 -> 20.3 s), + county event studies -12%/-13% (rust/python), with no scenario regressing. Rust outputs + move only at the cross-backend parity tolerances (REGISTRY-documented; never + bit-identical across backends). - **FE-absorption demeaning rewritten: factorize-once + `np.bincount` method of alternating projections** (`demean_by_groups` / `within_transform`). Each absorbed dimension is factorized once and group means are formed via `np.bincount` instead of re-hashing the group keys with diff --git a/TODO.md b/TODO.md index ca0e0fcf..9cd6b3fd 100644 --- a/TODO.md +++ b/TODO.md @@ -46,7 +46,6 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| -| Rust-backend `solve_ols` peak memory on wide designs: the faer path holds an owned copy of the stacked design plus SVD workspace, dominating fit-level peak RSS on wide absorbed fits (firm_churn-class: ~19-21 GB rust-backend vs ~13.4 GB numpy-backend for the same fit; measured PR-D 2026-07, corrected attribution in docs/performance-plan.md - the demean kernel's transients were NOT the dominator, and opt-in chunked dispatch caps them anyway). Candidates: borrow-view instead of owned copy into faer, workspace reuse, or routing wide designs to the numpy gelsd path. Interim escape hatch: DIFF_DIFF_BACKEND=python. | `rust/src/linalg.rs`, `diff_diff/linalg.py::solve_ols` | PR-D probe | Heavy | Medium | | `ImputationDiD` dense `(A0'A0).toarray()` scales `O((U+T+K)^2)` — OOM risk on large panels (only triggers when the sparse solver fails). Needs an alternative dense fallback or richer sparse strategy. | `imputation.py` | #141 | Heavy | Medium | | CR2 Bell-McCaffrey DOF uses a naive `O(n²k)` per-coefficient loop over cluster pairs; Pustejovsky-Tipton (2018) Appendix B has a scores-based formulation avoiding the full `n×n` `M`. Switch when a user hits a large-`n` cluster-robust design. | `linalg.py::_compute_cr2_bm` | Phase 1a | Heavy | Low | | Rust-backend HC2: the Rust path only supports HC1; HC2 and CR2 Bell-McCaffrey fall through to NumPy. Noticeable for large-`n` fits. | `rust/src/linalg.rs` | Phase 1a | Mid | Low | @@ -54,7 +53,6 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Wild cluster bootstrap CI inversion calls `_t_star(r)` ~O(100) times, each materializing a fresh `(B×n)` `y_star` + `(k×B)` refit + `(n×B)` residual arrays. Acceptable for the few-cluster regime; for large-`n`/large-`B`, chunk `_t_star` over draws or precompute the `r`-independent cluster-level pieces (restricted residuals are linear in `r`). | `utils.py::wild_bootstrap_se._t_star` | #543 | Mid | Low | | `SpilloverDiD` sparse cKDTree path for the staggered nearest-treated-distance helper (mirrors the static helper's sparse branch). `_compute_nearest_treated_distance_staggered` always builds dense `(n_units, n_treated_by_onset)` matrices per cohort; add a sparse branch gated on `n > _CONLEY_SPARSE_N_THRESHOLD`. | `spillover.py` | Wave B | Mid | Low | | `HeterogeneousAdoptionDiD` Phase 3 Stute: Appendix-D vectorized form replaces the per-iteration OLS refit with a single precomputed `M = I - X(X'X)^{-1}X'` applied to `eps*eta` (~2× faster, functionally identical). Shipped the literal-refit form to match paper text. | `had_pretests.py::stute_test` | Phase 3 | Mid | Low | -| Rust faer SVD ndarray-to-faer conversion overhead (minimal vs SVD cost). | `rust/src/linalg.rs:67` | #115 | Quick | Low | | Multiplier-bootstrap weight chunking (CallawaySantAnna, EfficientDiD, and HAD — all wired through `diff_diff/bootstrap_chunking.py`) covers the **unstratified** survey-PSU generation (the default unit-level bootstrap — `cluster=None`, equivalently `cluster="unit"` — the large-`n_units` OOM case). Remaining gap: the **stratified** survey-PSU generator (`generate_survey_multiplier_weights_batch`, per-stratum + lonely-PSU pooling + FPC) still materializes the full `(n_bootstrap × n_psu)` matrix (consumed via sliced blocks). Stratified designs have few PSUs so this rarely OOMs; tile per-stratum generation over draws (each stratum's draws are independent → contiguous draw-blocks reproduce the stream bit-identically) if a large-PSU stratified design hits memory. | `diff_diff/bootstrap_chunking.py::iter_survey_multiplier_weight_blocks` | follow-up | Mid | Low | | `ImputationDiD._iterative_demean` (`imputation.py:1064`) and its near-identical twin `TwoStageDiD._iterative_demean` (`two_stage.py:2201`) rebuild a `pd.Series`+`groupby().transform()` every alternating-projection iteration. Precompute the unit/time group codes once and use `np.bincount` for the per-iteration group means. **Not bit-identical** on pandas 3.0 (`np.bincount` is naive accumulation; pandas' `group_mean` is Kahan-compensated → ~5.8e-11, the same order as the demean's `tol=1e-10`), so this needs a tolerance/REGISTRY-Note treatment + golden re-validation. Modest peak effect (the per-iteration Series are transient), analytical-path-only (the bootstrap reuses #562's cached projection). Optimize both twins together (cross-surface-twins). | `imputation.py`, `two_stage.py` | PR-C deferral | Mid | Low | @@ -124,6 +122,7 @@ Doable in principle, but no current caller and/or explicitly out of paper scope. | Issue | Location | PR | Priority | |-------|----------|----|----------| | `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 | | 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` survey-design API consolidation (**scheduled: next minor bump**): drop the deprecated `survey=` / `weights=` kwargs on all 8 HAD surfaces; only `survey_design=` remains. Also fold the legacy back-end `weights=` routing into the unified `_resolve_survey_for_fit` path. DeprecationWarning has shipped; the removal is ~50 LoC gated on the semver bump. | `had.py`, `had_pretests.py` | next minor | Medium | diff --git a/benchmarks/speed_review/baselines/fe_absorption_after.json b/benchmarks/speed_review/baselines/fe_absorption_after.json index 30e2ec59..2e7dff91 100644 --- a/benchmarks/speed_review/baselines/fe_absorption_after.json +++ b/benchmarks/speed_review/baselines/fe_absorption_after.json @@ -25,14 +25,14 @@ "checksum": 6315.071127371084, "att": 0.2803785413990301, "se": 0.031244864610215972, - "fit_median_s": 2.3882, - "fit_min_s": 2.3505, - "fit_max_s": 2.4948, - "fit_cv": 0.0257, + "fit_median_s": 2.0868, + "fit_min_s": 2.0496, + "fit_max_s": 2.1839, + "fit_cv": 0.0255, "noisy": false, "n_timed_fits": 9, - "datagen_s": 0.004, - "peak_rss_mb": 3177.5, + "datagen_s": 0.005, + "peak_rss_mb": 2855.8, "backend_resolved": "python", "warnings": [] }, @@ -42,14 +42,14 @@ "checksum": 240997.06223418942, "att": 0.3144327184497798, "se": 0.005837281701253927, - "fit_median_s": 49.502, - "fit_min_s": 48.4532, - "fit_max_s": 50.3118, - "fit_cv": 0.0189, + "fit_median_s": 46.1275, + "fit_min_s": 44.4378, + "fit_max_s": 46.8508, + "fit_cv": 0.027, "noisy": false, "n_timed_fits": 3, - "datagen_s": 0.07, - "peak_rss_mb": 15900.8, + "datagen_s": 0.072, + "peak_rss_mb": 15072.3, "backend_resolved": "python", "warnings": [] }, @@ -59,14 +59,14 @@ "checksum": 76832.22571739106, "att": 0.251164131495566, "se": 0.002414938742160613, - "fit_median_s": 0.9834, - "fit_min_s": 0.9825, - "fit_max_s": 0.9845, - "fit_cv": 0.001, + "fit_median_s": 0.9799, + "fit_min_s": 0.9765, + "fit_max_s": 0.9907, + "fit_cv": 0.0076, "noisy": false, "n_timed_fits": 3, - "datagen_s": 0.077, - "peak_rss_mb": 976.9, + "datagen_s": 0.076, + "peak_rss_mb": 902.4, "backend_resolved": "python", "warnings": [] }, @@ -76,14 +76,14 @@ "checksum": 443696.3834000407, "att": 0.25194372599916653, "se": 0.0017886314723251327, - "fit_median_s": 0.973, - "fit_min_s": 0.968, - "fit_max_s": 0.9759, - "fit_cv": 0.0041, + "fit_median_s": 0.9421, + "fit_min_s": 0.9417, + "fit_max_s": 0.9425, + "fit_cv": 0.0004, "noisy": false, "n_timed_fits": 3, - "datagen_s": 0.062, - "peak_rss_mb": 1664.8, + "datagen_s": 0.061, + "peak_rss_mb": 1668.0, "backend_resolved": "python", "warnings": [ "Rank-deficient design matrix: dropping 2 of 4 columns (column 1, column 2). Coefficients for these columns are", @@ -96,14 +96,14 @@ "checksum": -176482.6014001448, "att": 0.24554568144130481, "se": 0.007662698712656388, - "fit_median_s": 4.0767, - "fit_min_s": 4.0736, - "fit_max_s": 4.0778, - "fit_cv": 0.0005, + "fit_median_s": 4.0394, + "fit_min_s": 4.0254, + "fit_max_s": 4.0462, + "fit_cv": 0.0026, "noisy": false, "n_timed_fits": 3, "datagen_s": 0.107, - "peak_rss_mb": 3589.2, + "peak_rss_mb": 3578.0, "backend_resolved": "python", "warnings": [ "Rank-deficient design matrix: dropping 2 of 4 columns (column 1, column 2). Coefficients for these columns are", @@ -116,14 +116,14 @@ "checksum": 376703.7248141762, "att": 0.2502082560083505, "se": 0.004392830594147583, - "fit_median_s": 31.7648, - "fit_min_s": 31.4957, - "fit_max_s": 31.8372, - "fit_cv": 0.0057, + "fit_median_s": 31.656, + "fit_min_s": 31.5351, + "fit_max_s": 31.7636, + "fit_cv": 0.0036, "noisy": false, "n_timed_fits": 3, "datagen_s": 0.067, - "peak_rss_mb": 1840.7, + "peak_rss_mb": 1839.7, "backend_resolved": "python", "warnings": [ "Rank-deficient design matrix: dropping 2 of 4 columns (column 1, column 2). Coefficients for these columns are", @@ -136,14 +136,14 @@ "checksum": -571.6190664591859, "att": 0.2515087011451766, "se": 0.028726073333820948, - "fit_median_s": 0.0035, - "fit_min_s": 0.0034, - "fit_max_s": 0.0039, - "fit_cv": 0.0437, + "fit_median_s": 0.0036, + "fit_min_s": 0.0035, + "fit_max_s": 0.0042, + "fit_cv": 0.0487, "noisy": false, "n_timed_fits": 21, "datagen_s": 0.001, - "peak_rss_mb": 145.1, + "peak_rss_mb": 148.1, "backend_resolved": "python", "warnings": [] } diff --git a/benchmarks/speed_review/baselines/fe_absorption_rust.json b/benchmarks/speed_review/baselines/fe_absorption_rust.json index 58e4518e..1919ce74 100644 --- a/benchmarks/speed_review/baselines/fe_absorption_rust.json +++ b/benchmarks/speed_review/baselines/fe_absorption_rust.json @@ -23,16 +23,16 @@ "scenario": "county_policy", "n_obs": 177289, "checksum": 6315.071127371084, - "att": 0.2803785413990308, - "se": 0.03124486461021617, - "fit_median_s": 1.7567, - "fit_min_s": 1.7285, - "fit_max_s": 1.8617, - "fit_cv": 0.0302, + "att": 0.28037854139903057, + "se": 0.031244864610215826, + "fit_median_s": 1.5536, + "fit_min_s": 1.4852, + "fit_max_s": 1.6031, + "fit_cv": 0.0286, "noisy": false, "n_timed_fits": 9, "datagen_s": 0.005, - "peak_rss_mb": 4237.1, + "peak_rss_mb": 3264.3, "backend_resolved": "rust", "warnings": [] }, @@ -40,16 +40,16 @@ "scenario": "firm_churn", "n_obs": 2400000, "checksum": 240997.06223418942, - "att": 0.31443271844971477, - "se": 0.005837281701253889, - "fit_median_s": 25.8387, - "fit_min_s": 24.6121, - "fit_max_s": 26.7088, - "fit_cv": 0.041, + "att": 0.314432718449715, + "se": 0.005837281701253807, + "fit_median_s": 20.341, + "fit_min_s": 19.7481, + "fit_max_s": 20.7773, + "fit_cv": 0.0255, "noisy": false, "n_timed_fits": 3, - "datagen_s": 0.069, - "peak_rss_mb": 20963.6, + "datagen_s": 0.07, + "peak_rss_mb": 17101.5, "backend_resolved": "rust", "warnings": [] }, @@ -57,16 +57,16 @@ "scenario": "scanner_twfe", "n_obs": 3255000, "checksum": 76832.22571739106, - "att": 0.25116413149501093, - "se": 0.0024149387421606118, - "fit_median_s": 0.6239, - "fit_min_s": 0.6215, - "fit_max_s": 0.6252, - "fit_cv": 0.003, + "att": 0.25116413149499167, + "se": 0.0024149387421606113, + "fit_median_s": 0.6118, + "fit_min_s": 0.6118, + "fit_max_s": 0.6154, + "fit_cv": 0.0034, "noisy": false, "n_timed_fits": 3, "datagen_s": 0.076, - "peak_rss_mb": 1477.0, + "peak_rss_mb": 1327.8, "backend_resolved": "rust", "warnings": [] }, @@ -76,14 +76,14 @@ "checksum": 443696.3834000407, "att": 0.25194372599916653, "se": 0.0017886314723251327, - "fit_median_s": 0.4839, - "fit_min_s": 0.483, - "fit_max_s": 0.4845, - "fit_cv": 0.0016, + "fit_median_s": 0.4525, + "fit_min_s": 0.4515, + "fit_max_s": 0.4566, + "fit_cv": 0.006, "noisy": false, "n_timed_fits": 3, "datagen_s": 0.062, - "peak_rss_mb": 1744.8, + "peak_rss_mb": 1741.9, "backend_resolved": "rust", "warnings": [ "Rank-deficient design matrix: dropping 2 of 4 columns (column 1, column 2). Coefficients for these columns are", @@ -96,14 +96,14 @@ "checksum": -176482.6014001448, "att": 0.24554568144130481, "se": 0.007662698712656388, - "fit_median_s": 2.8354, - "fit_min_s": 2.8354, - "fit_max_s": 2.8357, - "fit_cv": 0.0001, + "fit_median_s": 2.778, + "fit_min_s": 2.7747, + "fit_max_s": 2.7822, + "fit_cv": 0.0014, "noisy": false, "n_timed_fits": 3, - "datagen_s": 0.106, - "peak_rss_mb": 3587.3, + "datagen_s": 0.105, + "peak_rss_mb": 3586.4, "backend_resolved": "rust", "warnings": [ "Rank-deficient design matrix: dropping 2 of 4 columns (column 1, column 2). Coefficients for these columns are", @@ -116,14 +116,14 @@ "checksum": 376703.7248141762, "att": 0.2502082560083505, "se": 0.004392830594147583, - "fit_median_s": 11.365, - "fit_min_s": 11.3007, - "fit_max_s": 11.4856, - "fit_cv": 0.0082, + "fit_median_s": 11.3811, + "fit_min_s": 11.3491, + "fit_max_s": 11.4257, + "fit_cv": 0.0034, "noisy": false, "n_timed_fits": 3, - "datagen_s": 0.066, - "peak_rss_mb": 1841.8, + "datagen_s": 0.067, + "peak_rss_mb": 1842.5, "backend_resolved": "rust", "warnings": [ "Rank-deficient design matrix: dropping 2 of 4 columns (column 1, column 2). Coefficients for these columns are", @@ -134,16 +134,16 @@ "scenario": "guard_small", "n_obs": 20000, "checksum": -571.6190664591859, - "att": 0.2515087011451809, - "se": 0.028726073333820955, - "fit_median_s": 0.0034, + "att": 0.2515087011451813, + "se": 0.028726073333820944, + "fit_median_s": 0.0035, "fit_min_s": 0.0032, - "fit_max_s": 0.0036, - "fit_cv": 0.0346, + "fit_max_s": 0.0039, + "fit_cv": 0.0512, "noisy": false, "n_timed_fits": 21, "datagen_s": 0.001, - "peak_rss_mb": 148.0, + "peak_rss_mb": 148.3, "backend_resolved": "rust", "warnings": [] } diff --git a/diff_diff/linalg.py b/diff_diff/linalg.py index 38edf953..d122784a 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -145,7 +145,10 @@ def _detect_rank_deficiency( def _rank_and_pivot(M: np.ndarray) -> Tuple[int, np.ndarray]: # Pivoted QR: M @ P = Q @ R. The rank threshold is anchored to the # largest pivot diagonal |R[0,0]| (decreasing after pivoting). - _Q, R, piv = qr(M, mode="economic", pivoting=True) + # mode="r" skips forming the (unused) Q: R and the pivot come from + # the same dgeqp3 factorization either way, so rank decisions are + # bit-identical; only the dorgqr Q-formation work is saved. + R, piv = qr(M, mode="r", pivoting=True) r_diag = np.abs(np.diag(R)) if r_diag[0] == 0: return 0, piv @@ -308,8 +311,20 @@ def _equilibrated_lstsq(X: np.ndarray, y: np.ndarray) -> np.ndarray: """ col_norms = np.sqrt(np.einsum("ij,ij->j", X, X)) safe_norms = np.where(col_norms > 0, col_norms, 1.0) + # Materialize the scaled temporary F-order and let gelsd consume it + # in place: scipy honors overwrite_a only for F-contiguous input (it + # silently re-copies C-order), and LAPACK wants F-order anyway, so this + # skips the one internal copy lstsq would otherwise make. The temporary + # is exclusively ours (never reused after the call), and the VALUES are + # identical to the previous `X / safe_norms`, so results are bit-equal. + x_scaled = np.divide(X, safe_norms, out=np.empty(X.shape, order="F")) coef_scaled = scipy_lstsq( - X / safe_norms, y, lapack_driver="gelsd", check_finite=False, cond=1e-07 + x_scaled, + y, + lapack_driver="gelsd", + check_finite=False, + cond=1e-07, + overwrite_a=True, )[0] return coef_scaled / safe_norms @@ -459,7 +474,7 @@ def _ret(inv, n_dropped, n_keep, dropped): # IF multiplier leaves range(A) — e.g. a treated-cell mean). Equilibrating the # selection (rather than a raw pivot) keeps the cell-only-aliasing SE equal to # the well-conditioned near-collinear limit (se_ratio ~ 1). - _Q, _R, piv = qr(A_eq, mode="economic", pivoting=True) + _R, piv = qr(A_eq, mode="r", pivoting=True) # only the pivot is used kept = np.sort(piv[:n_keep]) A_eq_ginv = np.zeros((k, k), dtype=float) A_eq_ginv[np.ix_(kept, kept)] = np.linalg.inv(A_eq[np.ix_(kept, kept)]) diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 94bcb3ec..0abdfecd 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -547,7 +547,7 @@ The multiplier bootstrap uses random weights w_i with E[w]=0 and Var(w)=1: - Detection: Pivoted QR decomposition with tolerance `1e-07` (R's `qr()` default), with a **column-equilibration re-check** (unit 2-norm) that makes the rank *count* invariant to per-column scaling; the dropped-column selection is unchanged for well-scaled collinear designs (a scale-induced under-count instead adopts the scale-corrected equilibrated selection) - 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. 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:** 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). - 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 ad2b8538..e2a245b8 100644 --- a/docs/performance-plan.md +++ b/docs/performance-plan.md @@ -76,35 +76,49 @@ submit gate), tail_stress 31.74s -> 11.37s (2.79x), geo_experiment 0.976s -> 0.484s (2.02x, now ahead of the pyfixest yardstick's 0.623s), scanner 1.59x, survey 1.44x, county 1.36x, guard_small unregressed. Identity vs the committed after-baselines: 1e-13-1e-16 on every scenario incl. tail_stress. -Known trade-off: the rust-backend firm_churn fit peaks at ~21 GB RSS vs -~13.4 GB under the numpy backend. **Corrected attribution (2026-07, PR-D -probes):** the peak is dominated by the downstream SOLVER phase (the Rust -faer path holds an owned copy of the ~2.5 GB stacked design plus SVD -workspace), not by the demean kernel's dispatch marshalling as this -paragraph originally claimed. Measured evidence: an isolated width-16 -chunked dispatch cuts the kernel's transients to near-numpy footprint, yet -the end-to-end fit still peaks ~19-21 GB (only ~5-12% below unchunked) at a -+2-7% wall-clock cost - so chunked dispatch shipped as an opt-in env knob -(`DIFF_DIFF_DEMEAN_CHUNK_COLS`, default off) rather than default-on, and -the solver-phase peak is tracked as its own TODO row. `DIFF_DIFF_BACKEND= -python` remains the OOM escape hatch. Measurement note: single-run -`ru_maxrss` on macOS is unreliable under ambient memory pressure (the -compressor deflates resident peaks - one probe read 15.1 GB for a -configuration that reproducibly peaks ~19-20 GB); gate memory claims on -repeated runs under matched machine state. +**Solver-phase memory: final diagnosis + fix (2026-07, PR-D/PR-E).** The +fit-level peak on wide absorbed designs lives in the SOLVE phase on BOTH +backends. There is NO systematic rust-vs-python peak-RSS gap: an +interleaved A/B (3 rounds, alternating fresh subprocesses) measured rust +17.8 GB mean vs python 18.1 GB on firm_churn, each swinging +-1.5 GB +run-to-run - the historical "13.4 vs 21.0 GB backend gap" was a run-order +artifact (macOS's memory compressor evicts cold pages from the resident +set, so single-run `ru_maxrss` readings are machine-state lottery; gate +memory claims on allocation-level instruments or interleaved repeats, +never one RSS pair). The genuine waste was marshalling: the rust +`solve_ols` held ~6 concurrent n x k blocks (defensive input copy, scaled +clone, faer conversion copy, thin-SVD U, an ndarray COPY of U made only +for a k-vector dot, all scope-held under the vcov scores block) and the +python path formed a discarded Q in the rank-detection QR plus let gelsd +re-copy a C-order temporary. The PR-E slim (norms from the borrowed view, +equilibration fused into the single faer copy, U^T y off the faer factor, +early factor drops; python: `qr(mode="r")` + F-order `overwrite_a=True` +lstsq) cut the rust-side allocator high-water on the 2.4M x 130 clustered +solve from **15.32 GB to 7.81 GB** (6.14 -> 3.13 blocks, measured with the +feature-gated `alloc-profile` counting allocator) and, because the removed +copies also cost time, delivered CV-clear wall-clock wins on frozen-code +arms: firm_churn rust 25.84 -> 20.34 s (-21%), county 1.76 -> 1.55 s / +2.39 -> 2.09 s (rust/python, the skipped dorgqr), geo rust -6%, survey +rust -2%; python-arm estimates bit-identical (identity deltas exactly +0.0). Remaining floor is inherent to SVD-based lstsq (fused input copy + +thin-SVD U transient + vcov scores block + LAPACK gelsd internals); +further reduction = the tall-skinny-QR path parked with the QR-reuse row +in TODO.md. `DIFF_DIFF_BACKEND=python` remains an escape hatch, and the +opt-in `DIFF_DIFF_DEMEAN_CHUNK_COLS` knob (PR-D) still caps the demean +dispatch transients. ### FE-absorption suite results | Scenario | n rows | Before (s) | After (s) | Speedup | Rust (s) | Rust speedup | pyfixest (s) | |---|---:|---:|---:|---:|---:|---:|---:| -| 7. County policy event study (SunAbraham) | 177,289 | 3.865 (cv 2.5%) | 2.388 (cv 2.6%) | 1.6x | 1.757 (cv 3.0%) | 1.4x | 0.190 (cv 4.0%, proxy) | -| 8. Firm panel with churn (SunAbraham) | 2,400,000 | 92.957 (cv 0.9%) | 49.502 (cv 1.9%) | 1.9x | 25.839 (cv 4.1%) | 1.9x | 1.912 (cv 20.3%, noisy, proxy) | -| 9. Scanner store-week (TWFE) | 3,255,000 | 1.549 (cv 0.3%) | 0.983 (cv 0.1%) | 1.6x | 0.624 (cv 0.3%) | 1.6x | 0.644 (cv 12.3%, noisy) | -| 10. Geo experiment 5M orders (DiD absorb) | 5,000,000 | 2.630 (cv 0.6%) | 0.973 (cv 0.4%) | 2.7x | 0.484 (cv 0.2%) | 2.0x | 0.623 (cv 7.8%) | -| 11. Survey BRR replicates (DiD absorb) | 500,000 | 7.385 (cv 8.6%) | 4.077 (cv 0.1%) | 1.8x | 2.835 (cv 0.0%) | 1.4x | - | -| 12. Correlated-FE stress (DiD absorb) | 5,000,000 | 26.271 (cv 0.4%) | 31.765 (cv 0.6%) | 0.8x | 11.365 (cv 0.8%) | 2.8x | 3.075 (cv 1.5%) | -| 13. Small-panel guard (TWFE) | 20,000 | 0.005 (cv 4.1%) | 0.004 (cv 4.4%) | 1.3x | 0.003 (cv 3.5%) | 1.0x | 0.007 (cv 2.2%) | +| 7. County policy event study (SunAbraham) | 177,289 | 3.865 (cv 2.5%) | 2.087 (cv 2.5%) | 1.9x | 1.554 (cv 2.9%) | 1.3x | 0.190 (cv 4.0%, proxy) | +| 8. Firm panel with churn (SunAbraham) | 2,400,000 | 92.957 (cv 0.9%) | 46.127 (cv 2.7%) | 2.0x | 20.341 (cv 2.5%) | 2.3x | 1.912 (cv 20.3%, noisy, proxy) | +| 9. Scanner store-week (TWFE) | 3,255,000 | 1.549 (cv 0.3%) | 0.980 (cv 0.8%) | 1.6x | 0.612 (cv 0.3%) | 1.6x | 0.644 (cv 12.3%, noisy) | +| 10. Geo experiment 5M orders (DiD absorb) | 5,000,000 | 2.630 (cv 0.6%) | 0.942 (cv 0.0%) | 2.8x | 0.453 (cv 0.6%) | 2.1x | 0.623 (cv 7.8%) | +| 11. Survey BRR replicates (DiD absorb) | 500,000 | 7.385 (cv 8.6%) | 4.039 (cv 0.3%) | 1.8x | 2.778 (cv 0.1%) | 1.5x | - | +| 12. Correlated-FE stress (DiD absorb) | 5,000,000 | 26.271 (cv 0.4%) | 31.656 (cv 0.4%) | 0.8x | 11.381 (cv 0.3%) | 2.8x | 3.075 (cv 1.5%) | +| 13. Small-panel guard (TWFE) | 20,000 | 0.005 (cv 4.1%) | 0.004 (cv 4.9%) | 1.3x | 0.004 (cv 5.1%) | 1.0x | 0.007 (cv 2.2%) | *noisy* = CV above the 10% unusable threshold from the noise protocol. *Rust speedup* is vs the After column (numpy engine, same code). *proxy* = timing-only pyfixest stand-in: the Sun-Abraham scenarios run a saturated `i(rel_time)` event study there (pyfixest 0.60 has no `sunab()`) - comparable demeaning load, different estimand, so those cells are not exact-estimand comparisons (see `bench_fe_absorption_pyfixest.py`). diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 16dae3c5..e1ae59e8 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -19,6 +19,10 @@ extension-module = ["pyo3/extension-module"] # When enabled, ndarray's .dot() and general_mat_vec_mul dispatch to BLAS dgemv/dgemm accelerate = ["ndarray/blas", "dep:blas-src", "blas-src/accelerate"] openblas = ["ndarray/blas"] +# Counting global allocator for memory-footprint measurement builds ONLY +# (never enabled in wheels/CI): exposes reset_alloc_high_water / +# alloc_high_water_bytes pyfunctions. See rust/src/alloc_profile.rs. +alloc-profile = [] [dependencies] # PyO3 0.29 supports Python 3.9-3.14 diff --git a/rust/src/alloc_profile.rs b/rust/src/alloc_profile.rs new file mode 100644 index 00000000..6bff9e71 --- /dev/null +++ b/rust/src/alloc_profile.rs @@ -0,0 +1,70 @@ +//! Feature-gated counting allocator for memory-footprint measurement builds. +//! +//! Compiled ONLY under `--features alloc-profile` (never in shipped wheels). +//! Wraps the system allocator with atomic current/high-water byte counters so +//! before/after footprint claims can be verified at the allocation level - +//! resident-set metrics are unreliable for this purpose on macOS, where the +//! memory compressor evicts cold pages from RSS (see docs/performance-plan.md). +//! +//! Scope: counts RUST-side allocations in this extension only; numpy/Python +//! allocations are invisible - which is exactly the surface under measurement. + +use pyo3::prelude::*; +use std::alloc::{GlobalAlloc, Layout, System}; +use std::sync::atomic::{AtomicUsize, Ordering}; + +pub struct CountingAlloc; + +static CURRENT: AtomicUsize = AtomicUsize::new(0); +static HIGH_WATER: AtomicUsize = AtomicUsize::new(0); + +fn track_alloc(size: usize) { + let now = CURRENT.fetch_add(size, Ordering::Relaxed) + size; + // Racy max update is fine for a diagnostic: retry while below. + let mut hw = HIGH_WATER.load(Ordering::Relaxed); + while now > hw { + match HIGH_WATER.compare_exchange_weak(hw, now, Ordering::Relaxed, Ordering::Relaxed) { + Ok(_) => break, + Err(actual) => hw = actual, + } + } +} + +unsafe impl GlobalAlloc for CountingAlloc { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { System.alloc(layout) }; + if !ptr.is_null() { + track_alloc(layout.size()); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { System.dealloc(ptr, layout) }; + CURRENT.fetch_sub(layout.size(), Ordering::Relaxed); + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let new_ptr = unsafe { System.realloc(ptr, layout, new_size) }; + if !new_ptr.is_null() { + CURRENT.fetch_sub(layout.size(), Ordering::Relaxed); + track_alloc(new_size); + } + new_ptr + } +} + +#[global_allocator] +static GLOBAL: CountingAlloc = CountingAlloc; + +/// Reset the high-water mark to the current live-byte count. +#[pyfunction] +pub fn reset_alloc_high_water() { + HIGH_WATER.store(CURRENT.load(Ordering::Relaxed), Ordering::Relaxed); +} + +/// Peak rust-side allocated bytes since the last reset. +#[pyfunction] +pub fn alloc_high_water_bytes() -> usize { + HIGH_WATER.load(Ordering::Relaxed) +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 07e87b5d..30b6312f 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -14,6 +14,8 @@ extern crate blas_src; use pyo3::prelude::*; use std::collections::HashMap; +#[cfg(feature = "alloc-profile")] +mod alloc_profile; mod bootstrap; mod demean; mod linalg; @@ -63,6 +65,13 @@ fn _rust_backend(m: &Bound<'_, PyModule>) -> PyResult<()> { // Diagnostics m.add_function(wrap_pyfunction!(rust_backend_info, m)?)?; + // Allocation profiling (measurement builds only; --features alloc-profile) + #[cfg(feature = "alloc-profile")] + { + m.add_function(wrap_pyfunction!(alloc_profile::reset_alloc_high_water, m)?)?; + m.add_function(wrap_pyfunction!(alloc_profile::alloc_high_water_bytes, m)?)?; + } + // Version info m.add("__version__", env!("CARGO_PKG_VERSION"))?; diff --git a/rust/src/linalg.rs b/rust/src/linalg.rs index 9ce93887..c9940277 100644 --- a/rust/src/linalg.rs +++ b/rust/src/linalg.rs @@ -59,9 +59,13 @@ pub fn solve_ols<'py>( let k = x_arr.ncols(); // Solve using SVD with truncation for rank-deficient matrices - // This matches scipy's 'gelsd' behavior - let x_owned = x_arr.to_owned(); - let y_owned = y_arr.to_owned(); + // This matches scipy's 'gelsd' behavior. + // + // Memory discipline (wide designs make one n x k block ~GBs): work from + // the borrowed numpy views, materialize exactly ONE owned n x k copy + // (the equilibrated faer matrix below), and drop the SVD factors before + // the fitted/residual/vcov stage so U never coexists with the vcov + // scores block. Verified via the alloc-profile counting allocator. // Column equilibration: scale each column to unit 2-norm before the SVD so // rank detection (threshold = s_max * rcond, anchored to the largest singular @@ -71,68 +75,78 @@ pub fn solve_ols<'py>( // unscaled back to raw scale below, BEFORE fitted/residuals/vcov, so all // raw-scale quantities (x_arr) stay consistent. Mirrors the Python backend's // _detect_rank_deficiency / _equilibrated_lstsq equilibration. + // Norms are computed from the borrowed view in the same j-outer/i-inner + // accumulation order as before (bit-identical to the prior owned-copy scan). let mut safe_norms = Array1::::zeros(k); for j in 0..k { let mut acc = 0.0_f64; for i in 0..n { - let v = x_owned[[i, j]]; + let v = x_arr[[i, j]]; acc += v * v; } let norm = acc.sqrt(); safe_norms[j] = if norm > 0.0 { norm } else { 1.0 }; } - let mut x_scaled = x_owned.clone(); - for j in 0..k { - let s = safe_norms[j]; - for i in 0..n { - x_scaled[[i, j]] /= s; - } - } - // Convert ndarray to faer for SVD computation (on the equilibrated matrix) - let x_faer = ndarray_to_faer(&x_scaled); + // Equilibration fused into the faer conversion: the single owned n x k + // copy. Same per-element `x / norm` as the previous two-step + // clone-then-divide, so the SVD input values are unchanged. + let x_faer = faer::Mat::from_fn(n, k, |i, j| x_arr[[i, j]] / safe_norms[j]); // Compute thin SVD using faer: X = U * S * V^T let svd = match x_faer.thin_svd() { Ok(s) => s, Err(_) => { return Err(PyErr::new::( - "SVD computation failed" + "SVD computation failed", )) } }; + // The equilibrated input is not needed once the factors exist. + drop(x_faer); - // Extract U, S, V from faer SVD result (capitalized methods in faer 0.24) + // Extract from the faer SVD result (capitalized methods in faer 0.24) + // only the small pieces the solve needs - s (min(n,k)), V^T (k x k), + // and U^T y (min(n,k)) - then drop the factorization, releasing the + // n x k U storage before any downstream allocation. let u_faer = svd.U(); - let s_diag = svd.S(); // Returns diagonal view - let s_col = s_diag.column_vector(); // Get as column vector - let v_faer = svd.V(); // This is V, not V^T + let s_diag = svd.S(); // Returns diagonal view + let s_col = s_diag.column_vector(); // Get as column vector + let v_faer = svd.V(); // This is V, not V^T - // Convert back to ndarray let n_rows = u_faer.nrows(); let n_svd_cols = u_faer.ncols(); - let mut u = Array2::::zeros((n_rows, n_svd_cols)); - for i in 0..n_rows { - for j in 0..n_svd_cols { - u[[i, j]] = u_faer[(i, j)]; - } - } let s_len = s_col.nrows(); let mut s = Array1::::zeros(s_len); for i in 0..s_len { - s[i] = s_col[i]; // S column vector + s[i] = s_col[i]; // S column vector } let v_rows = v_faer.nrows(); let v_cols = v_faer.ncols(); - let mut vt = Array2::::zeros((v_cols, v_rows)); // V^T has shape (k, k) + let mut vt = Array2::::zeros((v_cols, v_rows)); // V^T has shape (k, k) for i in 0..v_rows { for j in 0..v_cols { - vt[[j, i]] = v_faer[(i, j)]; // Transpose V to get V^T + vt[[j, i]] = v_faer[(i, j)]; // Transpose V to get V^T } } + // U^T y computed directly off the faer factor - never materializes an + // ndarray copy of U. Sequential column-order accumulation (U is + // col-major): deterministic bits independent of BLAS/thread count. + let mut uty = Array1::::zeros(n_svd_cols); // (min(n,k),) + for j in 0..n_svd_cols { + let mut acc = 0.0_f64; + for i in 0..n_rows { + acc += u_faer[(i, j)] * y_arr[i]; + } + uty[j] = acc; + } + // Everything extracted; release the factorization (frees U's n x k + // storage) before fitted/residuals/vcov allocate. + drop(svd); + // Compute rcond threshold to match R's lm() behavior // R's qr() uses tol = 1e-07 by default, which is sqrt(eps) ≈ 1.49e-08 // We use 1e-07 for consistency with Python backend and R @@ -140,9 +154,8 @@ pub fn solve_ols<'py>( let s_max = s.iter().cloned().fold(0.0_f64, f64::max); let threshold = s_max * rcond; - // Compute truncated pseudoinverse solution: β = V * S^{-1} * U^T * y + // Truncated pseudoinverse solution: β = V * S^{-1} * (U^T y). // Singular values below threshold are treated as zero (truncated) - let uty = u.t().dot(&y_owned); // (min(n,k),) // Build S^{-1} with truncation and count effective rank // Note: s.len() = min(n, k) from thin SVD, so this handles underdetermined (n < k) correctly @@ -180,18 +193,20 @@ pub fn solve_ols<'py>( } else { // Full rank: compute robust vcov normally let cluster_arr = cluster_ids.as_ref().map(|c| c.as_array().to_owned()); - let vcov_arr = compute_robust_vcov_internal(&x_arr, &residuals.view(), cluster_arr.as_ref(), n, k)?; + let vcov_arr = compute_robust_vcov_internal( + &x_arr, + &residuals.view(), + cluster_arr.as_ref(), + n, + k, + )?; Some(vcov_arr.to_pyarray(py)) } } else { None }; - Ok(( - coefficients.to_pyarray(py), - residuals.to_pyarray(py), - vcov, - )) + Ok((coefficients.to_pyarray(py), residuals.to_pyarray(py), vcov)) } /// Compute HC1 or cluster-robust variance-covariance matrix. @@ -280,9 +295,10 @@ fn compute_robust_vcov_internal( let n_clusters = cluster_sums.len(); if n_clusters < 2 { - return Err(PyErr::new::( - format!("Need at least 2 clusters for cluster-robust SEs, got {}", n_clusters) - )); + return Err(PyErr::new::(format!( + "Need at least 2 clusters for cluster-robust SEs, got {}", + n_clusters + ))); } // Build cluster scores matrix (G, k) @@ -355,7 +371,7 @@ fn invert_symmetric(a: &Array2) -> PyResult> { if has_nan { return Err(PyErr::new::( - "Matrix inversion failed (singular matrix)" + "Matrix inversion failed (singular matrix)", )); } @@ -372,7 +388,11 @@ fn invert_symmetric(a: &Array2) -> PyResult> { min_pivot = min_pivot.min(pivot); } } - let pivot_ratio = if max_pivot > 0.0 { min_pivot / max_pivot } else { 0.0 }; + let pivot_ratio = if max_pivot > 0.0 { + min_pivot / max_pivot + } else { + 0.0 + }; // Only perform expensive residual check if pivots suggest potential instability. // Threshold of 1e-10 catches truly problematic matrices while avoiding @@ -396,13 +416,11 @@ fn invert_symmetric(a: &Array2) -> PyResult> { // while still producing usable results. Use 1e-4 * n as threshold. let threshold = 1e-4 * (n as f64); if max_residual > threshold { - return Err(PyErr::new::( - format!( - "Matrix inversion numerically unstable (residual={:.2e} > threshold={:.2e}). \ + return Err(PyErr::new::(format!( + "Matrix inversion numerically unstable (residual={:.2e} > threshold={:.2e}). \ Design matrix may be near-singular.", - max_residual, threshold - ) - )); + max_residual, threshold + ))); } } @@ -461,7 +479,11 @@ mod tests { // For n=2 < k=3: U is (2, 2), S has 2 values, V is (3, 2) assert_eq!(svd.U().nrows(), 2, "U should have n=2 rows"); assert_eq!(svd.U().ncols(), 2, "U should have min(n,k)=2 cols"); - assert_eq!(svd.S().column_vector().nrows(), 2, "S should have min(n,k)=2 singular values"); + assert_eq!( + svd.S().column_vector().nrows(), + 2, + "S should have min(n,k)=2 singular values" + ); assert_eq!(svd.V().nrows(), 3, "V should have k=3 rows"); assert_eq!(svd.V().ncols(), 2, "V should have min(n,k)=2 cols"); diff --git a/rust/src/trop.rs b/rust/src/trop.rs index 28cc2392..abd8bd74 100644 --- a/rust/src/trop.rs +++ b/rust/src/trop.rs @@ -52,10 +52,7 @@ pub fn compute_unit_distance_matrix<'py>( /// Internal implementation of unit distance matrix computation. /// /// Parallelizes over unit pairs using rayon. -fn compute_unit_distance_matrix_internal( - y: &ArrayView2, - d: &ArrayView2, -) -> Array2 { +fn compute_unit_distance_matrix_internal(y: &ArrayView2, d: &ArrayView2) -> Array2 { let n_periods = y.nrows(); let n_units = y.ncols(); @@ -195,29 +192,46 @@ fn univariate_loocv_search( let (lambda_time, lambda_unit, lambda_nn) = match param_type { 0 => { // Searching λ_time: use grid value directly (no inf expected) - (value, - fixed_unit, - if fixed_nn.is_infinite() { 1e10 } else { fixed_nn }) - }, + ( + value, + fixed_unit, + if fixed_nn.is_infinite() { + 1e10 + } else { + fixed_nn + }, + ) + } 1 => { // Searching λ_unit: use grid value directly (no inf expected) - (fixed_time, - value, - if fixed_nn.is_infinite() { 1e10 } else { fixed_nn }) - }, + ( + fixed_time, + value, + if fixed_nn.is_infinite() { + 1e10 + } else { + fixed_nn + }, + ) + } _ => { // Searching λ_nn: convert inf → 1e10 (factor model disabled) let value_converted = if value.is_infinite() { 1e10 } else { value }; - (fixed_time, - fixed_unit, - value_converted) - }, + (fixed_time, fixed_unit, value_converted) + } }; let (score, _, _) = loocv_score_for_params( - y, d, control_mask, time_dist, control_obs, - lambda_time, lambda_unit, lambda_nn, - max_iter, tol, + y, + d, + control_mask, + time_dist, + control_obs, + lambda_time, + lambda_unit, + lambda_nn, + max_iter, + tol, ); (value, score) }) @@ -260,22 +274,52 @@ fn cycling_parameter_search( for _cycle in 0..max_cycles { // Optimize λ_unit (fix λ_time, λ_nn) let (new_unit, _) = univariate_loocv_search( - y, d, control_mask, time_dist, control_obs, - lambda_unit_grid, lambda_time, 0.0, lambda_nn, 1, max_iter, tol, + y, + d, + control_mask, + time_dist, + control_obs, + lambda_unit_grid, + lambda_time, + 0.0, + lambda_nn, + 1, + max_iter, + tol, ); lambda_unit = new_unit; // Optimize λ_time (fix λ_unit, λ_nn) let (new_time, _) = univariate_loocv_search( - y, d, control_mask, time_dist, control_obs, - lambda_time_grid, 0.0, lambda_unit, lambda_nn, 0, max_iter, tol, + y, + d, + control_mask, + time_dist, + control_obs, + lambda_time_grid, + 0.0, + lambda_unit, + lambda_nn, + 0, + max_iter, + tol, ); lambda_time = new_time; // Optimize λ_nn (fix λ_unit, λ_time) let (new_nn, score) = univariate_loocv_search( - y, d, control_mask, time_dist, control_obs, - lambda_nn_grid, lambda_time, lambda_unit, 0.0, 2, max_iter, tol, + y, + d, + control_mask, + time_dist, + control_obs, + lambda_nn_grid, + lambda_time, + lambda_unit, + 0.0, + 2, + max_iter, + tol, ); lambda_nn = new_nn; @@ -355,38 +399,75 @@ pub fn loocv_grid_search<'py>( } // Get control observations for LOOCV - let control_obs = get_control_observations( - &y_arr, - &control_mask_arr, - ); + let control_obs = get_control_observations(&y_arr, &control_mask_arr); let n_attempted = control_obs.len(); // Stage 1: Univariate searches for initial values (paper footnote 2) // λ_time search: fix λ_unit=0, λ_nn=∞ (disabled) let (lambda_time_init, _) = univariate_loocv_search( - &y_arr, &d_arr, &control_mask_arr, &time_dist_arr, &control_obs, - &lambda_time_vec, 0.0, 0.0, f64::INFINITY, 0, max_iter, tol, + &y_arr, + &d_arr, + &control_mask_arr, + &time_dist_arr, + &control_obs, + &lambda_time_vec, + 0.0, + 0.0, + f64::INFINITY, + 0, + max_iter, + tol, ); // λ_nn search: fix λ_time=0 (uniform time weights), λ_unit=0 let (lambda_nn_init, _) = univariate_loocv_search( - &y_arr, &d_arr, &control_mask_arr, &time_dist_arr, &control_obs, - &lambda_nn_vec, 0.0, 0.0, 0.0, 2, max_iter, tol, + &y_arr, + &d_arr, + &control_mask_arr, + &time_dist_arr, + &control_obs, + &lambda_nn_vec, + 0.0, + 0.0, + 0.0, + 2, + max_iter, + tol, ); // λ_unit search: fix λ_nn=∞, λ_time=0 let (lambda_unit_init, _) = univariate_loocv_search( - &y_arr, &d_arr, &control_mask_arr, &time_dist_arr, &control_obs, - &lambda_unit_vec, 0.0, 0.0, f64::INFINITY, 1, max_iter, tol, + &y_arr, + &d_arr, + &control_mask_arr, + &time_dist_arr, + &control_obs, + &lambda_unit_vec, + 0.0, + 0.0, + f64::INFINITY, + 1, + max_iter, + tol, ); // Stage 2: Cycling refinement let (best_time, best_unit, best_nn) = cycling_parameter_search( - &y_arr, &d_arr, &control_mask_arr, &time_dist_arr, &control_obs, - &lambda_time_vec, &lambda_unit_vec, &lambda_nn_vec, - lambda_time_init, lambda_unit_init, lambda_nn_init, - max_iter, tol, 10, + &y_arr, + &d_arr, + &control_mask_arr, + &time_dist_arr, + &control_obs, + &lambda_time_vec, + &lambda_unit_vec, + &lambda_nn_vec, + lambda_time_init, + lambda_unit_init, + lambda_nn_init, + max_iter, + tol, + 10, ); // Convert λ_nn=∞ → 1e10 for final score computation (factor model disabled) @@ -396,13 +477,28 @@ pub fn loocv_grid_search<'py>( // Compute final score with converted values let (best_score, n_valid, first_failed) = loocv_score_for_params( - &y_arr, &d_arr, &control_mask_arr, &time_dist_arr, &control_obs, - best_time_eff, best_unit_eff, best_nn_eff, - max_iter, tol, + &y_arr, + &d_arr, + &control_mask_arr, + &time_dist_arr, + &control_obs, + best_time_eff, + best_unit_eff, + best_nn_eff, + max_iter, + tol, ); // Return ORIGINAL grid values (for user visibility) but score computed with converted - Ok((best_time, best_unit, best_nn, best_score, n_valid, n_attempted, first_failed)) + Ok(( + best_time, + best_unit, + best_nn, + best_score, + n_valid, + n_attempted, + first_failed, + )) } /// Get all valid control observations for LOOCV. @@ -526,9 +622,7 @@ fn compute_unit_distance_for_obs( continue; } // Both units must be control at this period and have valid values - if d[[t, i]] == 0.0 && d[[t, j]] == 0.0 - && y[[t, i]].is_finite() && y[[t, j]].is_finite() - { + if d[[t, i]] == 0.0 && d[[t, j]] == 0.0 && y[[t, i]].is_finite() && y[[t, j]].is_finite() { let diff = y[[t, i]] - y[[t, j]]; sum_sq += diff * diff; n_valid += 1; @@ -632,9 +726,8 @@ fn estimate_model( exclude_obs: Option<(usize, usize)>, ) -> Option<(Array1, Array1, Array2)> { // Create estimation mask - let mut est_mask = Array2::::from_shape_fn((n_periods, n_units), |(t, i)| { - control_mask[[t, i]] != 0 - }); + let mut est_mask = + Array2::::from_shape_fn((n_periods, n_units), |(t, i)| control_mask[[t, i]] != 0); if let Some((t_ex, i_ex)) = exclude_obs { est_mask[[t_ex, i_ex]] = false; @@ -656,7 +749,11 @@ fn estimate_model( // Lipschitz constant of ∇f is L_f = 2·max(W), so prox threshold = λ/(2·max(W)) let w_max = w_masked.iter().cloned().fold(0.0_f64, f64::max); - let prox_threshold = if w_max > 0.0 { lambda_nn / (2.0 * w_max) } else { lambda_nn / 2.0 }; + let prox_threshold = if w_max > 0.0 { + lambda_nn / (2.0 * w_max) + } else { + lambda_nn / 2.0 + }; // Weight sums per unit and time let weight_sum_per_unit: Array1 = w_masked.sum_axis(Axis(0)); @@ -730,12 +827,20 @@ fn estimate_model( // For W=0 cells, use current L instead of R (prevent absorbing treatment) let r_masked = Array2::from_shape_fn((n_periods, n_units), |(t, i)| { - if w_masked[[t, i]] > 0.0 { r_target[[t, i]] } else { l[[t, i]] } + if w_masked[[t, i]] > 0.0 { + r_target[[t, i]] + } else { + l[[t, i]] + } }); // Normalize weights: W_norm = W / W_max (max becomes 1) let w_norm = Array2::from_shape_fn((n_periods, n_units), |(t, i)| { - if w_max > 0.0 { w_masked[[t, i]] / w_max } else { w_masked[[t, i]] } + if w_max > 0.0 { + w_masked[[t, i]] / w_max + } else { + w_masked[[t, i]] + } }); // FISTA inner loop for L update @@ -757,7 +862,8 @@ fn estimate_model( let mut gradient_step = Array2::::zeros((n_periods, n_units)); for t in 0..n_periods { for i in 0..n_units { - gradient_step[[t, i]] = l_momentum[[t, i]] + w_norm[[t, i]] * (r_masked[[t, i]] - l_momentum[[t, i]]); + gradient_step[[t, i]] = l_momentum[[t, i]] + + w_norm[[t, i]] * (r_masked[[t, i]] - l_momentum[[t, i]]); } } @@ -810,9 +916,9 @@ fn soft_threshold_svd(m: &Array2, threshold: f64) -> Option> { }; let u_faer = svd.U(); - let s_diag = svd.S(); // Returns diagonal view - let s_col = s_diag.column_vector(); // Get as column vector - let v_faer = svd.V(); // This is V, not V^T + let s_diag = svd.S(); // Returns diagonal view + let s_col = s_diag.column_vector(); // Get as column vector + let v_faer = svd.V(); // This is V, not V^T let s_len = s_col.nrows(); @@ -1457,14 +1563,26 @@ fn solve_joint_with_lowrank( // Precompute normalized weights and threshold (constant across iterations) let delta_max = delta.iter().cloned().fold(0.0_f64, f64::max); - let threshold = if delta_max > 0.0 { lambda_nn / (2.0 * delta_max) } else { lambda_nn / 2.0 }; + let threshold = if delta_max > 0.0 { + lambda_nn / (2.0 * delta_max) + } else { + lambda_nn / 2.0 + }; // Precompute delta_norm (masked for NaN outcomes) let mut delta_norm = Array2::::zeros((n_periods, n_units)); for t in 0..n_periods { for i in 0..n_units { - let d_ti = if y[[t, i]].is_finite() { delta[[t, i]] } else { 0.0 }; - delta_norm[[t, i]] = if delta_max > 0.0 { d_ti / delta_max } else { d_ti }; + let d_ti = if y[[t, i]].is_finite() { + delta[[t, i]] + } else { + 0.0 + }; + delta_norm[[t, i]] = if delta_max > 0.0 { + d_ti / delta_max + } else { + d_ti + }; } } @@ -1476,7 +1594,7 @@ fn solve_joint_with_lowrank( // Step 1: Fix L, solve for (mu, alpha, beta) let y_adj = Array2::from_shape_fn((n_periods, n_units), |(t, i)| { - y[[t, i]] - l[[t, i]] // NaN - finite = NaN (preserves NaN info) + y[[t, i]] - l[[t, i]] // NaN - finite = NaN (preserves NaN info) }); let (mu, alpha, beta) = solve_joint_no_lowrank(&y_adj.view(), delta)?; @@ -1508,7 +1626,8 @@ fn solve_joint_with_lowrank( let mut gradient_step = Array2::::zeros((n_periods, n_units)); for t in 0..n_periods { for i in 0..n_units { - let l_mom = l_inner[[t, i]] + momentum * (l_inner[[t, i]] - l_inner_prev[[t, i]]); + let l_mom = + l_inner[[t, i]] + momentum * (l_inner[[t, i]] - l_inner_prev[[t, i]]); gradient_step[[t, i]] = l_mom + delta_norm[[t, i]] * (r_masked[[t, i]] - l_mom); } } @@ -1536,9 +1655,7 @@ fn solve_joint_with_lowrank( } // Final solve with converged L - let y_adj = Array2::from_shape_fn((n_periods, n_units), |(t, i)| { - y[[t, i]] - l[[t, i]] - }); + let y_adj = Array2::from_shape_fn((n_periods, n_units), |(t, i)| y[[t, i]] - l[[t, i]]); let (mu, alpha, beta) = solve_joint_no_lowrank(&y_adj.view(), delta)?; Some((mu, alpha, beta, l)) @@ -1585,11 +1702,10 @@ fn loocv_score_joint( delta_ex[[t_ex, i_ex]] = 0.0; let result = if lambda_nn >= 1e10 { - solve_joint_no_lowrank(y, &delta_ex.view()) - .map(|(mu, alpha, beta)| { - let l = Array2::::zeros((n_periods, n_units)); - (mu, alpha, beta, l) - }) + solve_joint_no_lowrank(y, &delta_ex.view()).map(|(mu, alpha, beta)| { + let l = Array2::::zeros((n_periods, n_units)); + (mu, alpha, beta, l) + }) } else { solve_joint_with_lowrank(y, &delta_ex.view(), lambda_nn, max_iter, tol) }; @@ -1597,7 +1713,8 @@ fn loocv_score_joint( match result { Some((mu, alpha, beta, l)) => { if y[[t_ex, i_ex]].is_finite() { - let tau_loocv = y[[t_ex, i_ex]] - mu - alpha[i_ex] - beta[t_ex] - l[[t_ex, i_ex]]; + let tau_loocv = + y[[t_ex, i_ex]] - mu - alpha[i_ex] - beta[t_ex] - l[[t_ex, i_ex]]; (sum + tau_loocv * tau_loocv, valid + 1, first_fail) } else { (sum, valid, first_fail) @@ -1749,7 +1866,15 @@ pub fn loocv_grid_search_global<'py>( let (best_lt, best_lu, best_ln, best_score, n_valid, first_failed) = best_result; - Ok((best_lt, best_lu, best_ln, best_score, n_valid, n_attempted, first_failed)) + Ok(( + best_lt, + best_lu, + best_ln, + best_score, + n_valid, + n_attempted, + first_failed, + )) } /// Compute bootstrap variance estimation for TROP global method in parallel. @@ -1883,7 +2008,11 @@ pub fn bootstrap_trop_variance_global<'py>( let treated_periods = n_periods.saturating_sub(first_treat_period); // Convert λ_nn=∞ → 1e10 (factor model disabled) - let ln_eff = if lambda_nn.is_infinite() { 1e10 } else { lambda_nn }; + let ln_eff = if lambda_nn.is_infinite() { + 1e10 + } else { + lambda_nn + }; // Run bootstrap iterations in parallel // RNG-canonical contract: control_indices and treated_indices are pre-generated @@ -1925,19 +2054,12 @@ pub fn bootstrap_trop_variance_global<'py>( ); let result = if ln_eff >= 1e10 { - solve_joint_no_lowrank(&y_boot.view(), &delta.view()) - .map(|(mu, alpha, beta)| { - let l = Array2::::zeros((n_periods, n_units)); - (mu, alpha, beta, l) - }) + solve_joint_no_lowrank(&y_boot.view(), &delta.view()).map(|(mu, alpha, beta)| { + let l = Array2::::zeros((n_periods, n_units)); + (mu, alpha, beta, l) + }) } else { - solve_joint_with_lowrank( - &y_boot.view(), - &delta.view(), - ln_eff, - max_iter, - tol, - ) + solve_joint_with_lowrank(&y_boot.view(), &delta.view(), ln_eff, max_iter, tol) }; // Post-hoc tau extraction: ATT = mean(Y - mu - alpha - beta - L) over treated @@ -2000,7 +2122,8 @@ mod tests { let valid_j = array![true, true, true, true]; let valid_i = array![true, true, true, true]; - let dist = compute_pair_distance(&y_j.view(), &y_i.view(), &valid_j.view(), &valid_i.view()); + let dist = + compute_pair_distance(&y_j.view(), &y_i.view(), &valid_j.view(), &valid_i.view()); // RMSE of constant difference 0.5 should be 0.5 assert!((dist - 0.5).abs() < 1e-10); @@ -2014,7 +2137,8 @@ mod tests { let valid_i = array![true, false, true, false]; // Only period 0 overlaps - let dist = compute_pair_distance(&y_j.view(), &y_i.view(), &valid_j.view(), &valid_i.view()); + let dist = + compute_pair_distance(&y_j.view(), &y_i.view(), &valid_j.view(), &valid_i.view()); // RMSE of single difference 0.5 should be 0.5 assert!((dist - 0.5).abs() < 1e-10); @@ -2027,7 +2151,8 @@ mod tests { let valid_j = array![true, true, false, false]; let valid_i = array![false, false, true, true]; - let dist = compute_pair_distance(&y_j.view(), &y_i.view(), &valid_j.view(), &valid_i.view()); + let dist = + compute_pair_distance(&y_j.view(), &y_i.view(), &valid_j.view(), &valid_i.view()); assert!(dist.is_infinite()); } diff --git a/rust/src/weights.rs b/rust/src/weights.rs index 196999c7..c4468beb 100644 --- a/rust/src/weights.rs +++ b/rust/src/weights.rs @@ -5,8 +5,8 @@ //! - Simplex projection //! - SDID unit and time weight computation -use ndarray::{s, Array1, Array2, ArrayView1, ArrayView2, Axis}; use ndarray::linalg::general_mat_vec_mul; +use ndarray::{s, Array1, Array2, ArrayView1, ArrayView2, Axis}; use numpy::{PyArray1, PyReadonlyArray1, PyReadonlyArray2, ToPyArray}; use pyo3::prelude::*; @@ -167,7 +167,8 @@ fn sc_weight_fw_gram( // Already at optimal vertex — compute objective for convergence check let xt_ata_x: f64 = ata_x.iter().zip(lam.iter()).map(|(&a, &b)| a * b).sum(); let atb_dot_lam: f64 = atb.iter().zip(lam.iter()).map(|(&a, &b)| a * b).sum(); - let val = zeta * zeta * lam_norm_sq + (xt_ata_x - 2.0 * atb_dot_lam + b_norm_sq) / n as f64; + let val = + zeta * zeta * lam_norm_sq + (xt_ata_x - 2.0 * atb_dot_lam + b_norm_sq) / n as f64; if t >= 1 && prev_val - val < min_decrease_sq { converged = true; break; @@ -183,7 +184,8 @@ fn sc_weight_fw_gram( let denom = d_err_sq + eta * d_x_norm_sq; if denom <= 0.0 { let atb_dot_lam: f64 = atb.iter().zip(lam.iter()).map(|(&a, &b)| a * b).sum(); - let val = zeta * zeta * lam_norm_sq + (xt_ata_x - 2.0 * atb_dot_lam + b_norm_sq) / n as f64; + let val = + zeta * zeta * lam_norm_sq + (xt_ata_x - 2.0 * atb_dot_lam + b_norm_sq) / n as f64; if t >= 1 && prev_val - val < min_decrease_sq { converged = true; break; @@ -461,8 +463,7 @@ fn sc_weight_fw_gram_weighted( } // Recompute weighted lam-norm after the in-place update - let lam_rw_norm_sq_new: f64 = - lam.iter().zip(reg_w.iter()).map(|(&l, &w)| w * l * l).sum(); + let lam_rw_norm_sq_new: f64 = lam.iter().zip(reg_w.iter()).map(|(&l, &w)| w * l * l).sum(); let xt_ata_x: f64 = ata_x.iter().zip(lam.iter()).map(|(&a, &b)| a * b).sum(); let atb_dot_lam: f64 = atb.iter().zip(lam.iter()).map(|(&a, &b)| a * b).sum(); let val = zeta * zeta * lam_rw_norm_sq_new @@ -576,8 +577,7 @@ fn sc_weight_fw_standard_weighted( let e = ax[k] - b[k]; err_sq += e * e; } - let lam_rw_norm_sq_new: f64 = - lam.iter().zip(reg_w.iter()).map(|(&l, &w)| w * l * l).sum(); + let lam_rw_norm_sq_new: f64 = lam.iter().zip(reg_w.iter()).map(|(&l, &w)| w * l * l).sum(); let val = zeta * zeta * lam_rw_norm_sq_new + err_sq / n as f64; if t >= 1 && prev_val - val < min_decrease_sq { @@ -722,11 +722,27 @@ fn sc_weight_fw_weighted_internal( let converged = if t0 < n { sc_weight_fw_gram_weighted( - &a, &b, &mut lam, &rw_view, eta, zeta, n, min_decrease_sq, max_iter, + &a, + &b, + &mut lam, + &rw_view, + eta, + zeta, + n, + min_decrease_sq, + max_iter, ) } else { sc_weight_fw_standard_weighted( - &a, &b, &mut lam, &rw_view, eta, zeta, n, min_decrease_sq, max_iter, + &a, + &b, + &mut lam, + &rw_view, + eta, + zeta, + n, + min_decrease_sq, + max_iter, ) }; @@ -990,7 +1006,15 @@ pub fn compute_time_weights<'py>( let y_pre = y_pre_control.as_array(); let y_post = y_post_control.as_array(); - let result = compute_time_weights_internal(&y_pre, &y_post, zeta_lambda, intercept, min_decrease, max_iter_pre_sparsify, max_iter); + let result = compute_time_weights_internal( + &y_pre, + &y_post, + zeta_lambda, + intercept, + min_decrease, + max_iter_pre_sparsify, + max_iter, + ); Ok(result.to_pyarray(py)) } @@ -1037,13 +1061,27 @@ pub(crate) fn compute_time_weights_internal( // `compute_time_weights` (Python wrapper in utils.py) with // `return_convergence=True`, which runs the two-pass in Python against // `sc_weight_fw_with_convergence`. - let (lam, _) = sc_weight_fw_internal(&y_time.view(), zeta_lambda, intercept, None, min_decrease, max_iter_pre_sparsify); + let (lam, _) = sc_weight_fw_internal( + &y_time.view(), + zeta_lambda, + intercept, + None, + min_decrease, + max_iter_pre_sparsify, + ); // Sparsify let lam_sparse = sparsify_internal(&lam); // Second pass: from sparsified initialization - let (lam2, _) = sc_weight_fw_internal(&y_time.view(), zeta_lambda, intercept, Some(&lam_sparse), min_decrease, max_iter); + let (lam2, _) = sc_weight_fw_internal( + &y_time.view(), + zeta_lambda, + intercept, + Some(&lam_sparse), + min_decrease, + max_iter, + ); lam2 } @@ -1076,8 +1114,13 @@ pub fn compute_sdid_unit_weights<'py>( let y_tr_mean = y_pre_treated_mean.as_array(); let result = compute_sdid_unit_weights_internal( - &y_pre, &y_tr_mean, zeta_omega, intercept, min_decrease, - max_iter_pre_sparsify, max_iter, + &y_pre, + &y_tr_mean, + zeta_omega, + intercept, + min_decrease, + max_iter_pre_sparsify, + max_iter, ); Ok(result.to_pyarray(py)) } @@ -1114,7 +1157,12 @@ pub(crate) fn compute_sdid_unit_weights_internal( // First pass: limited iterations. See note in compute_time_weights_internal // about convergence-tracking contract. let (omega, _) = sc_weight_fw_internal( - &y_unit.view(), zeta_omega, intercept, None, min_decrease, max_iter_pre_sparsify, + &y_unit.view(), + zeta_omega, + intercept, + None, + min_decrease, + max_iter_pre_sparsify, ); // Sparsify: zero out weights <= max/4, renormalize @@ -1122,7 +1170,12 @@ pub(crate) fn compute_sdid_unit_weights_internal( // Second pass: from sparsified initialization let (omega2, _) = sc_weight_fw_internal( - &y_unit.view(), zeta_omega, intercept, Some(&omega), min_decrease, max_iter, + &y_unit.view(), + zeta_omega, + intercept, + Some(&omega), + min_decrease, + max_iter, ); omega2 } @@ -1221,19 +1274,41 @@ mod tests { let y = array![[1.0, 2.0, 1.5], [3.0, 4.0, 3.5], [5.0, 6.0, 5.5]]; let (result, _converged) = sc_weight_fw_internal(&y.view(), 0.1, true, None, 1e-3, 100); let sum: f64 = result.sum(); - assert!((sum - 1.0).abs() < 1e-6, "FW weights should sum to 1, got {}", sum); - assert!(result.iter().all(|&w| w >= -1e-6), "FW weights should be non-negative"); + assert!( + (sum - 1.0).abs() < 1e-6, + "FW weights should sum to 1, got {}", + sum + ); + assert!( + result.iter().all(|&w| w >= -1e-6), + "FW weights should be non-negative" + ); } #[test] fn test_time_weights_on_simplex() { let y_pre = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]; let y_post = array![[10.0, 11.0, 12.0]]; - let result = compute_time_weights_internal(&y_pre.view(), &y_post.view(), 0.1, true, 1e-3, 100, 1000); + let result = compute_time_weights_internal( + &y_pre.view(), + &y_post.view(), + 0.1, + true, + 1e-3, + 100, + 1000, + ); assert_eq!(result.len(), 3); let sum: f64 = result.sum(); - assert!((sum - 1.0).abs() < 1e-6, "Time weights should sum to 1, got {}", sum); - assert!(result.iter().all(|&w| w >= -1e-6), "Time weights should be non-negative"); + assert!( + (sum - 1.0).abs() < 1e-6, + "Time weights should sum to 1, got {}", + sum + ); + assert!( + result.iter().all(|&w| w >= -1e-6), + "Time weights should be non-negative" + ); } #[test] @@ -1241,12 +1316,25 @@ mod tests { let y_pre = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]; let y_tr_mean = array![2.0, 5.0, 8.0]; let result = compute_sdid_unit_weights_internal( - &y_pre.view(), &y_tr_mean.view(), 0.5, true, 1e-3, 100, 1000, + &y_pre.view(), + &y_tr_mean.view(), + 0.5, + true, + 1e-3, + 100, + 1000, ); assert_eq!(result.len(), 3); let sum: f64 = result.sum(); - assert!((sum - 1.0).abs() < 1e-6, "Unit weights should sum to 1, got {}", sum); - assert!(result.iter().all(|&w| w >= -1e-6), "Unit weights should be non-negative"); + assert!( + (sum - 1.0).abs() < 1e-6, + "Unit weights should sum to 1, got {}", + sum + ); + assert!( + result.iter().all(|&w| w >= -1e-6), + "Unit weights should be non-negative" + ); } #[test] @@ -1254,7 +1342,13 @@ mod tests { let y_pre = array![[1.0], [2.0], [3.0]]; let y_tr_mean = array![1.5, 2.5, 3.5]; let result = compute_sdid_unit_weights_internal( - &y_pre.view(), &y_tr_mean.view(), 0.5, true, 1e-3, 100, 1000, + &y_pre.view(), + &y_tr_mean.view(), + 0.5, + true, + 1e-3, + 100, + 1000, ); assert_eq!(result.len(), 1); assert!((result[0] - 1.0).abs() < 1e-10); @@ -1292,15 +1386,25 @@ mod tests { assert!( (lam_gram[j] - lam_std[j]).abs() < 1e-10, "Gram and standard paths diverge at index {}: gram={}, std={}", - j, lam_gram[j], lam_std[j] + j, + lam_gram[j], + lam_std[j] ); } // Verify both are valid simplex weights let sum_gram: f64 = lam_gram.sum(); let sum_std: f64 = lam_std.sum(); - assert!((sum_gram - 1.0).abs() < 1e-6, "Gram weights should sum to 1, got {}", sum_gram); - assert!((sum_std - 1.0).abs() < 1e-6, "Standard weights should sum to 1, got {}", sum_std); + assert!( + (sum_gram - 1.0).abs() < 1e-6, + "Gram weights should sum to 1, got {}", + sum_gram + ); + assert!( + (sum_std - 1.0).abs() < 1e-6, + "Standard weights should sum to 1, got {}", + sum_std + ); } #[test] @@ -1314,8 +1418,15 @@ mod tests { // Verify valid simplex weights let sum: f64 = result.sum(); - assert!((sum - 1.0).abs() < 1e-6, "Weights should sum to 1, got {}", sum); - assert!(result.iter().all(|&w| w >= -1e-6), "Weights should be non-negative"); + assert!( + (sum - 1.0).abs() < 1e-6, + "Weights should sum to 1, got {}", + sum + ); + assert!( + result.iter().all(|&w| w >= -1e-6), + "Weights should be non-negative" + ); assert_eq!(result.len(), 8); } @@ -1323,10 +1434,12 @@ mod tests { fn test_incremental_ata_x_accuracy() { // Run 500+ iterations on a T0 < N problem and verify incremental ata_x // doesn't drift significantly from fresh computation. - let vals: Vec = (0..200).map(|i| { - let x = (i as f64) * 0.1; - x.sin() + ((i * 7) % 31) as f64 / 31.0 - }).collect(); + let vals: Vec = (0..200) + .map(|i| { + let x = (i as f64) * 0.1; + x.sin() + ((i * 7) % 31) as f64 / 31.0 + }) + .collect(); let y = Array2::from_shape_vec((20, 10), vals).unwrap(); // Run with enough iterations to exercise the refresh mechanism @@ -1334,8 +1447,15 @@ mod tests { // Verify valid result (convergence with correct weights) let sum: f64 = result.sum(); - assert!((sum - 1.0).abs() < 1e-6, "Weights should sum to 1, got {}", sum); - assert!(result.iter().all(|&w| w >= -1e-6), "Weights should be non-negative"); + assert!( + (sum - 1.0).abs() < 1e-6, + "Weights should sum to 1, got {}", + sum + ); + assert!( + result.iter().all(|&w| w >= -1e-6), + "Weights should be non-negative" + ); // Verify Gram path was used (T0=9 < N=20) assert_eq!(result.len(), 9); @@ -1351,8 +1471,15 @@ mod tests { let (result, _converged) = sc_weight_fw_internal(&y.view(), 0.2, true, None, 1e-5, 10000); let sum: f64 = result.sum(); - assert!((sum - 1.0).abs() < 1e-6, "Weights should sum to 1, got {}", sum); - assert!(result.iter().all(|&w| w >= -1e-6), "Weights should be non-negative"); + assert!( + (sum - 1.0).abs() < 1e-6, + "Weights should sum to 1, got {}", + sum + ); + assert!( + result.iter().all(|&w| w >= -1e-6), + "Weights should be non-negative" + ); assert_eq!(result.len(), 5); // T0 = 5 } @@ -1364,7 +1491,9 @@ mod tests { let eta = 0.5; // Deterministic test data - let a_vals: Vec = (0..(n * t0)).map(|i| ((i * 7 + 3) % 41) as f64 / 41.0).collect(); + let a_vals: Vec = (0..(n * t0)) + .map(|i| ((i * 7 + 3) % 41) as f64 / 41.0) + .collect(); let a = Array2::from_shape_vec((n, t0), a_vals).unwrap(); let ax: Array1 = (0..n).map(|i| ((i * 11 + 5) % 37) as f64 / 37.0).collect(); @@ -1395,7 +1524,9 @@ mod tests { assert!( (ref_grad[j] - new_grad[j]).abs() < 1e-12, "half_grad mismatch at index {}: manual={}, gemv={}", - j, ref_grad[j], new_grad[j] + j, + ref_grad[j], + new_grad[j] ); } } @@ -1406,18 +1537,34 @@ mod tests { // Gram path: N=15, T0=4 (T0 < N) let vals_gram: Vec = (0..75).map(|i| ((i * 3 + 1) % 37) as f64 / 37.0).collect(); let y_gram = Array2::from_shape_vec((15, 5), vals_gram).unwrap(); - let (result_gram, _converged_gram) = sc_weight_fw_internal(&y_gram.view(), 0.3, false, None, 1e-5, 10000); + let (result_gram, _converged_gram) = + sc_weight_fw_internal(&y_gram.view(), 0.3, false, None, 1e-5, 10000); let sum_gram: f64 = result_gram.sum(); - assert!((sum_gram - 1.0).abs() < 1e-6, "Gram intercept=false: weights should sum to 1, got {}", sum_gram); - assert!(result_gram.iter().all(|&w| w >= -1e-6), "Gram intercept=false: weights should be non-negative"); + assert!( + (sum_gram - 1.0).abs() < 1e-6, + "Gram intercept=false: weights should sum to 1, got {}", + sum_gram + ); + assert!( + result_gram.iter().all(|&w| w >= -1e-6), + "Gram intercept=false: weights should be non-negative" + ); // Standard path: N=4, T0=10 (T0 >= N) let vals_std: Vec = (0..44).map(|i| ((i * 5 + 2) % 29) as f64 / 29.0).collect(); let y_std = Array2::from_shape_vec((4, 11), vals_std).unwrap(); - let (result_std, _converged_std) = sc_weight_fw_internal(&y_std.view(), 0.3, false, None, 1e-5, 10000); + let (result_std, _converged_std) = + sc_weight_fw_internal(&y_std.view(), 0.3, false, None, 1e-5, 10000); let sum_std: f64 = result_std.sum(); - assert!((sum_std - 1.0).abs() < 1e-6, "Standard intercept=false: weights should sum to 1, got {}", sum_std); - assert!(result_std.iter().all(|&w| w >= -1e-6), "Standard intercept=false: weights should be non-negative"); + assert!( + (sum_std - 1.0).abs() < 1e-6, + "Standard intercept=false: weights should sum to 1, got {}", + sum_std + ); + assert!( + result_std.iter().all(|&w| w >= -1e-6), + "Standard intercept=false: weights should be non-negative" + ); } // ------------------------------------------------------------------------- @@ -1433,16 +1580,17 @@ mod tests { let (unweighted, conv_unweighted) = sc_weight_fw_internal(&y.view(), 0.3, true, None, 1e-5, 10000); - let (weighted, conv_weighted) = sc_weight_fw_weighted_internal( - &y.view(), 0.3, true, None, 1e-5, 10000, None, - ); + let (weighted, conv_weighted) = + sc_weight_fw_weighted_internal(&y.view(), 0.3, true, None, 1e-5, 10000, None); assert_eq!(conv_unweighted, conv_weighted, "convergence flags differ"); for j in 0..unweighted.len() { assert!( (unweighted[j] - weighted[j]).abs() < 1e-14, "reg_weights=None must delegate to unweighted; mismatch at {}: {} vs {}", - j, unweighted[j], weighted[j] + j, + unweighted[j], + weighted[j] ); } } @@ -1454,20 +1602,23 @@ mod tests { // agree to within machine precision (the weighted loop recomputes // lam_norm via the weighted code path each iteration so float // ordering can introduce tiny ULP-scale drift — stay at rel=1e-12). - let vals: Vec = (0..120).map(|i| ((i * 11 + 5) % 73) as f64 / 73.0).collect(); + let vals: Vec = (0..120) + .map(|i| ((i * 11 + 5) % 73) as f64 / 73.0) + .collect(); let y = Array2::from_shape_vec((20, 6), vals).unwrap(); let rw = Array1::from_elem(5, 1.0); // t0 = ncols - 1 = 5 let (unweighted, _) = sc_weight_fw_internal(&y.view(), 0.3, true, None, 1e-7, 10000); - let (weighted, _) = sc_weight_fw_weighted_internal( - &y.view(), 0.3, true, None, 1e-7, 10000, Some(&rw), - ); + let (weighted, _) = + sc_weight_fw_weighted_internal(&y.view(), 0.3, true, None, 1e-7, 10000, Some(&rw)); for j in 0..unweighted.len() { assert!( (unweighted[j] - weighted[j]).abs() < 1e-12, "uniform reg_weights must match unweighted at {}: {} vs {}", - j, unweighted[j], weighted[j] + j, + unweighted[j], + weighted[j] ); } } @@ -1477,18 +1628,26 @@ mod tests { // For arbitrary positive rw, the returned ω must lie on the standard // simplex (sums to 1, non-negative). Exercise both gram (T0 < N) and // standard (T0 >= N) paths. - let vals_gram: Vec = - (0..120).map(|i| ((i * 13 + 7) % 89) as f64 / 89.0).collect(); + let vals_gram: Vec = (0..120) + .map(|i| ((i * 13 + 7) % 89) as f64 / 89.0) + .collect(); let y_gram = Array2::from_shape_vec((20, 6), vals_gram).unwrap(); let rw_gram = Array1::from_vec(vec![0.5, 1.0, 1.5, 2.0, 0.8]); let (omega_gram, _) = sc_weight_fw_weighted_internal( - &y_gram.view(), 0.4, true, None, 1e-6, 10000, Some(&rw_gram), + &y_gram.view(), + 0.4, + true, + None, + 1e-6, + 10000, + Some(&rw_gram), ); let sum_gram: f64 = omega_gram.sum(); assert!( (sum_gram - 1.0).abs() < 1e-6, - "gram weighted-FW: ω must sum to 1, got {}", sum_gram + "gram weighted-FW: ω must sum to 1, got {}", + sum_gram ); assert!( omega_gram.iter().all(|&w| w >= -1e-9), @@ -1500,12 +1659,19 @@ mod tests { let rw_std = Array1::from_vec(vec![1.0, 0.5, 1.5, 2.0, 0.7, 1.2, 0.9, 1.3]); let (omega_std, _) = sc_weight_fw_weighted_internal( - &y_std.view(), 0.5, true, None, 1e-6, 10000, Some(&rw_std), + &y_std.view(), + 0.5, + true, + None, + 1e-6, + 10000, + Some(&rw_std), ); let sum_std: f64 = omega_std.sum(); assert!( (sum_std - 1.0).abs() < 1e-6, - "standard weighted-FW: ω must sum to 1, got {}", sum_std + "standard weighted-FW: ω must sum to 1, got {}", + sum_std ); assert!( omega_std.iter().all(|&w| w >= -1e-9), diff --git a/tests/test_linalg.py b/tests/test_linalg.py index caa7e44f..14c75e18 100644 --- a/tests/test_linalg.py +++ b/tests/test_linalg.py @@ -1599,6 +1599,66 @@ def test_rank_detection_mixed_scale_and_collinearity_keeps_identified_subset(sel vd = np.diag(vcov)[kept] assert np.all(np.isfinite(vd)) and np.all(vd >= 0) # valid kept-coef VCV + def test_rank_detection_qr_mode_r_matches_economic(self): + """Lock for the mode="r" refactor: rank detection reads only the R + diagonal and the pivot, and both are bit-identical between + mode="r" and mode="economic" (same dgeqp3 factorization; mode only + controls whether the unused Q is formed). Pivot values are locked by + same-session cross-mode equality, NOT hard-coded (dgeqp3 tie-breaks + are BLAS-dependent); only rank/dropped are asserted as values.""" + from scipy.linalg import qr + + from diff_diff.linalg import _detect_rank_deficiency + + rng = np.random.default_rng(7) + base = rng.standard_normal((80, 3)) + fixtures = { + # genuinely collinear, well-scaled: rank 3 of 4 + "collinear": (np.column_stack([base, base[:, 0] - base[:, 2]]), 3), + # scale artifact: full rank but one huge column + "scale": (np.column_stack([base, 1e9 * rng.standard_normal(80)]), 4), + } + for name, (X, expected_rank) in fixtures.items(): + rank, dropped, pivot = _detect_rank_deficiency(X) + assert rank == expected_rank, name + assert len(dropped) == X.shape[1] - expected_rank, name + # cross-mode equivalence: identical R diagonal and pivot + r_only = qr(X, mode="r", pivoting=True) + _q, r_eco, piv_eco = qr(X, mode="economic", pivoting=True) + np.testing.assert_array_equal( + np.abs(np.diag(r_only[0])), np.abs(np.diag(r_eco)), err_msg=name + ) + np.testing.assert_array_equal(r_only[1], piv_eco, err_msg=name) + + def test_equilibrated_lstsq_f_order_overwrite_matches_reference(self): + """Lock for the F-order + overwrite_a change: the in-place gelsd + consume must return bit-identical coefficients to a plain + non-overwriting C-order call (same values reach dgelsd either way), + and must not mutate the caller's X.""" + from scipy.linalg import lstsq as scipy_lstsq + + from diff_diff.linalg import _equilibrated_lstsq + + rng = np.random.default_rng(11) + X = rng.standard_normal((300, 6)) + X[:, 3] *= 1e7 # scale disparity exercises the equilibration + y = X[:, :3].sum(axis=1) + rng.standard_normal(300) + x_before = X.copy() + + coef = _equilibrated_lstsq(X, y) + + np.testing.assert_array_equal(X, x_before) # caller's X untouched + # like-for-like reference: SAME norm accumulation (einsum) and lstsq + # options as _equilibrated_lstsq, differing ONLY in the C-order + # non-overwriting call - the exact contract the F-order change claims + # to preserve. + norms = np.sqrt(np.einsum("ij,ij->j", X, X)) + ref = ( + scipy_lstsq(X / norms, y, lapack_driver="gelsd", check_finite=False, cond=1e-07)[0] + / norms + ) + np.testing.assert_array_equal(coef, ref) + class TestEstimatorIntegration: """Integration tests verifying estimators produce correct results.""" diff --git a/tests/test_rust_backend.py b/tests/test_rust_backend.py index 23759ad4..66337ed8 100644 --- a/tests/test_rust_backend.py +++ b/tests/test_rust_backend.py @@ -539,6 +539,45 @@ def test_solve_ols_coefficients_match(self): rust_resid, numpy_resid, decimal=8, err_msg="OLS residuals should match" ) + def test_solve_ols_underdetermined_match(self): + """n < k through the RAW rust backend's slimmed marshalling: the + thin-SVD U/V shapes flip (U is n x n, V is k x n), exercising the + uty and drop paths on the underdetermined branch. This is a + direct-kernel test - the PUBLIC solve_ols rejects n < k outright - + asserting the engines' shared residual/exact-fit contract.""" + from diff_diff._rust_backend import solve_ols as rust_fn + from diff_diff.linalg import _solve_ols_numpy as numpy_fn + + np.random.seed(7) + n, k = 6, 9 + X = np.random.randn(n, k) + y = np.random.randn(n) + + import warnings as _w + + rust_coeffs, rust_resid, _ = rust_fn(X, y, None, True) + with _w.catch_warnings(): + _w.simplefilter("ignore") # numpy path warns on the rank drop + _, numpy_resid, _ = numpy_fn(X, y, cluster_ids=None) + + # Coefficient CONVENTIONS legitimately differ here (rust kernel: + # truncated-SVD minimum-norm over all k; numpy path: column-drop with + # NaN); the public solve_ols never routes n < k to either engine (it + # rejects such designs up front), so like + # test_rank_deficient_ols_residuals_match this asserts the engines' + # shared contract: an exact fit with matching residuals. + assert np.all(np.isfinite(rust_coeffs)) + np.testing.assert_array_almost_equal( + rust_resid, np.zeros(n), decimal=10, err_msg="rust residuals ~0" + ) + np.testing.assert_array_almost_equal( + rust_resid, numpy_resid, decimal=8, err_msg="residuals should match" + ) + # the rust min-norm solution reproduces y exactly (fitted = X @ beta) + np.testing.assert_array_almost_equal( + X @ rust_coeffs, y, decimal=10, err_msg="exact fit expected" + ) + def test_solve_ols_with_clusters_match(self): """Test Rust and NumPy OLS with cluster SEs match.""" from diff_diff._rust_backend import solve_ols as rust_fn