1#include "raid/statistics/covariance.hpp"2#include "raid/core/error.hpp"3#include <cmath>4#include <map>56namespace raid::stats {78std::string to_string(CovType c) {9 switch (c) { case CovType::classical: return "classical"; case CovType::hc0: return "hc0"; case CovType::hc1: return "hc1"; case CovType::hc2: return "hc2"; case CovType::hc3: return "hc3"; case CovType::hac_bartlett: return "hac_bartlett"; case CovType::cluster: return "cluster"; }10 return "?";11}12CovType covtype_from_string(std::string_view s) {13 if (s == "classical" || s == "nonrobust") return CovType::classical;14 if (s == "hc0" || s == "HC0") return CovType::hc0; if (s == "hc1" || s == "HC1") return CovType::hc1;15 if (s == "hc2" || s == "HC2") return CovType::hc2; if (s == "hc3" || s == "HC3") return CovType::hc3;16 if (s == "hac" || s == "hac_bartlett" || s == "newey_west") return CovType::hac_bartlett;17 if (s == "cluster") return CovType::cluster;18 throw RaidError(ErrorCode::invalid_argument, "unknown covariance type '" + std::string(s) + "'");19}20int default_hac_lags(usize n) { return static_cast<int>(std::floor(4.0 * std::pow(static_cast<f64>(n) / 100.0, 2.0 / 9.0))); }2122Eigen::MatrixXd robust_covariance(const Eigen::Ref<const Eigen::MatrixXd>& X, const Eigen::Ref<const Eigen::VectorXd>& e,23 const Eigen::Ref<const Eigen::MatrixXd>& xtx_inv, const CovOptions& opt, const Eigen::VectorXd& h) {24 const Eigen::Index n = X.rows(), k = X.cols();25 RAID_REQUIRE(e.size() == n, invalid_argument, "robust_covariance: residual length");26 const f64 nd = static_cast<f64>(n), kd = static_cast<f64>(k);27 switch (opt.type) {28 case CovType::classical: {29 RAID_REQUIRE(n > k, insufficient_samples, "classical covariance needs n > k");30 const f64 s2 = e.squaredNorm() / (nd - kd);31 return s2 * xtx_inv;32 }33 case CovType::hc0: case CovType::hc1: case CovType::hc2: case CovType::hc3: {34 Eigen::VectorXd w = e.array().square();35 if (opt.type == CovType::hc1) w *= nd / (nd - kd);36 if (opt.type == CovType::hc2 || opt.type == CovType::hc3) {37 RAID_REQUIRE(h.size() == n, invalid_argument, "HC2/HC3 need leverage");38 for (Eigen::Index i = 0; i < n; ++i) { const f64 d = 1.0 - h[i]; w[i] /= (opt.type == CovType::hc2 ? d : d * d); }39 }40 const Eigen::MatrixXd meat = X.transpose() * w.asDiagonal() * X;41 return xtx_inv * meat * xtx_inv;42 }43 case CovType::hac_bartlett: {44 const int L = opt.hac_lags < 0 ? default_hac_lags(static_cast<usize>(n)) : opt.hac_lags;45 Eigen::MatrixXd S = Eigen::MatrixXd::Zero(k, k);46 Eigen::MatrixXd U = X.array().colwise() * e.array(); // scores n×k47 S = U.transpose() * U;48 for (int l = 1; l <= L; ++l) {49 const f64 wl = 1.0 - static_cast<f64>(l) / (static_cast<f64>(L) + 1.0);50 const Eigen::MatrixXd G = U.bottomRows(n - l).transpose() * U.topRows(n - l);51 S += wl * (G + G.transpose());52 }53 if (opt.hac_small_sample) S *= nd / (nd - kd);54 return xtx_inv * S * xtx_inv;55 }56 case CovType::cluster: {57 RAID_REQUIRE(static_cast<Eigen::Index>(opt.clusters.size()) == n, invalid_argument, "cluster ids length");58 std::map<i64, Eigen::VectorXd> sums;59 for (Eigen::Index i = 0; i < n; ++i) {60 auto it = sums.find(opt.clusters[static_cast<usize>(i)]);61 if (it == sums.end()) it = sums.emplace(opt.clusters[static_cast<usize>(i)], Eigen::VectorXd::Zero(k)).first;62 it->second += X.row(i).transpose() * e[i];63 }64 Eigen::MatrixXd S = Eigen::MatrixXd::Zero(k, k);65 for (auto& [g, s] : sums) S += s * s.transpose();66 const f64 G = static_cast<f64>(sums.size());67 f64 corr = 1.0;68 if (opt.cluster_small_sample) { RAID_REQUIRE(G > 1 && n > k, insufficient_samples, "cluster correction needs G > 1, n > k"); corr = (G / (G - 1.0)) * ((nd - 1.0) / (nd - kd)); }69 return corr * (xtx_inv * S * xtx_inv);70 }71 }72 throw RaidError(ErrorCode::invalid_argument, "unknown CovType");73}7475} // namespace raid::stats76