regression-diagnostics 0.2.0

//! Elastic-net and penalized-GLM diagnostics: fit an elastic net across the
//! lasso→ridge spectrum and watch the active set and effective degrees of
//! freedom respond, then ridge-penalize a logistic regression and see the
//! effective df shrink from p.
//!
//! Run with: `cargo run --example elastic_net_diagnostics`

use ndarray::{Array1, Array2};
use regression_diagnostics::regularized::{ElasticNetFit, PenalizedLogisticFit};

fn main() {
    // Linear problem: 6 predictors, only two truly active, two of the noise
    // predictors correlated with the signal (where elastic net beats lasso).
    let n = 50usize;
    let p = 7usize; // intercept + 6 predictors
    let mut x = Array2::<f64>::ones((n, p));
    let mut y = Array1::<f64>::zeros(n);
    for i in 0..n {
        let t = i as f64;
        let s1 = (t * 0.2).sin();
        let s2 = (t * 0.13).cos();
        x[(i, 1)] = s1;
        x[(i, 2)] = s1 + 0.1 * (t * 0.5).sin(); // correlated with col 1
        x[(i, 3)] = s2;
        x[(i, 4)] = (t * 0.07).sin();
        x[(i, 5)] = (t * 0.29).cos();
        x[(i, 6)] = ((i % 9) as f64 - 4.0) * 0.3;
        y[i] = 2.0 * s1 - 1.5 * s2 + 0.15 * ((i % 5) as f64 - 2.0);
    }

    println!("Elastic net across the lasso→ridge spectrum (λ = 0.1):\n");
    println!("{:<8}{:>10}{:>14}{:>10}", "alpha", "n_active", "effective_df", "AIC");
    for &alpha in &[1.0, 0.7, 0.4, 0.1] {
        let fit = ElasticNetFit::new(x.clone(), y.clone(), 0.1, alpha).unwrap();
        println!(
            "{:<8.1}{:>10}{:>14.2}{:>10.2}",
            alpha,
            fit.n_nonzero(),
            fit.effective_df(),
            fit.aic()
        );
    }
    let enet = ElasticNetFit::new(x.clone(), y.clone(), 0.1, 0.5).unwrap();
    println!("\n  At α = 0.5 the active set is columns {:?}.", enet.active_set());

    // --- Penalized logistic ------------------------------------------------
    let mut xb = Array2::<f64>::ones((n, 4));
    let mut yb = Array1::<f64>::zeros(n);
    for i in 0..n {
        let a = (i as f64) * 0.16 - 4.0;
        xb[(i, 1)] = a;
        xb[(i, 2)] = a + 0.2 * ((i % 4) as f64 - 1.5); // collinear with col 1
        xb[(i, 3)] = (a * 0.6).sin();
        yb[i] = if (i % 3 == 0) ^ (a > 0.0) { 1.0 } else { 0.0 };
    }
    println!("\nRidge-penalized logistic (p = {} incl. intercept):\n", xb.ncols());
    println!("{:<8}{:>14}{:>12}{:>10}", "lambda", "effective_df", "‖slopes‖", "AIC");
    for &lambda in &[0.0, 1.0, 5.0, 25.0] {
        let fit = PenalizedLogisticFit::new(xb.clone(), yb.clone(), lambda).unwrap();
        let slope_norm: f64 = (1..fit.n_parameters())
            .map(|j| fit.coefficients()[j].powi(2))
            .sum::<f64>()
            .sqrt();
        println!(
            "{:<8.1}{:>14.3}{:>12.3}{:>10.2}",
            lambda,
            fit.effective_df(),
            slope_norm,
            fit.aic()
        );
    }
    println!(
        "\n  As λ grows the effective df falls from {} toward the intercept and the\n  \
         coefficient norm shrinks — the bias/variance trade the penalty buys.",
        xb.ncols()
    );
}