regression-diagnostics 0.2.0

//! Generalized-linear-model diagnostics (Poisson / negative binomial / Gamma).
//! The correctness anchors are exact where a closed form exists: the Poisson MLE
//! satisfies the canonical score equations `Xᵀ(y − μ) = 0`, a group-indicator
//! design reproduces the group means on the log scale, deviance residuals square
//! to the residual deviance, GLM leverage sums to `p`, and the negative binomial
//! collapses onto Poisson as `θ → ∞`.

use ndarray::{Array1, Array2};
use regression_diagnostics::glm::{
    cooks_distance, deviance_residuals, fit_negative_binomial, leverage, pearson_residuals, Family,
    Gamma, GlmFit, NegativeBinomial, Poisson,
};
use regression_diagnostics::RegressionError;

/// A count response with genuine log-linear signal in one predictor.
fn count_dataset() -> (Array2<f64>, Array1<f64>) {
    let counts = [1.0, 1.0, 2.0, 2.0, 3.0, 4.0, 5.0, 6.0, 9.0, 12.0];
    let n = counts.len();
    let mut x = Array2::<f64>::ones((n, 2));
    let mut y = Array1::<f64>::zeros(n);
    for i in 0..n {
        x[(i, 1)] = i as f64 * 0.3;
        y[i] = counts[i];
    }
    (x, y)
}

#[test]
fn poisson_score_equations_hold_at_mle() {
    let (x, y) = count_dataset();
    let fit = GlmFit::new(Poisson, x.clone(), y.clone()).unwrap();
    // Canonical log link: the score is Xᵀ(y − μ), zero componentwise at the MLE.
    let mu = fit.fitted_means();
    let resid: Array1<f64> = &y - &mu;
    for j in 0..fit.n_parameters() {
        let g: f64 = (0..fit.n_observations())
            .map(|i| x[(i, j)] * resid[i])
            .sum();
        assert!(g.abs() < 1e-6, "score[{j}] = {g}");
    }
}

#[test]
fn poisson_group_indicator_reproduces_group_means() {
    // Two groups with a 0/1 indicator make the two-parameter model saturated,
    // so the fitted mean per group is exactly the group's sample mean, and on
    // the log scale intercept = ln(mean0), slope = ln(mean1/mean0).
    let group0 = [1.0, 2.0, 3.0, 2.0, 2.0]; // mean 2.0
    let group1 = [3.0, 5.0, 7.0, 5.0, 5.0]; // mean 5.0
    let n = group0.len() + group1.len();
    let mut x = Array2::<f64>::ones((n, 2));
    let mut y = Array1::<f64>::zeros(n);
    for (i, &c) in group0.iter().chain(group1.iter()).enumerate() {
        x[(i, 1)] = if i < group0.len() { 0.0 } else { 1.0 };
        y[i] = c;
    }
    let fit = GlmFit::new(Poisson, x, y).unwrap();
    let want_intercept = 2.0_f64.ln();
    let want_slope = (5.0_f64 / 2.0).ln();
    assert!(
        (fit.coefficients()[0] - want_intercept).abs() < 1e-8,
        "intercept {} vs {want_intercept}",
        fit.coefficients()[0]
    );
    assert!(
        (fit.coefficients()[1] - want_slope).abs() < 1e-8,
        "slope {} vs {want_slope}",
        fit.coefficients()[1]
    );
}

#[test]
fn deviance_residuals_square_to_residual_deviance() {
    for label in ["poisson", "gamma", "negbin"] {
        let (x, y) = count_dataset();
        let (dr, gof_dev) = match label {
            "poisson" => {
                let f = GlmFit::new(Poisson, x, y).unwrap();
                (deviance_residuals(&f), f.goodness_of_fit().residual_deviance)
            }
            "gamma" => {
                let f = GlmFit::new(Gamma, x, y).unwrap();
                (deviance_residuals(&f), f.goodness_of_fit().residual_deviance)
            }
            _ => {
                let f = GlmFit::new(NegativeBinomial::new(4.0).unwrap(), x, y).unwrap();
                (deviance_residuals(&f), f.goodness_of_fit().residual_deviance)
            }
        };
        let sum_sq: f64 = dr.iter().map(|d| d * d).sum();
        assert!(
            (sum_sq - gof_dev).abs() < 1e-8,
            "{label}: Σdev² = {sum_sq} vs residual_deviance = {gof_dev}"
        );
    }
}

#[test]
fn gamma_pearson_chi2_equals_dispersion_times_df() {
    let (x, y) = count_dataset();
    let fit = GlmFit::new(Gamma, x, y).unwrap();
    let pr = pearson_residuals(&fit);
    let chi2: f64 = pr.iter().map(|r| r * r).sum();
    let df = (fit.n_observations() - fit.n_parameters()) as f64;
    // The Gamma dispersion is defined as the Pearson χ² divided by n − p.
    assert!(
        (chi2 - fit.dispersion() * df).abs() < 1e-8,
        "χ² = {chi2} vs φ̂·df = {}",
        fit.dispersion() * df
    );
    assert!(fit.dispersion() > 0.0);
    // Estimated-dispersion family ⇒ Wald p-values come off Student's t, finite.
    assert!(fit.p_values().iter().all(|v| v.is_finite()));
}

#[test]
fn glm_leverage_sums_to_parameters() {
    for label in ["poisson", "gamma", "negbin"] {
        let (x, y) = count_dataset();
        let total = match label {
            "poisson" => leverage(&GlmFit::new(Poisson, x, y).unwrap()).sum(),
            "gamma" => leverage(&GlmFit::new(Gamma, x, y).unwrap()).sum(),
            _ => leverage(&GlmFit::new(NegativeBinomial::new(4.0).unwrap(), x, y).unwrap()).sum(),
        };
        assert!(
            (total - 2.0).abs() < 1e-6,
            "{label}: Σh = {total}, expected 2"
        );
    }
}

#[test]
fn goodness_of_fit_is_sane() {
    let (x, y) = count_dataset();
    let fit = GlmFit::new(Poisson, x, y).unwrap();
    let gof = fit.goodness_of_fit();
    assert!(gof.residual_deviance <= gof.null_deviance + 1e-9);
    assert!(gof.mcfadden_r2 > 0.0 && gof.mcfadden_r2 < 1.0);
    assert!(gof.aic.is_finite() && gof.bic.is_finite());
    assert_eq!(gof.dispersion, 1.0); // Poisson dispersion is fixed
    // Cook's distances are finite and non-negative.
    assert!(cooks_distance(&fit).iter().all(|v| v.is_finite() && *v >= 0.0));
}

#[test]
fn gamma_charges_a_parameter_for_estimated_dispersion() {
    // AIC = −2ℓ + 2k with k = p + 1 for Gamma (the estimated φ), vs k = p for
    // the fixed-dispersion families. Check the criterion matches the definition.
    let (x, y) = count_dataset();
    let fit = GlmFit::new(Gamma, x.clone(), y.clone()).unwrap();
    let gof = fit.goodness_of_fit();
    let k = fit.n_parameters() as f64 + 1.0;
    let want_aic = -2.0 * fit.log_likelihood() + 2.0 * k;
    assert!((gof.aic - want_aic).abs() < 1e-9, "{} vs {want_aic}", gof.aic);
}

#[test]
fn negative_binomial_approaches_poisson_as_theta_grows() {
    let (x, y) = count_dataset();
    let poisson = GlmFit::new(Poisson, x.clone(), y.clone()).unwrap();
    let nb = GlmFit::new(NegativeBinomial::new(1e8).unwrap(), x, y).unwrap();
    for j in 0..poisson.n_parameters() {
        let diff = (poisson.coefficients()[j] - nb.coefficients()[j]).abs();
        assert!(diff < 1e-4, "coef[{j}] diverges: {diff}");
    }
}

#[test]
fn predict_mean_reproduces_training_fits() {
    let (x, y) = count_dataset();
    let fit = GlmFit::new(Poisson, x.clone(), y).unwrap();
    let pred = fit.predict_mean(x.view());
    let mu = fit.fitted_means();
    for i in 0..fit.n_observations() {
        assert!((pred[i] - mu[i]).abs() < 1e-10);
    }
}

#[test]
fn out_of_support_responses_are_rejected() {
    let x = Array2::<f64>::ones((4, 2));
    // Negative count for Poisson.
    let neg = Array1::from(vec![1.0, 2.0, -1.0, 3.0]);
    assert!(matches!(
        GlmFit::new(Poisson, x.clone(), neg).unwrap_err(),
        RegressionError::InvalidResponse { .. }
    ));
    // Zero is outside the Gamma support (strictly positive).
    let zero = Array1::from(vec![1.0, 2.0, 0.0, 3.0]);
    assert!(matches!(
        GlmFit::new(Gamma, x, zero).unwrap_err(),
        RegressionError::InvalidResponse { .. }
    ));
}

#[test]
fn negative_binomial_theta_estimation() {
    // Over-dispersed counts: two observations per x level, deliberately spread.
    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 fit = fit_negative_binomial(x.clone(), y.clone()).unwrap();
    let theta_hat = fit.family().theta();
    // Over-dispersion ⇒ a finite, moderate θ (not pinned to the Poisson cap).
    assert!(theta_hat > 0.0 && theta_hat < 1e5, "θ̂ = {theta_hat}");
    // The estimate maximizes the profile: its log-lik beats a very different θ.
    let ll_hat = fit.log_likelihood();
    let ll_off = GlmFit::new(NegativeBinomial::new(0.05).unwrap(), x, y)
        .unwrap()
        .log_likelihood();
    assert!(ll_hat >= ll_off - 1e-6, "profile not maximized: {ll_hat} vs {ll_off}");
}

#[test]
fn negative_binomial_theta_must_be_positive() {
    assert!(matches!(
        NegativeBinomial::new(0.0).unwrap_err(),
        RegressionError::InvalidParameter { .. }
    ));
    assert!(matches!(
        NegativeBinomial::new(-2.0).unwrap_err(),
        RegressionError::InvalidParameter { .. }
    ));
    // Sanity: the variance function exceeds Poisson's μ for finite θ.
    let nb = NegativeBinomial::new(2.0).unwrap();
    assert!(nb.variance(3.0) > 3.0);
    assert_eq!(nb.theta(), 2.0);
}