regression-diagnostics 0.2.0

//! Internal linear-algebra layer.
//!
//! Everything here works in `nalgebra` types and is `pub(crate)` — it never
//! appears in the public API, which speaks `ndarray` (see the crate-root docs
//! and the README "Linear algebra dependency" section for why the boundary is
//! drawn this way).
//!
//! Two properties this module guarantees for the diagnostics that depend on it:
//!
//! * **QR, not the normal equations.** Coefficients come from a QR solve of `X`
//!   directly, never from inverting `XᵀX`. Forming `XᵀX` squares the condition
//!   number, which is precisely the quantity the multicollinearity diagnostics
//!   exist to measure — solving `XᵀX β = Xᵀy` would make VIF and the condition
//!   number untrustworthy exactly when they matter most.
//! * **No `n × n` hat matrix.** Leverage (`diag(H)`) is read off the thin `Q`
//!   factor as the squared norm of each row, which is `O(n·p)` memory instead of
//!   the `O(n²)` a materialized `H = X(XᵀX)⁻¹Xᵀ` would cost. Do **not** "simplify"
//!   this back to forming `H`.

use nalgebra::{DMatrix, DVector};

use crate::error::{RegressionError, Result};

/// The full set of quantities the main [`OlsFit`](crate::OlsFit) needs from a
/// single QR factorization, computed once and reused everywhere downstream.
pub(crate) struct QrFit {
    /// Estimated coefficients, one per design-matrix column.
    pub coef: DVector<f64>,
    /// Fitted values `X β`.
    pub fitted: DVector<f64>,
    /// Leverage vector `diag(H)`, read from the thin `Q` factor.
    pub leverage: DVector<f64>,
    /// `(XᵀX)⁻¹`, formed stably from the triangular `R` as `R⁻¹ R⁻ᵀ` (never by
    /// inverting `XᵀX` directly).
    pub xtx_inv: DMatrix<f64>,
    /// Singular values of `X` in descending order (for the condition number).
    pub singular_values: Vec<f64>,
}

/// Fit `y ~ X` by ordinary least squares via QR decomposition.
///
/// Returns [`RegressionError::RankDeficient`] if `X` is not full column rank.
pub(crate) fn ols_via_qr(x: &DMatrix<f64>, y: &DVector<f64>) -> Result<QrFit> {
    let n = x.nrows();
    let p = x.ncols();

    let qr = x.clone().qr();
    // Thin factors: Q is n × p, R is p × p upper-triangular.
    let q = qr.q();
    let r = qr.r();

    // A (near-)zero diagonal entry of R means a linearly dependent column —
    // treat as rank-deficient for the main fit.
    let r_diag_min = (0..p)
        .map(|i| r[(i, i)].abs())
        .fold(f64::INFINITY, f64::min);
    if r_diag_min <= f64::EPSILON * 100.0 * r[(0, 0)].abs().max(1.0) {
        return Err(RegressionError::RankDeficient);
    }
    let r_inv = r.try_inverse().ok_or(RegressionError::RankDeficient)?;

    // Least-squares solve for an overdetermined system: β = R⁻¹ (Qᵀ y).
    // (nalgebra's `QR::solve` only handles square systems, so solve explicitly.)
    let qty = q.transpose() * y;
    let coef = &r_inv * qty;

    // (XᵀX)⁻¹ = (RᵀR)⁻¹ = R⁻¹ R⁻ᵀ.
    let xtx_inv = &r_inv * r_inv.transpose();

    // leverage_i = ‖Q_i,:‖² = the i-th diagonal of H = Q Qᵀ.
    let mut leverage = DVector::<f64>::zeros(n);
    for i in 0..n {
        let mut s = 0.0;
        for j in 0..q.ncols() {
            s += q[(i, j)] * q[(i, j)];
        }
        leverage[i] = s;
    }

    let fitted = x * &coef;

    let singular_values = x
        .clone()
        .singular_values()
        .iter()
        .copied()
        .collect::<Vec<_>>();

    Ok(QrFit {
        coef,
        fitted,
        leverage,
        xtx_inv,
        singular_values,
    })
}

/// Centered coefficient of determination (`R²`) of a least-squares fit of
/// `y ~ X`, where `X` is assumed to contain an intercept column.
///
/// Used by the auxiliary regressions inside VIF, Breusch-Pagan and White's test.
/// Returns `None` only if the target has zero variance (a constant `y`), where
/// `R²` is undefined; collinear auxiliary designs are handled by the caller
/// (they yield `R² → 1`, i.e. an infinite VIF) rather than erroring here.
pub(crate) fn aux_r_squared(x: &DMatrix<f64>, y: &DVector<f64>) -> Option<f64> {
    let n = y.len() as f64;
    let mean = y.sum() / n;
    let tss: f64 = y.iter().map(|v| (v - mean).powi(2)).sum();
    if tss <= 0.0 {
        return None;
    }
    // A rank-deficient auxiliary design means perfect collinearity among the
    // regressors: the target predictor is an exact linear combination of the
    // others, so R² is 1 (infinite VIF). Signal that with R² = 1.
    let qr = x.clone().qr();
    let r = qr.r();
    // `r` is min(rows, cols) × cols; only the leading square block has a
    // diagonal. If the design has more columns than rows it is underdetermined
    // (more regressors than observations) — the auxiliary fit is not estimable.
    if x.ncols() > x.nrows() {
        return None;
    }
    let p = x.ncols();
    let r_diag_min = (0..p)
        .map(|i| r[(i, i)].abs())
        .fold(f64::INFINITY, f64::min);
    if r_diag_min <= f64::EPSILON * 100.0 * r[(0, 0)].abs().max(1.0) {
        return Some(1.0);
    }
    let r_inv = match r.try_inverse() {
        Some(inv) => inv,
        None => return Some(1.0),
    };
    let coef = &r_inv * (qr.q().transpose() * y);
    let fitted = x * &coef;
    let rss: f64 = y
        .iter()
        .zip(fitted.iter())
        .map(|(a, b)| (a - b).powi(2))
        .sum();
    Some((1.0 - rss / tss).clamp(0.0, 1.0))
}

/// Build a column-major [`DMatrix`] from a row-major slice (`n × p`, row order).
pub(crate) fn dmatrix_from_rows(n: usize, p: usize, data: &[f64]) -> DMatrix<f64> {
    DMatrix::from_row_iterator(n, p, data.iter().copied())
}

/// Build a [`DVector`] from a slice.
pub(crate) fn dvector_from_slice(data: &[f64]) -> DVector<f64> {
    DVector::from_iterator(data.len(), data.iter().copied())
}