regression-diagnostics 0.2.0

//! Elastic-net and penalized-GLM diagnostics. Anchored on exact reductions:
//! elastic net at `α = 1` is the lasso and at `λ = 0` is OLS; ridge-penalized
//! logistic at `λ = 0` is ordinary logistic. The shrinkage-aware degrees of
//! freedom behave monotonically in the penalty.

use ndarray::{Array1, Array2};
use regression_diagnostics::logistic::LogisticFit;
use regression_diagnostics::regularized::{ElasticNetFit, LassoFit, PenalizedLogisticFit};
use regression_diagnostics::{OlsFit, RegressionError};

/// A small design with an intercept and three correlated predictors.
fn design() -> (Array2<f64>, Array1<f64>) {
    let n = 30usize;
    let mut x = Array2::<f64>::ones((n, 4));
    let mut y = Array1::<f64>::zeros(n);
    for i in 0..n {
        let t = i as f64;
        x[(i, 1)] = (t * 0.3).sin();
        x[(i, 2)] = (t * 0.3).sin() + 0.05 * t; // correlated with col 1
        x[(i, 3)] = ((t * 0.17).cos()) * 2.0;
        y[i] = 1.0 + 2.0 * x[(i, 1)] - 1.5 * x[(i, 3)] + 0.1 * ((i % 5) as f64 - 2.0);
    }
    (x, y)
}

fn binary_design() -> (Array2<f64>, Array1<f64>) {
    let n = 40usize;
    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.2 - 4.0;
        x[(i, 1)] = a;
        x[(i, 2)] = (a * 0.7).cos();
        y[i] = if (i % 3 == 0) ^ (a > 0.0) { 1.0 } else { 0.0 };
    }
    (x, y)
}

#[test]
fn elastic_net_alpha_one_equals_lasso() {
    let (x, y) = design();
    let lambda = 0.1;
    let lasso = LassoFit::new(x.clone(), y.clone(), lambda).unwrap();
    let enet = ElasticNetFit::new(x, y, lambda, 1.0).unwrap();
    for j in 0..lasso.n_parameters() {
        let diff = (lasso.coefficients()[j] - enet.coefficients()[j]).abs();
        assert!(diff < 1e-8, "coef[{j}] enet vs lasso: {diff}");
    }
}

#[test]
fn elastic_net_lambda_zero_equals_ols() {
    let (x, y) = design();
    let ols = OlsFit::new(x.clone(), y.clone()).unwrap();
    // Any alpha at lambda = 0 is unpenalized least squares.
    let enet = ElasticNetFit::new(x, y, 0.0, 0.5).unwrap();
    for j in 0..ols.n_parameters() {
        let diff = (ols.coefficients()[j] - enet.coefficients()[j]).abs();
        assert!(diff < 1e-5, "coef[{j}] enet vs OLS: {diff}");
    }
}

#[test]
fn elastic_net_df_reflects_ridge_shrinkage() {
    let (x, y) = design();
    let lambda = 0.2;
    // Pure lasso: df is exactly the active-set size (+1 intercept).
    let lasso_like = ElasticNetFit::new(x.clone(), y.clone(), lambda, 1.0).unwrap();
    let expected = lasso_like.n_nonzero() as f64 + 1.0;
    assert!((lasso_like.effective_df() - expected).abs() < 1e-9);

    // Adding ridge (alpha < 1) spends fewer df than the active-set count for the
    // same surviving predictors, and never more than p.
    let blended = ElasticNetFit::new(x, y, lambda, 0.3).unwrap();
    let active_plus_intercept = blended.n_nonzero() as f64 + 1.0;
    assert!(blended.effective_df() <= active_plus_intercept + 1e-9);
    assert!(blended.effective_df() <= blended.n_parameters() as f64 + 1e-9);
    assert!(blended.effective_df() > 1.0); // more than just the intercept
}

#[test]
fn penalized_logistic_lambda_zero_equals_logistic() {
    let (x, y) = binary_design();
    let logit = LogisticFit::new(x.clone(), y.clone()).unwrap();
    let pen = PenalizedLogisticFit::new(x, y, 0.0).unwrap();
    for j in 0..logit.n_parameters() {
        let diff = (logit.coefficients()[j] - pen.coefficients()[j]).abs();
        assert!(diff < 1e-7, "coef[{j}] penalized vs logistic: {diff}");
    }
    // With no penalty the effective df is the full parameter count.
    assert!((pen.effective_df() - pen.n_parameters() as f64).abs() < 1e-6);
}

#[test]
fn penalized_logistic_shrinks_with_lambda() {
    let (x, y) = binary_design();
    let light = PenalizedLogisticFit::new(x.clone(), y.clone(), 1.0).unwrap();
    let heavy = PenalizedLogisticFit::new(x, y, 50.0).unwrap();

    // Effective df decreases monotonically toward the intercept as λ grows.
    assert!(heavy.effective_df() < light.effective_df());
    assert!(heavy.effective_df() < heavy.n_parameters() as f64);
    assert!(heavy.effective_df() > 0.0);

    // Penalized slopes are shrunk toward zero (compare the non-intercept norm).
    let norm = |f: &PenalizedLogisticFit| -> f64 {
        (1..f.n_parameters()).map(|j| f.coefficients()[j].powi(2)).sum()
    };
    assert!(norm(&heavy) < norm(&light));
    // AIC/SEs are finite.
    assert!(heavy.aic().is_finite());
    assert!(heavy.coefficient_standard_errors().iter().all(|v| v.is_finite()));
}

#[test]
fn invalid_parameters_are_rejected() {
    let (x, y) = design();
    assert!(matches!(
        ElasticNetFit::new(x.clone(), y.clone(), -0.1, 0.5).unwrap_err(),
        RegressionError::InvalidParameter { .. }
    ));
    assert!(matches!(
        ElasticNetFit::new(x.clone(), y.clone(), 0.1, 1.5).unwrap_err(),
        RegressionError::InvalidParameter { .. }
    ));
    let (xb, yb) = binary_design();
    assert!(matches!(
        PenalizedLogisticFit::new(xb, yb, -1.0).unwrap_err(),
        RegressionError::InvalidParameter { .. }
    ));
}