feat(rust): one-way HC2 robust vcov kernel

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

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_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_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.

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!(

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<Bound<'py, PyArray2<f64>>> {
+    let x_arr = x.as_array();
+    let residuals_arr = residuals.as_array();
+
+    if residuals_arr.len() != x_arr.nrows() {
+        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(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<f64> = (&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::<pyo3::exceptions::PyValueError, _>(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<f64> = 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<f64>,

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)