regression-diagnostics 0.2.0

//! Ridge and lasso diagnostics: known-answer checks anchored on the identities
//! that pin these estimators (λ = 0 reduces to OLS; shrinkage is monotone).

use ndarray::{array, Array1, Array2};
use regression_diagnostics::multicollinearity::vif;
use regression_diagnostics::regularized::{select_lambda_gcv, LassoFit, RidgeFit};
use regression_diagnostics::OlsFit;

fn collinear_dataset() -> (Array2<f64>, Array1<f64>) {
    // Intercept + two predictors, x2 strongly collinear with x1.
    let x = array![
        [1.0, 1.0, 2.05],
        [1.0, 2.0, 3.98],
        [1.0, 3.0, 6.02],
        [1.0, 4.0, 8.01],
        [1.0, 5.0, 9.97],
        [1.0, 6.0, 12.03],
        [1.0, 7.0, 13.98],
        [1.0, 8.0, 16.02],
        [1.0, 9.0, 17.99],
        [1.0, 10.0, 20.05],
    ];
    let y = array![2.1, 3.9, 6.2, 7.8, 10.1, 11.9, 14.2, 15.8, 18.1, 19.9];
    (x, y)
}

#[test]
fn ridge_at_zero_lambda_equals_ols() {
    let (x, y) = collinear_dataset();
    let ols = OlsFit::new(x.clone(), y.clone()).unwrap();
    let ridge = RidgeFit::new(x, y, 0.0).unwrap();
    for j in 0..ols.n_parameters() {
        assert!(
            (ols.coefficients()[j] - ridge.coefficients()[j]).abs() < 1e-7,
            "coef[{j}]: OLS {} vs ridge(0) {}",
            ols.coefficients()[j],
            ridge.coefficients()[j]
        );
    }
}

#[test]
fn ridge_effective_df_and_leverage_identity() {
    let (x, y) = collinear_dataset();
    let ridge0 = RidgeFit::new(x.clone(), y.clone(), 0.0).unwrap();
    // At λ = 0, effective df == number of parameters.
    assert!((ridge0.effective_df() - ridge0.n_parameters() as f64).abs() < 1e-6);

    // Leverage sums to the effective df (the ridge analogue of Σh = p).
    for &lam in &[0.0, 1.0, 50.0] {
        let r = RidgeFit::new(x.clone(), y.clone(), lam).unwrap();
        let sum: f64 = r.leverage().sum();
        assert!(
            (sum - r.effective_df()).abs() < 1e-6,
            "λ={lam}: Σh={sum} vs edf={}",
            r.effective_df()
        );
    }
}

#[test]
fn ridge_shrinks_with_lambda() {
    let (x, y) = collinear_dataset();
    let coef_norm = |lam: f64| {
        let r = RidgeFit::new(x.clone(), y.clone(), lam).unwrap();
        // Slope magnitude (exclude intercept at column 0).
        (r.coefficients()[1].powi(2) + r.coefficients()[2].powi(2)).sqrt()
    };
    let edf = |lam: f64| {
        RidgeFit::new(x.clone(), y.clone(), lam)
            .unwrap()
            .effective_df()
    };

    assert!(coef_norm(0.0) > coef_norm(10.0));
    assert!(coef_norm(10.0) > coef_norm(1000.0));
    assert!(edf(0.0) > edf(10.0));
    assert!(edf(10.0) > edf(1000.0));
}

#[test]
fn ridge_vif_matches_ols_at_zero_and_falls_with_lambda() {
    let (x, y) = collinear_dataset();
    let ols = OlsFit::new(x.clone(), y.clone()).unwrap();
    let ols_vif = vif(&ols);

    let r0 = RidgeFit::new(x.clone(), y.clone(), 0.0).unwrap();
    let rv0 = r0.ridge_vif();
    // At λ = 0 the ridge VIF equals the OLS VIF.
    for j in [1usize, 2] {
        assert!(
            (rv0[j] - ols_vif[j]).abs() < 1e-6,
            "col {j}: ridge_vif(0)={} vs ols_vif={}",
            rv0[j],
            ols_vif[j]
        );
    }
    // Regularization tames collinearity: VIF drops as λ grows.
    let rv_big = RidgeFit::new(x, y, 5.0).unwrap().ridge_vif();
    assert!(rv_big[1] < rv0[1], "ridge VIF should fall with λ");
    assert!(rv0[0].is_nan()); // intercept slot
}

#[test]
fn gcv_selects_a_finite_lambda() {
    let (x, y) = collinear_dataset();
    let grid: Vec<f64> = (0..7).map(|k: i32| 10f64.powi(k - 3)).collect(); // 1e-3 .. 1e3
    let best = select_lambda_gcv(x, y, &grid).unwrap();
    assert!(best.lambda().is_finite());
    assert!(best.gcv().is_finite());
}

#[test]
fn lasso_at_zero_lambda_matches_ols() {
    // A well-conditioned design: coordinate descent converges quickly to the OLS
    // solution at λ = 0. (On a strongly collinear design pure-OLS coordinate
    // descent converges only very slowly — that is a property of the algorithm,
    // and OLS is the right tool there anyway.)
    let n = 30usize;
    let mut x = Array2::<f64>::ones((n, 3));
    let mut y = Array1::<f64>::zeros(n);
    for i in 0..n {
        let a = (i as f64 * 0.37).sin() * 3.0;
        let b = (i as f64 * 0.91).cos() * 2.0;
        x[(i, 1)] = a;
        x[(i, 2)] = b;
        y[i] = 4.0 - 2.0 * a + 1.5 * b;
    }
    let ols = OlsFit::new(x.clone(), y.clone()).unwrap();
    let lasso = LassoFit::with_options(x, y, 0.0, 1e-11, 200_000).unwrap();
    for j in 0..ols.n_parameters() {
        assert!(
            (ols.coefficients()[j] - lasso.coefficients()[j]).abs() < 1e-4,
            "coef[{j}]: OLS {} vs lasso(0) {}",
            ols.coefficients()[j],
            lasso.coefficients()[j]
        );
    }
}

#[test]
fn lasso_large_lambda_zeros_all_predictors() {
    let (x, y) = collinear_dataset();
    let lasso = LassoFit::new(x, y.clone(), 1e6).unwrap();
    assert_eq!(lasso.n_nonzero(), 0);
    assert!(lasso.active_set().is_empty());
    // Only the intercept survives, at the response mean.
    let ybar = y.sum() / y.len() as f64;
    assert!((lasso.coefficients()[0] - ybar).abs() < 1e-6);
    assert!(lasso.coefficients()[1].abs() < 1e-12);
    assert!(lasso.coefficients()[2].abs() < 1e-12);
}

#[test]
fn lasso_sparsity_increases_with_lambda() {
    // A design where predictors contribute at different strengths, so lasso drops
    // them one at a time as λ rises.
    let n = 40usize;
    let mut x = Array2::<f64>::ones((n, 4));
    let mut y = Array1::<f64>::zeros(n);
    for i in 0..n {
        let a = (i as f64 * 0.3).sin();
        let b = (i as f64 * 0.11).cos();
        let c = ((i % 7) as f64) - 3.0;
        x[(i, 1)] = a;
        x[(i, 2)] = b;
        x[(i, 3)] = c;
        // y depends strongly on a, weakly on b, not at all on c.
        y[i] = 1.0 + 5.0 * a + 0.4 * b;
    }
    let small = LassoFit::new(x.clone(), y.clone(), 0.001).unwrap();
    let big = LassoFit::new(x, y, 0.5).unwrap();
    assert!(
        big.n_nonzero() <= small.n_nonzero(),
        "more penalty should not increase the active set: {} vs {}",
        big.n_nonzero(),
        small.n_nonzero()
    );
    assert_eq!(small.active_set().len(), small.n_nonzero());
}