regression-diagnostics 0.2.0

//! Generalized-linear-model diagnostics for count data: fit a Poisson model by
//! IRLS, print the coefficient table and goodness-of-fit summary, check for
//! over-dispersion, then refit as a negative binomial and compare by AIC.
//!
//! Run with: `cargo run --example glm_diagnostics`

use ndarray::{Array1, Array2};
use regression_diagnostics::glm::{
    cooks_distance, fit_negative_binomial, leverage, GlmFit, Poisson,
};

fn print_coef_table<F: regression_diagnostics::glm::Family>(fit: &GlmFit<F>, names: &[&str]) {
    let se = fit.coefficient_standard_errors();
    let stat = fit.wald_statistics();
    let pv = fit.p_values();
    println!(
        "{:<8}{:>10}{:>10}{:>9}{:>9}",
        "", "coef", "std err", "stat", "P>|.|"
    );
    println!("{:-<46}", "");
    for j in 0..fit.n_parameters() {
        println!(
            "{:<8}{:>10.4}{:>10.4}{:>9.3}{:>9.3}",
            names[j],
            fit.coefficients()[j],
            se[j],
            stat[j],
            pv[j]
        );
    }
}

fn main() {
    // Synthetic counts with an increasing log-linear trend but far more scatter
    // at each covariate level than Poisson (mean = variance) can absorb — two
    // observations per x, deliberately spread apart, so the fit is
    // over-dispersed and the negative binomial should win on AIC.
    let rows: [(f64, f64); 16] = [
        (0.0, 0.0),
        (0.0, 4.0),
        (0.5, 1.0),
        (0.5, 6.0),
        (1.0, 0.0),
        (1.0, 8.0),
        (1.5, 2.0),
        (1.5, 10.0),
        (2.0, 1.0),
        (2.0, 13.0),
        (2.5, 3.0),
        (2.5, 17.0),
        (3.0, 2.0),
        (3.0, 22.0),
        (3.5, 5.0),
        (3.5, 28.0),
    ];
    let n = rows.len();
    let mut x = Array2::<f64>::ones((n, 2));
    let mut y = Array1::<f64>::zeros(n);
    for (i, &(xi, yi)) in rows.iter().enumerate() {
        x[(i, 1)] = xi;
        y[i] = yi;
    }
    let names = ["const", "x"];

    // --- Poisson fit -------------------------------------------------------
    let poisson = GlmFit::new(Poisson, x.clone(), y.clone()).expect("converges");
    println!("Poisson regression (log link, {} IRLS iters)\n", poisson.iterations());
    print_coef_table(&poisson, &names);

    let gof = poisson.goodness_of_fit();
    println!("\nGoodness of fit:");
    println!(
        "  Null deviance:     {:>8.3}  (df {})",
        gof.null_deviance, gof.df_null as usize
    );
    println!(
        "  Residual deviance: {:>8.3}  (df {})",
        gof.residual_deviance, gof.df_residual as usize
    );
    println!("  McFadden pseudo-R²:{:>8.4}", gof.mcfadden_r2);
    println!("  AIC / BIC:         {:>8.2} / {:.2}", gof.aic, gof.bic);

    // Over-dispersion check: residual deviance / df ≫ 1 means Poisson's
    // mean = variance assumption is violated.
    let dispersion_ratio = gof.residual_deviance / gof.df_residual;
    println!(
        "\n  Deviance/df = {:.2}  ({})",
        dispersion_ratio,
        if dispersion_ratio > 1.5 {
            "over-dispersed — try negative binomial"
        } else {
            "no strong over-dispersion"
        }
    );

    // --- Negative-binomial refit, estimating θ ----------------------------
    // fit_negative_binomial profiles the likelihood over the dispersion θ rather
    // than requiring it up front. Compare models by AIC (lower is better).
    let nb = fit_negative_binomial(x.clone(), y.clone()).expect("converges");
    let nb_gof = nb.goodness_of_fit();
    println!(
        "\nNegative binomial refit (estimated θ̂ = {:.3})\n",
        nb.family().theta()
    );
    print_coef_table(&nb, &names);
    println!(
        "\n  AIC: Poisson {:.2}  vs  NegBin {:.2}  →  {} preferred",
        gof.aic,
        nb_gof.aic,
        if nb_gof.aic < gof.aic { "NegBin" } else { "Poisson" }
    );

    // --- Influence ---------------------------------------------------------
    let cd = cooks_distance(&poisson);
    let lev = leverage(&poisson);
    let mut idx: Vec<usize> = (0..n).collect();
    idx.sort_by(|&a, &b| cd[b].partial_cmp(&cd[a]).unwrap());
    println!("\nMost influential observations, Poisson fit (Cook's distance):");
    println!("  obs   leverage    cooks_d       y");
    for &i in idx.iter().take(3) {
        println!("  {:>3}   {:>8.4}   {:>8.4}   {:>5.0}", i, lev[i], cd[i], y[i]);
    }
}