regression-diagnostics 0.2.0

//! Linear-mixed-model diagnostics. The exact anchor is the classical result that
//! for a **balanced** one-way random-effects design the REML variance components
//! equal the ANOVA (method-of-moments) estimators: `σ̂²_e = MSE` and
//! `σ̂²_b = (MSB − MSE)/m`. The GLS intercept is the grand mean, the BLUPs sum to
//! zero, and with no between-group signal the fit collapses to OLS.

use ndarray::{Array1, Array2};
use regression_diagnostics::glm::{GlmFit, Poisson};
use regression_diagnostics::mixed::{
    GlmmFamily, GlmmFit, LinearMixedModel, Method, MixedModel, RandomEffect,
};
use regression_diagnostics::{OlsFit, RegressionError};

/// A balanced one-way design: `g` groups of `m` observations, group `j` centered
/// at `10·j` with a fixed within-group pattern that sums to zero.
fn balanced() -> (Array2<f64>, Array1<f64>, Vec<usize>, usize, usize) {
    let g = 5usize;
    let m = 6usize;
    let noise = [-1.0, -0.6, -0.2, 0.2, 0.6, 1.0]; // sums to 0, so ȳ_j = 10j
    let n = g * m;
    let x = Array2::<f64>::ones((n, 1)); // intercept only
    let mut y = Array1::<f64>::zeros(n);
    let mut groups = vec![0usize; n];
    let mut r = 0;
    for j in 0..g {
        for &e in noise.iter() {
            y[r] = 10.0 * j as f64 + e;
            groups[r] = j;
            r += 1;
        }
    }
    (x, y, groups, g, m)
}

/// ANOVA variance-component estimators for a balanced one-way design.
fn anova_components(y: &Array1<f64>, groups: &[usize], g: usize, m: usize) -> (f64, f64, f64) {
    let n = y.len();
    let grand: f64 = y.sum() / n as f64;
    let mut group_sum = vec![0.0; g];
    for i in 0..n {
        group_sum[groups[i]] += y[i];
    }
    let group_mean: Vec<f64> = group_sum.iter().map(|s| s / m as f64).collect();
    let ssb: f64 = m as f64 * group_mean.iter().map(|gm| (gm - grand).powi(2)).sum::<f64>();
    let ssw: f64 = (0..n).map(|i| (y[i] - group_mean[groups[i]]).powi(2)).sum();
    let msb = ssb / (g - 1) as f64;
    let mse = ssw / (n - g) as f64;
    let sigma2_b = (msb - mse) / m as f64;
    (mse, sigma2_b, grand)
}

#[test]
fn reml_matches_anova_on_balanced_data() {
    let (x, y, groups, g, m) = balanced();
    let (mse, sigma2_b, grand) = anova_components(&y, &groups, g, m);

    let fit = LinearMixedModel::new(x, y, &groups).unwrap();
    assert!(
        (fit.residual_variance() - mse).abs() < 1e-3,
        "σ̂²_e = {}, ANOVA MSE = {mse}",
        fit.residual_variance()
    );
    assert!(
        (fit.group_variance() - sigma2_b).abs() < 0.1,
        "σ̂²_b = {}, ANOVA = {sigma2_b}",
        fit.group_variance()
    );
    // GLS intercept is the grand mean for a balanced design.
    assert!((fit.coefficients()[0] - grand).abs() < 1e-6);
    // ICC is high (groups are far apart) and a valid proportion.
    assert!(fit.icc() > 0.9 && fit.icc() < 1.0);
    assert_eq!(fit.n_groups(), g);
    // BLUPs sum to zero for a balanced design.
    assert!(fit.random_effects().sum().abs() < 1e-8);
}

#[test]
fn blups_track_group_means() {
    let (x, y, groups, g, m) = balanced();
    let (_, _, grand) = anova_components(&y, &groups, g, m);
    let fit = LinearMixedModel::new(x, y.clone(), &groups).unwrap();
    let b = fit.random_effects();
    // With very high ICC the shrinkage factor ≈ 1, so BLUP_j ≈ ȳ_j − grand.
    for j in 0..g {
        let group_mean = 10.0 * j as f64; // ȳ_j by construction
        assert!(
            (b[j] - (group_mean - grand)).abs() < 0.05,
            "BLUP[{j}] = {}, expected ≈ {}",
            b[j],
            group_mean - grand
        );
    }
}

#[test]
fn collapses_to_ols_without_group_signal() {
    // y driven by a covariate, group labels unrelated to y ⇒ σ̂²_b ≈ 0 and the
    // fixed effects match OLS.
    let n = 40usize;
    let mut x = Array2::<f64>::ones((n, 2));
    let mut y = Array1::<f64>::zeros(n);
    let mut groups = vec![0usize; n];
    for i in 0..n {
        let xi = (i as f64) * 0.25 - 5.0;
        x[(i, 1)] = xi;
        y[i] = 3.0 - 1.5 * xi + 0.4 * ((i % 7) as f64 - 3.0); // no group term
        groups[i] = i % 8; // 8 groups, unrelated to y
    }
    let ols = OlsFit::new(x.clone(), y.clone()).unwrap();
    let lmm = LinearMixedModel::new(x, y, &groups).unwrap();
    assert!(lmm.icc() < 0.15, "ICC = {}", lmm.icc());
    for j in 0..ols.n_parameters() {
        assert!(
            (ols.coefficients()[j] - lmm.coefficients()[j]).abs() < 0.05,
            "coef[{j}] OLS vs LMM"
        );
    }
}

#[test]
fn ml_gives_smaller_group_variance_than_reml() {
    // ML variance components are biased downward relative to REML.
    let (x, y, groups, _, _) = balanced();
    let reml = LinearMixedModel::with_method(x.clone(), y.clone(), &groups, Method::Reml).unwrap();
    let ml = LinearMixedModel::with_method(x, y, &groups, Method::Ml).unwrap();
    assert!(ml.group_variance() <= reml.group_variance() + 1e-9);
    assert!(reml.log_likelihood().is_finite() && ml.log_likelihood().is_finite());
    assert!(reml.aic().is_finite());
}

#[test]
fn general_model_reduces_to_random_intercept() {
    // A single intercept term in the general engine must match the closed-form
    // LinearMixedModel (variance components, fixed effects, log-likelihood).
    let (x, y, groups, _, _) = balanced();
    let closed = LinearMixedModel::new(x.clone(), y.clone(), &groups).unwrap();
    let general = MixedModel::new(x, y, vec![RandomEffect::intercept(&groups)]).unwrap();

    assert!(
        (general.residual_variance() - closed.residual_variance()).abs() < 1e-2,
        "σ²_e: {} vs {}",
        general.residual_variance(),
        closed.residual_variance()
    );
    assert!(
        (general.term_covariance(0)[(0, 0)] - closed.group_variance()).abs() < 0.5,
        "σ²_b: {} vs {}",
        general.term_covariance(0)[(0, 0)],
        closed.group_variance()
    );
    assert!((general.coefficients()[0] - closed.coefficients()[0]).abs() < 1e-3);
    assert!((general.log_likelihood() - closed.log_likelihood()).abs() < 1e-2);
    assert_eq!(general.n_terms(), 1);
}

#[test]
fn random_slope_model_fits() {
    // Grouped data with a group-varying slope on a covariate.
    let g = 6usize;
    let per = 12usize;
    let slope_shift = [-1.5, -0.8, 0.2, 0.7, 1.1, 0.3];
    let n = g * per;
    let mut x = Array2::<f64>::ones((n, 2)); // fixed: intercept + covariate
    let mut z = Array2::<f64>::ones((n, 2)); // random: intercept + same covariate
    let mut y = Array1::<f64>::zeros(n);
    let mut groups = vec![0usize; n];
    let mut r = 0;
    for (j, &shift) in slope_shift.iter().enumerate() {
        for i in 0..per {
            let cov = (i as f64) * 0.3 - 1.5;
            x[(r, 1)] = cov;
            z[(r, 1)] = cov;
            y[r] = 2.0 + (3.0 + shift) * cov + ((i % 3) as f64 - 1.0) * 0.3;
            groups[r] = j;
            r += 1;
        }
    }
    let fit = MixedModel::new(x, y, vec![RandomEffect::new(&groups, z)]).unwrap();
    // 2×2 covariance, both variances non-negative and finite.
    let cov = fit.term_covariance(0);
    assert_eq!(cov.shape(), [2, 2]);
    assert!(cov[(0, 0)] >= 0.0 && cov[(1, 1)] >= 0.0);
    assert!(cov.iter().all(|v| v.is_finite()));
    // BLUPs: one (intercept, slope) pair per group.
    assert_eq!(fit.random_effects(0).shape(), [g, 2]);
    assert_eq!(fit.n_random_effects(), g * 2);
    assert!(fit.residual_variance() > 0.0 && fit.log_likelihood().is_finite());
    // The slope variance should be clearly positive (slopes really do vary).
    assert!(cov[(1, 1)] > 0.1, "slope variance = {}", cov[(1, 1)]);
}

#[test]
fn crossed_random_effects_fit() {
    // Two independent grouping factors (crossed), each a random intercept.
    let n = 48usize;
    let mut x = Array2::<f64>::ones((n, 1));
    let mut y = Array1::<f64>::zeros(n);
    let mut fa = vec![0usize; n];
    let mut fb = vec![0usize; n];
    for i in 0..n {
        fa[i] = i % 4; // factor A: 4 levels
        fb[i] = i % 3; // factor B: 3 levels (crossed with A)
        let a_eff = [2.0, -1.0, 0.5, -1.5][fa[i]];
        let b_eff = [1.0, -0.5, -0.5][fb[i]];
        x[(i, 0)] = 1.0;
        y[i] = 10.0 + a_eff + b_eff + ((i % 5) as f64 - 2.0) * 0.4;
    }
    let fit = MixedModel::new(
        x,
        y,
        vec![RandomEffect::intercept(&fa), RandomEffect::intercept(&fb)],
    )
    .unwrap();
    assert_eq!(fit.n_terms(), 2);
    assert!(fit.term_covariance(0)[(0, 0)] >= 0.0);
    assert!(fit.term_covariance(1)[(0, 0)] >= 0.0);
    assert_eq!(fit.random_effects(0).shape(), [4, 1]);
    assert_eq!(fit.random_effects(1).shape(), [3, 1]);
    assert!(fit.log_likelihood().is_finite());
}

#[test]
fn poisson_glmm_approaches_glm_without_group_signal() {
    // Counts follow a covariate-driven log-mean; group labels are unrelated, so
    // the random-intercept SD should be small and β close to the plain GLM.
    let n = 48usize;
    let mut x = Array2::<f64>::ones((n, 2));
    let mut y = Array1::<f64>::zeros(n);
    let mut groups = vec![0usize; n];
    for i in 0..n {
        let xi = (i as f64) * 0.08 - 1.9;
        x[(i, 1)] = xi;
        let mu = (0.6 + 0.5 * xi).exp();
        // Low-noise counts around the covariate trend, no group offset.
        y[i] = (mu + ((i % 3) as f64 - 1.0) * 0.5).round().max(0.0);
        groups[i] = i % 8; // interleaved, unrelated to the trend
    }
    let glm = GlmFit::new(Poisson, x.clone(), y.clone()).unwrap();
    let glmm = GlmmFit::new(x, y, &groups, GlmmFamily::Poisson).unwrap();
    assert!(glmm.sigma_b() < 0.5, "σ_b = {}", glmm.sigma_b());
    for j in 0..glm.n_parameters() {
        assert!(
            (glm.coefficients()[j] - glmm.coefficients()[j]).abs() < 0.25,
            "coef[{j}] GLM {} vs GLMM {}",
            glm.coefficients()[j],
            glmm.coefficients()[j]
        );
    }
    assert!(glmm.log_likelihood().is_finite());
}

#[test]
fn poisson_glmm_detects_group_variance() {
    // Strong group-level offsets ⇒ a clearly positive σ_b, larger than the
    // no-signal case.
    let g = 8usize;
    let per = 8usize;
    let offset = [-0.9, 0.8, -0.5, 0.6, -0.7, 0.9, -0.3, 0.4];
    let n = g * per;
    let mut x = Array2::<f64>::ones((n, 2));
    let mut y = Array1::<f64>::zeros(n);
    let mut groups = vec![0usize; n];
    let mut r = 0;
    for (j, &off) in offset.iter().enumerate() {
        for i in 0..per {
            let xi = (i as f64) * 0.2 - 0.7;
            x[(r, 1)] = xi;
            let mu = (0.8 + 0.3 * xi + off).exp();
            y[r] = mu.round().max(0.0);
            groups[r] = j;
            r += 1;
        }
    }
    let glmm = GlmmFit::new(x, y, &groups, GlmmFamily::Poisson).unwrap();
    assert!(glmm.sigma_b() > 0.25, "σ_b = {}", glmm.sigma_b());
    assert!(glmm.coefficient_standard_errors().iter().all(|v| v.is_finite()));
    assert!(glmm.aic().is_finite());
    assert_eq!(glmm.random_effects().len(), g);
}

#[test]
fn binomial_glmm_fits() {
    let g = 10usize;
    let per = 10usize;
    let offset = [-2.0, 1.5, -1.0, 2.0, -1.5, 1.0, -0.5, 0.5, -2.0, 1.5];
    let n = g * per;
    let mut x = Array2::<f64>::ones((n, 2));
    let mut y = Array1::<f64>::zeros(n);
    let mut groups = vec![0usize; n];
    let mut r = 0;
    for (j, &off) in offset.iter().enumerate() {
        for i in 0..per {
            let xi = (i as f64) * 0.3 - 1.5;
            x[(r, 1)] = xi;
            let eta = 0.4 * xi + off;
            y[r] = if eta + ((i % 2) as f64 - 0.5) > 0.0 { 1.0 } else { 0.0 };
            groups[r] = j;
            r += 1;
        }
    }
    let glmm = GlmmFit::new(x, y, &groups, GlmmFamily::Binomial).unwrap();
    assert!(glmm.sigma_b() > 0.0);
    assert!(glmm.log_likelihood().is_finite());
    assert_eq!(glmm.n_groups(), g);
    assert!(glmm.p_values().iter().all(|v| v.is_finite()));
}

#[test]
fn glmm_rejects_bad_input() {
    let x = Array2::<f64>::ones((4, 1));
    // Negative Poisson count.
    let bad = Array1::from(vec![1.0, 2.0, -1.0, 0.0]);
    assert!(matches!(
        GlmmFit::new(x.clone(), bad, &[0, 1, 0, 1], GlmmFamily::Poisson).unwrap_err(),
        RegressionError::InvalidResponse { .. }
    ));
    // Only one group.
    let ok = Array1::from(vec![0.0, 1.0, 1.0, 0.0]);
    assert!(matches!(
        GlmmFit::new(x, ok, &[3, 3, 3, 3], GlmmFamily::Binomial).unwrap_err(),
        RegressionError::InvalidResponse { .. }
    ));
}

#[test]
fn rejects_degenerate_input() {
    let x = Array2::<f64>::ones((4, 1));
    let y = Array1::from(vec![1.0, 2.0, 3.0, 4.0]);
    // Only one group.
    let one_group = vec![0usize, 0, 0, 0];
    assert!(matches!(
        LinearMixedModel::new(x.clone(), y.clone(), &one_group).unwrap_err(),
        RegressionError::InvalidResponse { .. }
    ));
    // n <= p.
    let x_wide = Array2::<f64>::ones((2, 3));
    let y2 = Array1::from(vec![1.0, 2.0]);
    let groups2 = vec![0usize, 1];
    assert!(matches!(
        LinearMixedModel::new(x_wide, y2, &groups2).unwrap_err(),
        RegressionError::NoResidualDegreesOfFreedom { .. }
    ));
}