regression-diagnostics 0.2.0

//! Random-intercept linear mixed model on grouped data: recover the variance
//! components, the intraclass correlation, and the shrinkage BLUPs, and see how
//! REML fixed-effect standard errors differ from the (too-optimistic) OLS ones.
//!
//! Run with: `cargo run --example mixed_model`

use ndarray::{Array1, Array2};
use regression_diagnostics::mixed::{GlmmFamily, GlmmFit, LinearMixedModel};
use regression_diagnostics::OlsFit;

fn main() {
    // Six schools, ten students each; a school-level intercept shift plus a
    // student-level covariate (hours studied) and within-school noise.
    let n_groups = 6usize;
    let per = 10usize;
    let school_effect = [-6.0, -2.0, 1.0, 3.0, -1.0, 5.0]; // between-group spread
    let n = n_groups * 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, &effect) in school_effect.iter().enumerate() {
        for i in 0..per {
            let hours = 1.0 + (i as f64) * 0.4; // 1.0 .. 4.6
            let noise = ((i * 3 + j) % 5) as f64 - 2.0; // small within-school wiggle
            x[(r, 1)] = hours;
            y[r] = 50.0 + 4.0 * hours + effect + noise;
            groups[r] = j;
            r += 1;
        }
    }

    let lmm = LinearMixedModel::new(x.clone(), y.clone(), &groups).unwrap();
    let ols = OlsFit::new(x, y).unwrap();

    println!("Random-intercept linear mixed model (REML)\n");
    let se = lmm.coefficient_standard_errors();
    let ols_se = ols.coefficient_standard_errors();
    println!("Fixed effects:");
    println!("{:<10}{:>10}{:>12}{:>12}", "", "estimate", "SE (LMM)", "SE (OLS)");
    for (j, name) in ["intercept", "hours"].iter().enumerate() {
        println!(
            "{:<10}{:>10.3}{:>12.3}{:>12.3}",
            name,
            lmm.coefficients()[j],
            se[j],
            ols_se[j]
        );
    }

    println!("\nVariance components:");
    println!("  between-group σ²_b = {:.3}", lmm.group_variance());
    println!("  residual     σ²_e = {:.3}", lmm.residual_variance());
    println!(
        "  intraclass correlation ICC = {:.3}  ({:.0}% of variance is between schools)",
        lmm.icc(),
        100.0 * lmm.icc()
    );
    println!("  log-likelihood (REML) = {:.3},  AIC = {:.2}", lmm.log_likelihood(), lmm.aic());

    println!("\nBLUPs (predicted school intercepts, shrunk toward 0):");
    let b = lmm.random_effects();
    for j in 0..lmm.n_groups() {
        println!("  school {j}:  {:>6.3}   (true effect {:>5.1})", b[j], school_effect[j]);
    }
    println!(
        "\n  Note: ignoring the grouping, OLS reports a smaller SE for the intercept\n  \
         ({:.3} vs {:.3}) — it double-counts correlated within-school observations.",
        ols_se[0], se[0]
    );

    // --- A generalized mixed model (Poisson counts) -----------------------
    // Same grouping idea, but the response is a count with a per-school random
    // intercept — fit by Laplace approximation.
    let mut xc = Array2::<f64>::ones((n, 2));
    let mut counts = Array1::<f64>::zeros(n);
    let mut grp = vec![0usize; n];
    let mut r = 0;
    for (j, &effect) in school_effect.iter().enumerate() {
        for i in 0..per {
            let cov = (i as f64) * 0.2 - 1.0;
            xc[(r, 1)] = cov;
            let mu = (1.2 + 0.3 * cov + 0.1 * effect).exp();
            counts[r] = mu.round().max(0.0);
            grp[r] = j;
            r += 1;
        }
    }
    let glmm = GlmmFit::new(xc, counts, &grp, GlmmFamily::Poisson).unwrap();
    println!("\nPoisson GLMM (random school intercept, Laplace):");
    println!(
        "  covariate coef = {:.3} (se {:.3}),  σ_b = {:.3},  AIC = {:.2}",
        glmm.coefficients()[1],
        glmm.coefficient_standard_errors()[1],
        glmm.sigma_b(),
        glmm.aic()
    );
}