regression-diagnostics 0.2.0

//! Survival analysis on the classic Freireich (1963) leukemia remission data:
//! a Kaplan–Meier curve per treatment arm, a Cox proportional-hazards fit, and
//! the proportional-hazards check via Schoenfeld residuals.
//!
//! Run with: `cargo run --example survival_diagnostics`

use ndarray::{Array1, Array2};
use regression_diagnostics::survival::{
    martingale_residuals, schoenfeld_residuals, AftDistribution, AftFit, CoxFit, KaplanMeier,
};

fn main() {
    // Remission times (weeks); placebo group has all events, 6-MP group is
    // heavily censored. x = 1 marks the placebo arm.
    let placebo_t = [
        1.0, 1.0, 2.0, 2.0, 3.0, 4.0, 4.0, 5.0, 5.0, 8.0, 8.0, 8.0, 8.0, 11.0, 11.0, 12.0, 12.0,
        15.0, 17.0, 22.0, 23.0,
    ];
    let mp_t = [
        6.0, 6.0, 6.0, 7.0, 10.0, 13.0, 16.0, 22.0, 23.0, 6.0, 9.0, 10.0, 11.0, 17.0, 19.0, 20.0,
        25.0, 32.0, 32.0, 34.0, 35.0,
    ];
    let mp_e = [
        1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
        0.0, 0.0, 0.0,
    ];

    let n = placebo_t.len() + mp_t.len();
    let mut time = Array1::<f64>::zeros(n);
    let mut event = Array1::<f64>::zeros(n);
    let mut x = Array2::<f64>::zeros((n, 1));
    let mut r = 0;
    for &t in placebo_t.iter() {
        time[r] = t;
        event[r] = 1.0;
        x[(r, 0)] = 1.0;
        r += 1;
    }
    for (k, &t) in mp_t.iter().enumerate() {
        time[r] = t;
        event[r] = mp_e[k];
        x[(r, 0)] = 0.0;
        r += 1;
    }

    // --- Kaplan–Meier per arm ---------------------------------------------
    let km_placebo = KaplanMeier::new(
        Array1::from(placebo_t.to_vec()),
        Array1::from(vec![1.0; placebo_t.len()]),
    )
    .unwrap();
    let km_mp = KaplanMeier::new(Array1::from(mp_t.to_vec()), Array1::from(mp_e.to_vec())).unwrap();
    println!("Kaplan–Meier survival at selected weeks:");
    println!("  week   placebo    6-MP");
    for &w in &[5.0, 10.0, 15.0, 20.0] {
        println!(
            "  {:>4}   {:>7.3}   {:>5.3}",
            w as usize,
            km_placebo.survival_at(w),
            km_mp.survival_at(w)
        );
    }
    println!(
        "  median survival: placebo = {:?}, 6-MP = {:?}",
        km_placebo.median_survival(),
        km_mp.median_survival()
    );

    // --- Cox proportional hazards -----------------------------------------
    let cox = CoxFit::new(time, event, x).unwrap();
    let beta = cox.coefficients()[0];
    let se = cox.coefficient_standard_errors()[0];
    println!("\nCox proportional-hazards (Efron ties, {} iters):", cox.iterations());
    println!(
        "  placebo log-hazard-ratio β = {:.4}  (se {:.4}, p = {:.2e})",
        beta,
        se,
        cox.p_values()[0]
    );
    println!(
        "  hazard ratio exp(β) = {:.3}  → placebo hazard is ~{:.1}× the 6-MP hazard",
        cox.hazard_ratios()[0],
        cox.hazard_ratios()[0]
    );
    println!("  concordance (C-index) = {:.3}", cox.concordance());
    println!("  log partial likelihood = {:.3},  AIC = {:.2}", cox.log_partial_likelihood(), cox.aic());

    // --- Proportional-hazards check ---------------------------------------
    let m = martingale_residuals(&cox);
    println!("\n  Σ martingale residuals = {:.2e}  (≈ 0 at the MLE)", m.sum());
    let (times, sch) = schoenfeld_residuals(&cox);
    // A trend of Schoenfeld residuals against time would signal a time-varying
    // effect (PH violation); here we just report the correlation with time.
    let mean_t = times.iter().sum::<f64>() / times.len() as f64;
    let mean_r = (0..sch.nrows()).map(|i| sch[(i, 0)]).sum::<f64>() / sch.nrows() as f64;
    let mut num = 0.0;
    let mut dt = 0.0;
    let mut dr = 0.0;
    for i in 0..sch.nrows() {
        let a = times[i] - mean_t;
        let b = sch[(i, 0)] - mean_r;
        num += a * b;
        dt += a * a;
        dr += b * b;
    }
    let corr = num / (dt.sqrt() * dr.sqrt());
    println!(
        "  Schoenfeld–time correlation = {:.3}  ({})",
        corr,
        if corr.abs() < 0.3 {
            "no strong evidence against proportional hazards"
        } else {
            "possible PH violation — inspect further"
        }
    );

    // --- Parametric alternative: Weibull AFT ------------------------------
    // Rebuild time/event/x (they were consumed by CoxFit).
    let mut time2 = Array1::<f64>::zeros(n);
    let mut event2 = Array1::<f64>::zeros(n);
    let mut xa = Array2::<f64>::ones((n, 2)); // AFT needs an intercept column
    let mut r2 = 0;
    for &t in placebo_t.iter() {
        time2[r2] = t;
        event2[r2] = 1.0;
        xa[(r2, 1)] = 1.0;
        r2 += 1;
    }
    for (k, &t) in mp_t.iter().enumerate() {
        time2[r2] = t;
        event2[r2] = mp_e[k];
        xa[(r2, 1)] = 0.0;
        r2 += 1;
    }
    let aft = AftFit::new(time2, event2, xa, AftDistribution::Weibull).unwrap();
    println!("\nWeibull AFT (parametric alternative):");
    println!(
        "  placebo coefficient β = {:.3}  (se {:.3})  — negative: placebo shortens survival",
        aft.coefficients()[1],
        aft.coefficient_standard_errors()[1]
    );
    println!("  scale σ = {:.3},  log-likelihood = {:.2},  AIC = {:.2}", aft.scale(), aft.log_likelihood(), aft.aic());
}