regression-diagnostics 0.2.0

//! Longley dataset — the standard textbook multicollinearity / OLS benchmark.
//!
//! Reference values are `statsmodels`' published `OLS.summary()` output for the
//! same data (16 observations, response `TOTEMP`, six predictors plus a
//! constant). The design is severely ill-conditioned, which is exactly why it is
//! the canonical stress test for both the coefficient solve (QR vs normal
//! equations) and the multicollinearity diagnostics.

use ndarray::{Array1, Array2};
use regression_diagnostics::fit_statistics::{
    adjusted_r_squared, aic, bic, f_statistic, log_likelihood, r_squared,
};
use regression_diagnostics::multicollinearity::{condition_number, vif};
use regression_diagnostics::residuals::durbin_watson;
use regression_diagnostics::OlsFit;

/// Columns: GNPDEFL, GNP, UNEMP, ARMED, POP, YEAR.
///
/// These are the exact `statsmodels.datasets.longley` values (the raw scale its
/// published coefficients correspond to), not R's rescaled `longley` frame.
const PREDICTORS: [[f64; 6]; 16] = [
    [83.0, 234289.0, 2356.0, 1590.0, 107608.0, 1947.0],
    [88.5, 259426.0, 2325.0, 1456.0, 108632.0, 1948.0],
    [88.2, 258054.0, 3682.0, 1616.0, 109773.0, 1949.0],
    [89.5, 284599.0, 3351.0, 1650.0, 110929.0, 1950.0],
    [96.2, 328975.0, 2099.0, 3099.0, 112075.0, 1951.0],
    [98.1, 346999.0, 1932.0, 3594.0, 113270.0, 1952.0],
    [99.0, 365385.0, 1870.0, 3547.0, 115094.0, 1953.0],
    [100.0, 363112.0, 3578.0, 3350.0, 116219.0, 1954.0],
    [101.2, 397469.0, 2904.0, 3048.0, 117388.0, 1955.0],
    [104.6, 419180.0, 2822.0, 2857.0, 118734.0, 1956.0],
    [108.4, 442769.0, 2936.0, 2798.0, 120445.0, 1957.0],
    [110.8, 444546.0, 4681.0, 2637.0, 121950.0, 1958.0],
    [112.6, 482704.0, 3813.0, 2552.0, 123366.0, 1959.0],
    [114.2, 502601.0, 3931.0, 2514.0, 125368.0, 1960.0],
    [115.7, 518173.0, 4806.0, 2572.0, 127852.0, 1961.0],
    [116.9, 554894.0, 4007.0, 2827.0, 130081.0, 1962.0],
];

const TOTEMP: [f64; 16] = [
    60323.0, 61122.0, 60171.0, 61187.0, 63221.0, 63639.0, 64989.0, 63761.0, 66019.0, 67857.0,
    68169.0, 66513.0, 68655.0, 69564.0, 69331.0, 70551.0,
];

fn longley_fit() -> OlsFit {
    // Design with an explicit leading intercept column of ones.
    let mut x = Array2::<f64>::ones((16, 7));
    for i in 0..16 {
        for j in 0..6 {
            x[(i, j + 1)] = PREDICTORS[i][j];
        }
    }
    let y = Array1::from(TOTEMP.to_vec());
    OlsFit::new(x, y).unwrap()
}

fn rel_close(got: f64, want: f64, rel: f64) -> bool {
    (got - want).abs() <= rel * want.abs().max(1.0)
}

#[test]
fn coefficients_match_statsmodels() {
    let fit = longley_fit();
    let coef = fit.coefficients();
    // statsmodels OLS coefficients (const, GNPDEFL, GNP, UNEMP, ARMED, POP, YEAR).
    let want = [
        -3_482_258.634_596,
        15.061_872,
        -0.035_819,
        -2.020_230,
        -1.033_227,
        -0.051_104,
        1_829.151_465,
    ];
    for (j, &w) in want.iter().enumerate() {
        assert!(
            rel_close(coef[j], w, 1e-4),
            "coef[{j}] = {} but expected ~{w}",
            coef[j]
        );
    }
}

#[test]
fn fit_statistics_match_statsmodels() {
    let fit = longley_fit();
    assert!(rel_close(r_squared(&fit), 0.995479, 1e-5));
    assert!(rel_close(adjusted_r_squared(&fit), 0.992465, 1e-5));

    let f = f_statistic(&fit);
    assert!(rel_close(f.statistic, 330.285, 1e-3), "F = {}", f.statistic);
    assert!(f.p_value < 1e-8);

    assert!(rel_close(log_likelihood(&fit), -109.6174, 1e-4));
    assert!(rel_close(aic(&fit), 233.2349, 1e-4));
    assert!(rel_close(bic(&fit), 238.643, 1e-4));
}

#[test]
fn durbin_watson_matches_statsmodels() {
    let fit = longley_fit();
    assert!(rel_close(durbin_watson(&fit), 2.559, 1e-3));
}

#[test]
fn vif_flags_severe_multicollinearity() {
    let fit = longley_fit();
    let v = vif(&fit);
    // Intercept slot is NaN by convention.
    assert!(v[0].is_nan());
    // statsmodels-published per-predictor VIFs (regressing each on the others
    // including the constant). Longley's are famously enormous.
    let want = [135.53, 1788.51, 33.62, 3.59, 399.15, 758.98];
    for (k, &w) in want.iter().enumerate() {
        let got = v[k + 1];
        assert!(
            rel_close(got, w, 2e-2),
            "VIF for predictor {k} = {got} but expected ~{w}"
        );
    }
}

#[test]
fn condition_number_is_large() {
    let fit = longley_fit();
    // The precise value depends on the (documented) unscaled-design convention;
    // what matters diagnostically is that it is enormous, flagging the design as
    // severely ill-conditioned.
    let c = condition_number(&fit);
    assert!(c > 1e6, "condition number = {c}, expected very large");
}

#[test]
fn leverage_sums_to_number_of_parameters() {
    let fit = longley_fit();
    let total: f64 = fit.leverage().sum();
    assert!(
        (total - fit.n_parameters() as f64).abs() < 1e-6,
        "leverage sum = {total}, expected {}",
        fit.n_parameters()
    );
}