regression-diagnostics 0.2.0

use ndarray::{Array1, Array2};

use super::CoxFit;

/// **Martingale residuals** `Mᵢ = δᵢ − Ĥᵢ`, the event indicator minus the
/// estimated cumulative hazard `Ĥᵢ = exp(ηᵢ)·Ĥ₀(tᵢ)`.
///
/// They measure the difference between the observed number of events for subject
/// `i` (0 or 1) and the number the model expected. At the MLE they **sum to
/// zero**. Martingale residuals are highly skewed (bounded above by 1, unbounded
/// below), so they are best for assessing functional form of a covariate rather
/// than for symmetric outlier detection — for that, prefer
/// [`deviance_residuals`].
pub fn martingale_residuals(fit: &CoxFit) -> Array1<f64> {
    let event = fit.event();
    Array1::from_shape_fn(fit.n_observations(), |i| {
        event[i] - fit.cumulative_hazard_at(i)
    })
}

/// **Deviance residuals** — a symmetrizing transform of the martingale
/// residuals,
///
/// `dᵢ = sign(Mᵢ)·√(−2[ Mᵢ + δᵢ ln(δᵢ − Mᵢ) ])`.
///
/// Roughly mean-zero and symmetric, so large `|dᵢ|` flags poorly-fit subjects
/// (positive = died sooner than expected, negative = survived longer). The
/// `δᵢ − Mᵢ` term is exactly the estimated cumulative hazard `Ĥᵢ`.
pub fn deviance_residuals(fit: &CoxFit) -> Array1<f64> {
    let m = martingale_residuals(fit);
    let event = fit.event();
    Array1::from_shape_fn(fit.n_observations(), |i| {
        let mi = m[i];
        let di = event[i];
        let hi = di - mi; // estimated cumulative hazard (>= 0)
        let inner = if di > 0.0 && hi > 0.0 {
            mi + di * hi.ln()
        } else {
            mi
        };
        let sign = if mi >= 0.0 { 1.0 } else { -1.0 };
        sign * (-2.0 * inner).max(0.0).sqrt()
    })
}

/// **Schoenfeld residuals**, defined only at event times: for the subject with
/// an event at `tᵢ`, `rᵢₐ = xᵢₐ − x̄ₐ(tᵢ)`, where `x̄ₐ(t)` is the risk-weighted
/// mean of covariate `a` over the risk set `R(t)`.
///
/// Returns `(event_times, residuals)` with `residuals` of shape
/// `n_events × p`, one row per event (ordered by time). Schoenfeld residuals are
/// the building block for testing the **proportional-hazards assumption** — a
/// trend against time in column `a` signals a time-varying effect for predictor
/// `a`. For each covariate they **sum to zero** at the MLE (the score
/// equations). Tied events share the same risk-set mean.
pub fn schoenfeld_residuals(fit: &CoxFit) -> (Vec<f64>, Array2<f64>) {
    let x = fit.design_matrix();
    let beta = fit.coef_slice();
    let time = fit.time();
    let event = fit.event();
    let n = fit.n_observations();
    let p = fit.n_parameters();

    let w: Vec<f64> = (0..n)
        .map(|i| (0..p).map(|j| x[(i, j)] * beta[j]).sum::<f64>().exp())
        .collect();

    // Events ordered by time.
    let mut ev: Vec<usize> = (0..n).filter(|&i| event[i] == 1.0).collect();
    ev.sort_by(|&a, &b| time[a].partial_cmp(&time[b]).unwrap());

    let mut times = Vec::with_capacity(ev.len());
    let mut resid = Array2::<f64>::zeros((ev.len(), p));
    for (row, &i) in ev.iter().enumerate() {
        let t = time[i];
        let s = fit.stratum_of(i);
        // Risk-weighted covariate mean over the (stratum-restricted) risk set R(t).
        let mut denom = 0.0;
        let mut num = vec![0.0; p];
        for jj in 0..n {
            if fit.at_risk(jj, t, s) {
                denom += w[jj];
                for a in 0..p {
                    num[a] += w[jj] * x[(jj, a)];
                }
            }
        }
        times.push(t);
        for a in 0..p {
            let mean_a = if denom > 0.0 { num[a] / denom } else { 0.0 };
            resid[(row, a)] = x[(i, a)] - mean_a;
        }
    }
    (times, resid)
}