solow-gee 0.2.0

//! Cross-validation of the nominal / ordinal categorical GEE estimators
//! against golden reference values frozen in `tests/fixtures/gee_ext.json`
//! (generated by `tools/reference/gen_gee_ext.py`).

use ndarray::{Array1, Array2};
use serde_json::Value;
use solow_gee::{CategoricalCov, CategoricalGeeResults, NominalGee, OrdinalGee};
use std::fs;

fn load() -> Value {
    let p = concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/../../tests/fixtures/gee_ext.json"
    );
    let s = fs::read_to_string(p).expect("fixture present (run tools/reference/gen_gee_ext.py)");
    serde_json::from_str(&s).unwrap()
}

fn mat(v: &Value) -> Array2<f64> {
    let rows: Vec<Vec<f64>> = v
        .as_array()
        .unwrap()
        .iter()
        .map(|r| {
            r.as_array()
                .unwrap()
                .iter()
                .map(|x| x.as_f64().unwrap())
                .collect()
        })
        .collect();
    let (m, n) = (rows.len(), rows[0].len());
    Array2::from_shape_vec((m, n), rows.into_iter().flatten().collect()).unwrap()
}

fn vec1(v: &Value) -> Array1<f64> {
    Array1::from_vec(
        v.as_array()
            .unwrap()
            .iter()
            .map(|x| x.as_f64().unwrap())
            .collect(),
    )
}

fn rel(got: f64, want: f64) -> f64 {
    (got - want).abs() / (1.0 + want.abs())
}

fn check_vec(label: &str, got: &Array1<f64>, exp: &Value, key: &str, tol: f64) {
    let want = vec1(&exp[key]);
    assert_eq!(
        got.len(),
        want.len(),
        "{label}.{key}: length {} != {}",
        got.len(),
        want.len()
    );
    for i in 0..got.len() {
        let e = rel(got[i], want[i]);
        assert!(
            e <= tol,
            "{label}.{key}[{i}]: rel-err {e:.3e} (got {}, want {})",
            got[i],
            want[i]
        );
    }
}

fn check_scalar(label: &str, got: f64, exp: &Value, key: &str, tol: f64) {
    let want = exp[key].as_f64().unwrap();
    let e = rel(got, want);
    assert!(
        e <= tol,
        "{label}.{key}: rel-err {e:.3e} (got {got}, want {want})"
    );
}

fn check_mat(label: &str, got: &Array2<f64>, exp: &Value, key: &str, tol: f64) {
    let want = mat(&exp[key]);
    assert_eq!(got.dim(), want.dim(), "{label}.{key}: shape");
    for i in 0..got.nrows() {
        for j in 0..got.ncols() {
            let e = rel(got[[i, j]], want[[i, j]]);
            assert!(
                e <= tol,
                "{label}.{key}[{i}][{j}]: rel-err {e:.3e} (got {}, want {})",
                got[[i, j]],
                want[[i, j]]
            );
        }
    }
}

fn cov_for(name: &str) -> CategoricalCov {
    match name {
        "independence" => CategoricalCov::Independence,
        "global_odds_ratio" => CategoricalCov::GlobalOddsRatio,
        other => panic!("unknown cov_struct {other}"),
    }
}

fn verify(label: &str, res: &CategoricalGeeResults, exp: &Value) {
    // Mean parameters and the working association solve the score equations
    // essentially exactly; match the reference to MLE-grade precision.
    check_vec(label, &res.params, exp, "params", 1e-7);
    check_vec(label, &res.bse, exp, "bse", 1e-6);
    check_vec(label, &res.tvalues, exp, "tvalues", 1e-6);
    // p-values pass through the normal survival function.
    check_vec(label, &res.pvalues, exp, "pvalues", 1e-6);
    check_vec(label, &res.fittedvalues, exp, "fittedvalues", 1e-6);

    check_mat(label, &res.cov_robust, exp, "cov_robust", 1e-6);
    check_mat(label, &res.cov_naive, exp, "cov_naive", 1e-6);

    check_scalar(label, res.dep_params, exp, "dep_params", 1e-6);
    check_scalar(label, res.scale, exp, "scale", 1e-10);
}

#[test]
fn categorical_gee_matches_reference() {
    let fx = load();
    for c in fx["cases"].as_array().unwrap() {
        let name = c["name"].as_str().unwrap();
        let kind = c["kind"].as_str().unwrap();
        let cov = cov_for(c["cov_struct"].as_str().unwrap());
        let y = vec1(&c["endog"]);
        let x = mat(&c["exog"]);
        let groups: Vec<i64> = c["groups"]
            .as_array()
            .unwrap()
            .iter()
            .map(|g| g.as_i64().unwrap())
            .collect();

        let res = match kind {
            "nominal" => NominalGee::new(y, x, &groups, cov)
                .unwrap()
                .ctol(1e-10)
                .maxiter(300)
                .fit()
                .unwrap(),
            "ordinal" => OrdinalGee::new(y, x, &groups, cov)
                .unwrap()
                .ctol(1e-10)
                .maxiter(300)
                .fit()
                .unwrap(),
            other => panic!("unknown kind {other}"),
        };
        assert!(
            res.converged,
            "{name}: did not converge (score_norm too large)"
        );
        verify(name, &res, &c["expected"]);
    }
}