From fef419713ff9ef554fe00f455b6cb05dd6b9e21c Mon Sep 17 00:00:00 2001 From: igerber Date: Tue, 7 Jul 2026 19:12:44 -0400 Subject: [PATCH] feat(rust): one-way HC2 robust vcov kernel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Rust vcov path supported only HC1/CR1; HC2 fell through to NumPy (noticeable at large n). Adds compute_robust_vcov_hc2 mirroring the NumPy hc2 branch exactly: hat diagonals off the same bread, u^2/max(1-h, 1e-10) leverage meat, no n/(n-k) factor — matches NumPy at ~1e-15 across a seed grid. The near-singular hat-diagonal guard stays Python-side: the kernel returns a sentinel error and the documented warn-and-fall-back-to-HC1 fires in the dispatcher, identical to the NumPy branch (locked by a monkeypatched-sentinel test + an exact-h=1 clamp-parity test). The symbol imports independently (mixed-version safe: a stale extension degrades HC2 to NumPy without disabling older Rust accelerations). return_dof / weighted / clustered / CR2-BM requests stay on NumPy; the TODO row is narrowed to CR2-BM. Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 10 ++++ TODO.md | 2 +- diff_diff/_backend.py | 13 +++++ diff_diff/linalg.py | 63 ++++++++++++++++++++ rust/src/lib.rs | 1 + rust/src/linalg.rs | 68 ++++++++++++++++++++++ tests/test_rust_backend.py | 115 +++++++++++++++++++++++++++++++++++++ 7 files changed, 271 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2f7c8dab..c84c4b41 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -434,6 +434,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `docs/methodology/REGISTRY.md` (both estimator sections + "Absorbed Fixed Effects"). ### Performance +- **Rust-backend HC2 vcov.** The Rust vcov path supported only HC1/CR1; one-way + (unclustered, unweighted) HC2 now dispatches to a new `compute_robust_vcov_hc2` kernel + mirroring the NumPy branch exactly — hat diagonals off the same bread, + `u²/max(1−h, 1e-10)` leverage meat, no n/(n−k) factor — matching NumPy at ~1e-15 on a + seed grid. The near-singular hat-diagonal guard stays Python-side: the kernel returns a + sentinel error and the documented warn-and-fall-back-to-HC1 fires in the dispatcher, + identical to the NumPy branch (locked by a monkeypatched-sentinel test plus an exact + h=1 clamp-parity test). The symbol is imported independently (mixed-version safe — a + stale extension degrades HC2 to NumPy without disabling older Rust accelerations). + `return_dof` / weighted / CR2-BM requests stay on NumPy (CR2-BM tracked in TODO). - **`CallawaySantAnna` per-(g,t) IF scatters converted from `np.add.at` to fancy `+=`** (`staggered.py::_cluster_robust_se_from_per_gt_if` — runs once per (g,t) cell when `cluster=` is set — and the general combined-IF assembly path in diff --git a/TODO.md b/TODO.md index e55b68c0..f95e5234 100644 --- a/TODO.md +++ b/TODO.md @@ -48,7 +48,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | 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 | | `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 | +| Rust-backend CR2 Bell-McCaffrey: falls through to NumPy (the leverage/Satterthwaite-DOF path needs `return_dof` support, which the Rust vcov dispatch excludes). The one-way HC2 kernel landed 2026-07-07 (`compute_robust_vcov_hc2`, mirrors the NumPy hc2 branch at ~1e-15; near-singular hat-diagonal sentinel + Python-side warn-and-HC1-fallback). | `rust/src/linalg.rs` | Phase 1a | Mid | Low | | 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 | diff --git a/diff_diff/_backend.py b/diff_diff/_backend.py index 2824fb90..9175fcbb 100644 --- a/diff_diff/_backend.py +++ b/diff_diff/_backend.py @@ -82,6 +82,17 @@ except ImportError: _rust_batched_ridge_chol_solve = None +# HC2 (leverage-corrected) robust vcov: imported independently for the same +# mixed-version reason as demean_map (a stale extension missing only this +# newer symbol degrades HC2 to the NumPy path without disabling the older +# Rust accelerations). +try: + from diff_diff._rust_backend import ( + compute_robust_vcov_hc2 as _rust_compute_robust_vcov_hc2, + ) +except ImportError: + _rust_compute_robust_vcov_hc2 = None + # Determine final backend based on environment variable and availability if _backend_env == "python": # Force pure Python mode - disable Rust even if available @@ -94,6 +105,8 @@ _rust_demean_map = None # Batched ridge-regularized SPD solve _rust_batched_ridge_chol_solve = None + # HC2 robust vcov + _rust_compute_robust_vcov_hc2 = None # TROP estimator acceleration (local method) _rust_unit_distance_matrix = None _rust_loocv_grid_search = None diff --git a/diff_diff/linalg.py b/diff_diff/linalg.py index c779caf5..b0b70e09 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -47,6 +47,7 @@ from diff_diff._backend import ( HAS_RUST_BACKEND, _rust_compute_robust_vcov, + _rust_compute_robust_vcov_hc2, _rust_solve_ols, ) @@ -1712,6 +1713,68 @@ def compute_robust_vcov( if weights is not None: weights = _validate_weights(weights, weight_type, X.shape[0]) + # Rust HC2 (one-way, unweighted, no DOF): mirrors the NumPy hc2 branch + # exactly (leverage meat, no n/(n-k) factor). The near-singular + # hat-diagonal guard stays Python-side: the kernel returns a sentinel + # error and the documented warn-and-fall-back-to-HC1 fires here, + # identical to the NumPy branch's behavior. Imported independently + # (mixed-version safe) — None on a stale extension. + if ( + HAS_RUST_BACKEND + and _rust_compute_robust_vcov_hc2 is not None + and weights is None + and vcov_type == "hc2" + and cluster_ids is None + and not return_dof + ): + X_c = np.ascontiguousarray(X, dtype=np.float64) + residuals_c = np.ascontiguousarray(residuals, dtype=np.float64) + try: + return _rust_compute_robust_vcov_hc2(X_c, residuals_c) + except ValueError as e: + error_msg = str(e) + if "Hat-matrix diagonal exceeds 1" in error_msg: + warnings.warn( + f"{error_msg} Falling back to HC1.", + UserWarning, + stacklevel=2, + ) + return _compute_robust_vcov_numpy( + X, + residuals, + cluster_ids=None, + weights=None, + weight_type=weight_type, + vcov_type="hc1", + return_dof=return_dof, + ) + if "Matrix inversion failed" in error_msg: + raise ValueError( + "Design matrix is rank-deficient (singular X'X matrix). " + "This indicates perfect multicollinearity. Check your fixed effects " + "and covariates for linear dependencies." + ) from e + if "numerically unstable" in error_msg.lower(): + # Mirror the HC1 dispatch: fall back to the NumPy HC2 branch + # (which applies its own hat-diagonal guard semantics) rather + # than hard-erroring where the pre-kernel path would not. + warnings.warn( + f"Rust backend detected numerical instability: {e}. " + "Falling back to Python backend for variance computation.", + UserWarning, + stacklevel=2, + ) + return _compute_robust_vcov_numpy( + X, + residuals, + cluster_ids=None, + weights=None, + weight_type=weight_type, + vcov_type="hc2", + return_dof=return_dof, + ) + raise + # Use Rust backend if available AND no weights AND the requested path is # the unchanged HC1/CR1 dispatch AND the caller does not need DOF. Any # other combination falls through to the NumPy implementation below. diff --git a/rust/src/lib.rs b/rust/src/lib.rs index ceebb5b9..10f25dbd 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -53,6 +53,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::compute_robust_vcov, m)?)?; + m.add_function(wrap_pyfunction!(linalg::compute_robust_vcov_hc2, m)?)?; // Batched ridge-regularized SPD solve (EfficientDiD per-unit weights) m.add_function(wrap_pyfunction!( diff --git a/rust/src/linalg.rs b/rust/src/linalg.rs index c9940277..29db7c80 100644 --- a/rust/src/linalg.rs +++ b/rust/src/linalg.rs @@ -236,6 +236,74 @@ pub fn compute_robust_vcov<'py>( Ok(vcov.to_pyarray(py)) } +/// HC2 (leverage-corrected) heteroskedasticity-robust vcov, one-way only. +/// +/// Mirrors the NumPy `_compute_robust_vcov_numpy` unweighted `hc2` branch +/// exactly (sandwich::vcovHC type="HC2" convention): +/// h_i = x_i' (X'X)^{-1} x_i +/// meat = X' diag(u_i^2 / max(1 - h_i, 1e-10)) X +/// vcov = (X'X)^{-1} meat (X'X)^{-1} (NO n/(n-k) factor) +/// A hat diagonal exceeding 1 + 1e-6 signals a near-singular design; this +/// returns the sentinel error "Hat-matrix diagonal exceeds 1" so the Python +/// dispatcher can reproduce the documented warn-and-fall-back-to-HC1 +/// behavior (the guard decision stays in one place, Python-side). +/// +/// # Arguments +/// * `x` - Design matrix (n, k) +/// * `residuals` - OLS residuals (n,) +/// +/// # Returns +/// Variance-covariance matrix (k, k) +#[pyfunction] +#[pyo3(signature = (x, residuals))] +pub fn compute_robust_vcov_hc2<'py>( + py: Python<'py>, + x: PyReadonlyArray2<'py, f64>, + residuals: PyReadonlyArray1<'py, f64>, +) -> PyResult>> { + let x_arr = x.as_array(); + let residuals_arr = residuals.as_array(); + + if residuals_arr.len() != x_arr.nrows() { + return Err(PyErr::new::(format!( + "residuals length ({}) must match design rows ({})", + residuals_arr.len(), + x_arr.nrows() + ))); + } + + let xtx = x_arr.t().dot(&x_arr); + let xtx_inv = invert_symmetric(&xtx)?; + + // Hat diagonals h_i = x_i' (X'X)^{-1} x_i via one GEMM + rowwise dot: + // H_diag = rowsum((X (X'X)^{-1}) * X). + let x_bread = x_arr.dot(&xtx_inv); // (n, k) + let h_diag: Array1 = (&x_bread * &x_arr).sum_axis(Axis(1)); + + let h_max = h_diag.iter().cloned().fold(f64::NEG_INFINITY, f64::max); + if h_max > 1.0 + 1e-6 { + return Err(PyErr::new::(format!( + "Hat-matrix diagonal exceeds 1 (max={:.6}); the design is near-singular.", + h_max + ))); + } + + // meat = X' diag(u^2 / max(1 - h, 1e-10)) X + let factor: Array1 = residuals_arr + .iter() + .zip(h_diag.iter()) + .map(|(u, h)| u * u / (1.0 - h).max(1e-10)) + .collect(); + let factor_col = factor.insert_axis(Axis(1)); // (n, 1) + let x_weighted = &x_arr * &factor_col; // (n, k) + let meat = x_arr.t().dot(&x_weighted); // (k, k) + + // Sandwich WITHOUT DOF adjustment (HC2's leverage correction replaces it). + let temp = xtx_inv.dot(&meat); + let vcov = temp.dot(&xtx_inv); + Ok(vcov.to_pyarray(py)) +} + /// Internal implementation of robust variance-covariance computation. fn compute_robust_vcov_internal( x: &ArrayView2, diff --git a/tests/test_rust_backend.py b/tests/test_rust_backend.py index 210ebb84..1275ae24 100644 --- a/tests/test_rust_backend.py +++ b/tests/test_rust_backend.py @@ -3507,3 +3507,118 @@ def test_shape_validation_errors(self): _batched_chol_symbol(np.zeros((2, 3, 4)), np.zeros(2)) with pytest.raises(ValueError, match="ridge length"): _batched_chol_symbol(np.zeros((2, 3, 3)), np.zeros(5)) + + +@pytest.mark.skipif(not HAS_RUST_BACKEND, reason="Rust backend not available") +class TestRustHC2Vcov: + """Rust HC2 (leverage-corrected) vcov parity with the NumPy hc2 branch. + + The kernel mirrors `_compute_robust_vcov_numpy`'s unweighted `hc2` path + exactly (hat diagonals off the same bread, `u^2 / max(1 - h, 1e-10)` meat, + NO n/(n-k) factor); the near-singular hat-diagonal guard stays Python-side + (sentinel error -> documented warn-and-fall-back-to-HC1).""" + + @staticmethod + def _design(n=400, k=5, seed=0): + rng = np.random.default_rng(seed) + X = np.column_stack([np.ones(n), rng.normal(size=(n, k - 1))]) + beta = rng.normal(size=k) + e = rng.normal(size=n) * (1.0 + 0.5 * np.abs(X[:, 1])) # heteroskedastic + y = X @ beta + e + resid = y - X @ np.linalg.lstsq(X, y, rcond=None)[0] + return X, resid + + def test_hc2_matches_numpy(self): + from diff_diff.linalg import _compute_robust_vcov_numpy, compute_robust_vcov + + for seed in (0, 1, 2): + X, resid = self._design(seed=seed) + v_rust = compute_robust_vcov(X, resid, vcov_type="hc2") + v_py = _compute_robust_vcov_numpy(X, resid, None, vcov_type="hc2") + np.testing.assert_allclose(v_rust, v_py, rtol=1e-12, atol=1e-15) + + def test_hc2_kernel_direct(self): + from diff_diff._rust_backend import compute_robust_vcov_hc2 + + X, resid = self._design() + v = compute_robust_vcov_hc2(np.ascontiguousarray(X), np.ascontiguousarray(resid)) + assert v.shape == (X.shape[1], X.shape[1]) + np.testing.assert_allclose(v, v.T, rtol=0, atol=1e-12) + assert np.all(np.diag(v) > 0) + + def test_exact_unit_leverage_clamp_parity(self): + """h_ii == 1 exactly (a one-obs dummy is its own perfect predictor) + does NOT trip the > 1 + 1e-6 guard in either backend — both take the + max(1 - h, 1e-10) clamp path and must agree.""" + from diff_diff.linalg import _compute_robust_vcov_numpy, compute_robust_vcov + + n = 60 + rng = np.random.default_rng(3) + d = np.zeros(n) + d[0] = 1.0 + X = np.column_stack([np.ones(n), d, rng.normal(size=n)]) + y = X @ np.array([1.0, 2.0, 0.5]) + rng.normal(size=n) + resid = y - X @ np.linalg.lstsq(X, y, rcond=None)[0] + + v_rust_path = compute_robust_vcov(X, resid, vcov_type="hc2") + v_numpy_path = _compute_robust_vcov_numpy(X, resid, None, vcov_type="hc2") + np.testing.assert_allclose(v_rust_path, v_numpy_path, rtol=1e-9, atol=1e-12) + + def test_sentinel_error_falls_back_to_hc1_with_warning(self, monkeypatch): + """The kernel's near-singular sentinel error must reproduce the NumPy + branch's warn-and-fall-back-to-HC1 through the dispatcher (the guard + decision is Python-side; the kernel only signals).""" + import diff_diff.linalg as la + + X, resid = self._design() + + def _sentinel(*a, **k): + raise ValueError( + "Hat-matrix diagonal exceeds 1 (max=1.000010); the design is near-singular." + ) + + monkeypatch.setattr(la, "_rust_compute_robust_vcov_hc2", _sentinel) + with pytest.warns(UserWarning, match="Falling back to HC1"): + v = la.compute_robust_vcov(X, resid, vcov_type="hc2") + v_hc1 = la._compute_robust_vcov_numpy(X, resid, None, vcov_type="hc1") + np.testing.assert_allclose(v, v_hc1, rtol=1e-12, atol=1e-15) + + def test_dispatch_declined_for_dof_and_weights(self): + """return_dof / weights / cluster requests stay on the NumPy path + (values equal by construction; this locks that the kernel's absence + of those features cannot change results).""" + from diff_diff.linalg import compute_robust_vcov + + X, resid = self._design() + v, dof = compute_robust_vcov(X, resid, vcov_type="hc2", return_dof=True) + assert dof.shape == (X.shape[1],) + assert np.all(dof == X.shape[0] - X.shape[1]) + w = np.ones(X.shape[0]) + v_w = compute_robust_vcov(X, resid, weights=w, weight_type="pweight", vcov_type="hc2") + np.testing.assert_allclose(v_w, v, rtol=1e-10, atol=1e-14) + + def test_kernel_rejects_length_mismatch(self): + """Malformed inputs must fail loudly, not silently truncate.""" + from diff_diff._rust_backend import compute_robust_vcov_hc2 + + X, resid = self._design() + too_long = np.concatenate([resid, [1.0, 2.0]]) + with pytest.raises(ValueError, match="must match design rows"): + compute_robust_vcov_hc2(np.ascontiguousarray(X), np.ascontiguousarray(too_long)) + + def test_numerical_instability_falls_back_to_numpy_hc2(self, monkeypatch): + """The HC1 dispatch's numerically-unstable fallback is mirrored: the + dispatcher warns and returns the NumPy HC2 result (not a hard error, + and not HC1).""" + import diff_diff.linalg as la + + X, resid = self._design() + + def _unstable(*a, **k): + raise ValueError("Matrix inversion numerically unstable (residual check failed)") + + monkeypatch.setattr(la, "_rust_compute_robust_vcov_hc2", _unstable) + with pytest.warns(UserWarning, match="numerical instability"): + v = la.compute_robust_vcov(X, resid, vcov_type="hc2") + v_py = la._compute_robust_vcov_numpy(X, resid, None, vcov_type="hc2") + np.testing.assert_allclose(v, v_py, rtol=1e-12, atol=1e-15)