1#include "raid/estimators/ols.hpp"2#include "raid/core/error.hpp"3#include <chrono>4#include <cmath>56namespace raid::est {78namespace {9// Two-sided p-value from a t statistic with df degrees of freedom (df<=0 → normal). Regularised incomplete beta via10// Lentz continued fraction (Numerical Recipes), FP64. Validated against scipy.stats.t.sf in tests/reference.11f64 betacf(f64 a, f64 b, f64 x) {12 const int MAXIT = 300; const f64 EPS = 3e-16, FPMIN = 1e-300;13 f64 qab = a + b, qap = a + 1, qam = a - 1, c = 1, d = 1 - qab * x / qap;14 if (std::abs(d) < FPMIN) d = FPMIN; d = 1 / d; f64 h = d;15 for (int m = 1; m <= MAXIT; ++m) {16 const f64 m2 = 2.0 * m; f64 aa = m * (b - m) * x / ((qam + m2) * (a + m2));17 d = 1 + aa * d; if (std::abs(d) < FPMIN) d = FPMIN; c = 1 + aa / c; if (std::abs(c) < FPMIN) c = FPMIN; d = 1 / d; h *= d * c;18 aa = -(a + m) * (qab + m) * x / ((a + m2) * (qap + m2));19 d = 1 + aa * d; if (std::abs(d) < FPMIN) d = FPMIN; c = 1 + aa / c; if (std::abs(c) < FPMIN) c = FPMIN; d = 1 / d;20 const f64 del = d * c; h *= del; if (std::abs(del - 1) < EPS) break;21 }22 return h;23}24f64 betai(f64 a, f64 b, f64 x) {25 if (x <= 0) return 0; if (x >= 1) return 1;26 const f64 bt = std::exp(std::lgamma(a + b) - std::lgamma(a) - std::lgamma(b) + a * std::log(x) + b * std::log(1 - x));27 if (x < (a + 1) / (a + b + 2)) return bt * betacf(a, b, x) / a;28 return 1 - bt * betacf(b, a, 1 - x) / b;29}30} // namespace3132f64 t_pvalue_two_sided(f64 t, f64 df) {33 if (!std::isfinite(t)) return std::numeric_limits<f64>::quiet_NaN();34 if (df <= 0) return std::erfc(std::abs(t) / std::sqrt(2.0));35 return betai(df / 2.0, 0.5, df / (df + t * t));36}3738OlsResult ols(const Eigen::Ref<const Eigen::MatrixXd>& X0, const Eigen::Ref<const Eigen::VectorXd>& y0, const OlsOptions& opt) {39 const auto t0 = std::chrono::steady_clock::now();40 const Eigen::Index n = X0.rows(), k = X0.cols();41 RAID_REQUIRE(y0.size() == n, invalid_argument, "ols: X rows != y size");42 RAID_REQUIRE(n > 0 && k > 0, insufficient_samples, "ols: empty design");43 RAID_REQUIRE(X0.allFinite() && y0.allFinite(), data_quality, "ols: non-finite values in X or y (drop or impute explicitly)");44 OlsResult r; r.estimate.model = opt.ridge_lambda > 0 ? "ridge" : (opt.weights.empty() ? "ols" : "wls");45 // Weighted transform: X̃ = √W X, ỹ = √W y46 Eigen::MatrixXd X = X0; Eigen::VectorXd y = y0;47 if (!opt.weights.empty()) {48 RAID_REQUIRE(static_cast<Eigen::Index>(opt.weights.size()) == n, invalid_argument, "wls: weights length");49 for (Eigen::Index i = 0; i < n; ++i) { RAID_REQUIRE(opt.weights[static_cast<usize>(i)] >= 0, invalid_argument, "wls: negative weight"); const f64 s = std::sqrt(opt.weights[static_cast<usize>(i)]); X.row(i) *= s; y[i] *= s; }50 }51 // Ridge via augmentation: append √λ·P rows (P = identity on penalised columns)52 Eigen::MatrixXd Xa = X; Eigen::VectorXd ya = y;53 if (opt.ridge_lambda > 0) {54 Eigen::VectorXd pen = Eigen::VectorXd::Constant(k, std::sqrt(opt.ridge_lambda));55 for (usize u : opt.unpenalized) { RAID_REQUIRE(u < static_cast<usize>(k), invalid_argument, "ridge: bad unpenalized index"); pen[static_cast<Eigen::Index>(u)] = 0; }56 Xa.conservativeResize(n + k, k); ya.conservativeResize(n + k);57 Xa.bottomRows(k) = pen.asDiagonal(); ya.tail(k).setZero();58 }59 Eigen::VectorXd beta(k);60 if (opt.solver == Solver::svd) {61 Eigen::JacobiSVD<Eigen::MatrixXd> svd(Xa, Eigen::ComputeThinU | Eigen::ComputeThinV);62 const auto& s = svd.singularValues();63 r.rank = svd.rank(); r.rcond = s[s.size() - 1] / s[0];64 beta = svd.solve(ya);65 Eigen::VectorXd inv_s2 = (s.array() > s[0] * 1e-14).select(1.0 / s.array().square(), 0.0);66 r.xtx_inv = svd.matrixV() * inv_s2.asDiagonal() * svd.matrixV().transpose();67 } else if (opt.solver == Solver::normal_cholesky) {68 const Eigen::MatrixXd A = Xa.transpose() * Xa; Eigen::LLT<Eigen::MatrixXd> llt(A);69 RAID_REQUIRE(llt.info() == Eigen::Success, numerical, "ols: X'X not positive definite");70 beta = llt.solve(Xa.transpose() * ya); r.xtx_inv = llt.solve(Eigen::MatrixXd::Identity(k, k)); r.rank = k;71 Eigen::JacobiSVD<Eigen::MatrixXd> svd(A); r.rcond = svd.singularValues()[k - 1] / svd.singularValues()[0];72 } else {73 Eigen::ColPivHouseholderQR<Eigen::MatrixXd> qr(Xa);74 qr.setThreshold(1e-12);75 r.rank = qr.rank();76 const Eigen::MatrixXd R = qr.matrixR().topLeftCorner(k, k).template triangularView<Eigen::Upper>();77 r.rcond = std::abs(R(k - 1, k - 1)) / std::abs(R(0, 0));78 if (r.rank < k) {79 r.estimate.warnings.push_back("design is rank deficient (rank " + std::to_string(r.rank) + " < " + std::to_string(k) + "); minimum-norm solution via SVD");80 Eigen::JacobiSVD<Eigen::MatrixXd> svd(Xa, Eigen::ComputeThinU | Eigen::ComputeThinV);81 beta = svd.solve(ya);82 const auto& s = svd.singularValues(); Eigen::VectorXd inv_s2 = (s.array() > s[0] * 1e-14).select(1.0 / s.array().square(), 0.0);83 r.xtx_inv = svd.matrixV() * inv_s2.asDiagonal() * svd.matrixV().transpose();84 } else {85 beta = qr.solve(ya);86 // (X'X)^{-1} = P R^{-1} R^{-T} P'87 const Eigen::MatrixXd Rinv = R.template triangularView<Eigen::Upper>().solve(Eigen::MatrixXd::Identity(k, k));88 r.xtx_inv = qr.colsPermutation() * (Rinv * Rinv.transpose()) * qr.colsPermutation().transpose();89 }90 }91 if (r.rcond < opt.rcond_warn) r.estimate.warnings.push_back("ill-conditioned design: rcond = " + std::to_string(r.rcond));92 r.fitted = X0 * beta; // unweighted fitted values93 r.residuals = y0 - r.fitted;94 Eigen::VectorXd e = y - X * beta; // weighted residuals for inference95 const f64 rss = e.squaredNorm();96 const f64 df_resid = static_cast<f64>(n - r.rank);97 r.sigma2 = df_resid > 0 ? rss / df_resid : std::numeric_limits<f64>::quiet_NaN();98 const f64 ybar = y.mean(); const f64 tss = (y.array() - ybar).square().sum();99 r.r2 = tss > 0 ? 1.0 - rss / tss : std::numeric_limits<f64>::quiet_NaN();100 r.r2_adj = tss > 0 && df_resid > 0 ? 1.0 - (rss / df_resid) / (tss / static_cast<f64>(n - 1)) : std::numeric_limits<f64>::quiet_NaN();101 // leverage h_ii = x_i' (X'X)^{-1} x_i102 r.leverage = (X * r.xtx_inv).cwiseProduct(X).rowwise().sum();103 // inference104 InferenceResult inf;105 stats::CovOptions co = opt.cov;106 inf.covariance = stats::robust_covariance(X, e, r.xtx_inv, co, r.leverage);107 inf.covariance_type = stats::to_string(co.type);108 if (co.type == stats::CovType::hac_bartlett) inf.covariance_type += "(" + std::to_string(co.hac_lags < 0 ? stats::default_hac_lags(static_cast<usize>(n)) : co.hac_lags) + ")";109 inf.std_errors.resize(static_cast<usize>(k)); inf.t_stats.resize(static_cast<usize>(k)); inf.p_values.resize(static_cast<usize>(k));110 const f64 df_t = (co.type == stats::CovType::classical) ? df_resid : df_resid; // statsmodels uses t(df_resid) for OLS by default111 for (Eigen::Index j = 0; j < k; ++j) {112 const f64 se = std::sqrt(std::max(0.0, inf.covariance(j, j)));113 inf.std_errors[static_cast<usize>(j)] = se; inf.t_stats[static_cast<usize>(j)] = se > 0 ? beta[j] / se : std::numeric_limits<f64>::quiet_NaN();114 inf.p_values[static_cast<usize>(j)] = t_pvalue_two_sided(inf.t_stats[static_cast<usize>(j)], df_t);115 }116 r.estimate.inference = inf;117 r.estimate.coefficients.assign(beta.data(), beta.data() + k);118 r.estimate.names = opt.names; if (r.estimate.names.size() != static_cast<usize>(k)) { r.estimate.names.clear(); for (Eigen::Index j = 0; j < k; ++j) r.estimate.names.push_back("x" + std::to_string(j)); }119 r.estimate.loss = rss; r.estimate.n_obs = static_cast<usize>(n);120 r.estimate.convergence = Convergence{true, 1, 1, 0.0, 0, "direct solve"};121 r.estimate.extra = {{"sigma2", r.sigma2}, {"r2", r.r2}, {"r2_adj", r.r2_adj}, {"rank", r.rank}, {"rcond", r.rcond}, {"df_resid", df_resid}, {"ridge_lambda", opt.ridge_lambda}};122 r.estimate.runtime_seconds = std::chrono::duration<f64>(std::chrono::steady_clock::now() - t0).count();123 return r;124}125126SufficientStats accumulate(const Eigen::Ref<const Eigen::MatrixXd>& X, const Eigen::Ref<const Eigen::VectorXd>& y) {127 SufficientStats s; s.xtx = X.transpose() * X; s.xty = X.transpose() * y; s.yty = y.squaredNorm(); s.n = static_cast<usize>(X.rows());128 return s;129}130131OlsResult ols_from_stats(const SufficientStats& s, const OlsOptions& opt) {132 const auto t0 = std::chrono::steady_clock::now();133 const Eigen::Index k = s.xtx.rows();134 OlsResult r; r.estimate.model = "ols.sufficient_stats";135 Eigen::MatrixXd A = s.xtx;136 if (opt.ridge_lambda > 0) { for (Eigen::Index j = 0; j < k; ++j) { bool pen = true; for (usize u : opt.unpenalized) if (static_cast<Eigen::Index>(u) == j) pen = false; if (pen) A(j, j) += opt.ridge_lambda; } }137 Eigen::LDLT<Eigen::MatrixXd> ldlt(A);138 Eigen::VectorXd beta;139 Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> es(A);140 const auto& ev = es.eigenvalues();141 r.rcond = ev[0] / ev[k - 1]; r.rank = (ev.array() > ev[k - 1] * 1e-13).count();142 if (ldlt.info() == Eigen::Success && ldlt.isPositive() && r.rank == k) { beta = ldlt.solve(s.xty); r.xtx_inv = ldlt.solve(Eigen::MatrixXd::Identity(k, k)); }143 else {144 r.estimate.warnings.push_back("X'X singular or indefinite: pseudo-inverse solution");145 Eigen::VectorXd inv = (ev.array() > ev[k - 1] * 1e-13).select(1.0 / ev.array(), 0.0);146 r.xtx_inv = es.eigenvectors() * inv.asDiagonal() * es.eigenvectors().transpose(); beta = r.xtx_inv * s.xty;147 }148 if (r.rcond < opt.rcond_warn) r.estimate.warnings.push_back("ill-conditioned X'X: rcond = " + std::to_string(r.rcond));149 const f64 rss = std::max(0.0, s.yty - 2.0 * beta.dot(s.xty) + beta.dot(s.xtx * beta));150 const f64 df_resid = static_cast<f64>(s.n) - static_cast<f64>(r.rank);151 r.sigma2 = df_resid > 0 ? rss / df_resid : std::numeric_limits<f64>::quiet_NaN();152 InferenceResult inf; inf.covariance = r.sigma2 * r.xtx_inv; inf.covariance_type = "classical";153 for (Eigen::Index j = 0; j < k; ++j) { const f64 se = std::sqrt(std::max(0.0, inf.covariance(j, j))); inf.std_errors.push_back(se); inf.t_stats.push_back(se > 0 ? beta[j] / se : std::numeric_limits<f64>::quiet_NaN()); inf.p_values.push_back(t_pvalue_two_sided(inf.t_stats.back(), df_resid)); }154 r.estimate.inference = inf; r.estimate.coefficients.assign(beta.data(), beta.data() + k);155 r.estimate.names = opt.names; if (r.estimate.names.size() != static_cast<usize>(k)) { r.estimate.names.clear(); for (Eigen::Index j = 0; j < k; ++j) r.estimate.names.push_back("x" + std::to_string(j)); }156 r.estimate.loss = rss; r.estimate.n_obs = s.n; r.estimate.convergence = Convergence{true, 1, 1, 0.0, 0, "sufficient statistics"};157 r.estimate.extra = {{"sigma2", r.sigma2}, {"rank", r.rank}, {"rcond", r.rcond}, {"df_resid", df_resid}};158 r.estimate.runtime_seconds = std::chrono::duration<f64>(std::chrono::steady_clock::now() - t0).count();159 return r;160}161162RollingOlsResult rolling_ols(const Eigen::Ref<const Eigen::MatrixXd>& X, const Eigen::Ref<const Eigen::VectorXd>& y, usize window, usize step) {163 const usize n = static_cast<usize>(X.rows()), k = static_cast<usize>(X.cols());164 RAID_REQUIRE(window >= k + 1 && window <= n, insufficient_samples, "rolling_ols: window must satisfy k+1 <= window <= n");165 RAID_REQUIRE(step >= 1, invalid_argument, "rolling_ols: step >= 1");166 const usize nw = (n - window) / step + 1;167 RollingOlsResult r; r.window = window; r.coefficients.resize(static_cast<Eigen::Index>(k), static_cast<Eigen::Index>(nw)); r.rss.resize(nw); r.ranks.resize(nw);168 for (usize j = 0; j < nw; ++j) {169 const usize b = j * step;170 Eigen::ColPivHouseholderQR<Eigen::MatrixXd> qr(X.middleRows(static_cast<Eigen::Index>(b), static_cast<Eigen::Index>(window)));171 qr.setThreshold(1e-12);172 const Eigen::VectorXd beta = qr.solve(y.segment(static_cast<Eigen::Index>(b), static_cast<Eigen::Index>(window)));173 r.coefficients.col(static_cast<Eigen::Index>(j)) = beta;174 r.rss[j] = (y.segment(static_cast<Eigen::Index>(b), static_cast<Eigen::Index>(window)) - X.middleRows(static_cast<Eigen::Index>(b), static_cast<Eigen::Index>(window)) * beta).squaredNorm();175 r.ranks[j] = qr.rank();176 }177 return r;178}179180} // namespace raid::est181