regression-diagnostics 0.2.0

use crate::OlsFit;

/// Total sum of squares used as the `R²` denominator.
///
/// With an intercept this is the **centered** TSS `Σ (yᵢ − ȳ)²`; without one it
/// is the **uncentered** TSS `Σ yᵢ²`. The distinction matters: reporting a
/// centered `R²` for a through-the-origin model would overstate fit, so the
/// choice follows whether the model actually estimated a constant.
pub fn total_sum_of_squares(fit: &OlsFit) -> f64 {
    let y = fit.response();
    if fit.has_intercept() {
        let mean = y.sum() / y.len() as f64;
        y.iter().map(|v| (v - mean).powi(2)).sum()
    } else {
        y.iter().map(|v| v * v).sum()
    }
}

/// Coefficient of determination `R² = 1 − RSS / TSS`.
///
/// Uses centered TSS when the model has an intercept, uncentered otherwise (see
/// [`total_sum_of_squares`]).
pub fn r_squared(fit: &OlsFit) -> f64 {
    let tss = total_sum_of_squares(fit);
    if tss <= 0.0 {
        return f64::NAN;
    }
    1.0 - fit.residual_sum_of_squares() / tss
}

/// Adjusted `R²`, penalizing model complexity:
///
/// `R̄² = 1 − (RSS / dfₑ) / (TSS / df_total)`
///
/// where `dfₑ = n − p` and `df_total = n − 1` with an intercept (`n` without).
///
/// Adding an irrelevant predictor can only nudge plain `R²` upward, but it costs
/// a residual degree of freedom, so adjusted `R²` frequently *falls* — that
/// asymmetry is exactly what makes it the honest "did this predictor earn its
/// place" statistic.
pub fn adjusted_r_squared(fit: &OlsFit) -> f64 {
    let tss = total_sum_of_squares(fit);
    if tss <= 0.0 {
        return f64::NAN;
    }
    let df_total = if fit.has_intercept() {
        (fit.n_observations() - 1) as f64
    } else {
        fit.n_observations() as f64
    };
    let df_resid = fit.df_residual();
    if df_resid <= 0.0 || df_total <= 0.0 {
        return f64::NAN;
    }
    1.0 - (fit.residual_sum_of_squares() / df_resid) / (tss / df_total)
}