regression-diagnostics 0.2.0

use super::{lm_heteroskedasticity_test, non_intercept_columns};
use crate::OlsFit;

/// Result of White's heteroskedasticity test (LM form).
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct WhiteTest {
    /// Lagrange-multiplier statistic `n·R²_aux`.
    pub statistic: f64,
    /// Degrees of freedom (number of auxiliary regressors).
    pub df: usize,
    /// Upper-tail p-value under χ²(df).
    pub p_value: f64,
}

/// White's test for heteroskedasticity.
///
/// Like [`breusch_pagan`](super::breusch_pagan) it regresses the squared
/// residuals on auxiliary regressors, but the auxiliary set is richer: the
/// original regressors, their **squares**, and all **pairwise cross-products**.
/// This makes White's test sensitive to more general (including nonlinear) forms
/// of heteroskedasticity and, because it includes cross-products, it doubles as a
/// specification check.
///
/// # Tradeoff versus Breusch-Pagan
///
/// The extra terms cost degrees of freedom. In **small samples** that spreads the
/// test's power thin, so White's can fail to flag heteroskedasticity that the
/// tighter Breusch-Pagan catches. White's is the more general test; Breusch-Pagan
/// has more power when you already suspect a linear variance relationship. The
/// crate offers both rather than picking one for you.
///
/// With `k` non-intercept regressors the auxiliary regression has `k` linear
/// terms, `k` squares, and `k(k−1)/2` cross-products; duplicate columns that can
/// arise (e.g. the square of a 0/1 dummy equals the dummy) are dropped to keep
/// the auxiliary design full rank.
///
/// Because that regressor count grows quadratically in `k`, small samples can
/// leave the auxiliary regression without enough residual degrees of freedom to
/// estimate. When there are too few observations the statistic and p-value are
/// returned as `NaN` rather than a fabricated number.
pub fn white_test(fit: &OlsFit) -> WhiteTest {
    let n = fit.n_observations();
    let cols = non_intercept_columns(fit);
    let x = fit.design_matrix();

    let base: Vec<Vec<f64>> = cols.iter().map(|&j| x.column(j).to_vec()).collect();
    let k = base.len();

    let mut regressors: Vec<Vec<f64>> = Vec::new();
    // Linear terms.
    for col in &base {
        regressors.push(col.clone());
    }
    // Squares.
    for col in &base {
        regressors.push(col.iter().map(|v| v * v).collect());
    }
    // Pairwise cross-products.
    for a in 0..k {
        for b in (a + 1)..k {
            let cross: Vec<f64> = base[a]
                .iter()
                .zip(base[b].iter())
                .map(|(u, v)| u * v)
                .collect();
            regressors.push(cross);
        }
    }

    // Drop near-duplicate / near-constant auxiliary columns to protect the
    // auxiliary design's rank (e.g. the square of a 0/1 dummy).
    let regressors = dedup_regressors(n, regressors);

    let resid_sq: Vec<f64> = fit.residuals().iter().map(|e| e * e).collect();

    let (statistic, df, p_value) = lm_heteroskedasticity_test(n, &regressors, &resid_sq);
    WhiteTest {
        statistic,
        df,
        p_value,
    }
}

/// Remove constant columns and columns that duplicate an earlier one (up to a
/// tight relative tolerance), which would otherwise make the auxiliary design
/// rank-deficient.
fn dedup_regressors(n: usize, cols: Vec<Vec<f64>>) -> Vec<Vec<f64>> {
    let mut kept: Vec<Vec<f64>> = Vec::new();
    for col in cols {
        let first = col[0];
        let is_constant = col
            .iter()
            .all(|&v| (v - first).abs() <= 1e-12 * first.abs().max(1.0));
        if is_constant {
            continue;
        }
        let dup = kept
            .iter()
            .any(|k| (0..n).all(|i| (k[i] - col[i]).abs() <= 1e-12 * col[i].abs().max(1.0)));
        if !dup {
            kept.push(col);
        }
    }
    kept
}