regression-diagnostics 0.2.0

//! [`OlsFit`] — the fitted-model type every diagnostic in this crate operates
//! on.

use ndarray::{Array1, Array2, ArrayView1, ArrayView2};

use crate::error::{RegressionError, Result};
use crate::linalg::{self, dmatrix_from_rows, dvector_from_slice};

/// A fitted ordinary-least-squares model: the shared representation every
/// diagnostic in this crate is computed from.
///
/// # Intercept convention
///
/// **The caller owns the design matrix `X`, including any intercept column.**
/// This crate does *not* silently prepend a column of ones. There are two
/// constructors, and the choice is explicit at the call site:
///
/// * [`OlsFit::new`] fits exactly the columns you pass. If one of them is
///   constant it is auto-detected and treated as the intercept (this drives the
///   centered-vs-uncentered `R²` choice, and excludes that column from VIF); if
///   none is constant the model is fit **through the origin**.
/// * [`OlsFit::with_intercept`] prepends a ones column for you and marks it as
///   the intercept.
///
/// Getting this wrong silently corrupts every downstream diagnostic, so the
/// convention is stated here rather than buried in the implementation.
///
/// # What is computed and cached
///
/// Construction performs a single QR factorization of `X` and caches the
/// coefficients, fitted values, residuals, leverage vector `diag(H)`, `(XᵀX)⁻¹`,
/// the design's singular values, and the residual variance. Diagnostics read
/// these cached quantities rather than refactorizing.
///
/// Coefficients are obtained by a QR **solve of `X`**, never by inverting `XᵀX`
/// — that matters for the multicollinearity diagnostics specifically, since
/// forming `XᵀX` squares the condition number they exist to measure. Leverage is
/// read from the thin `Q` factor, so the full `n × n` hat matrix is never formed.
///
/// # Example
///
/// ```
/// use ndarray::{array, Array2};
/// use regression_diagnostics::OlsFit;
///
/// // y = 1 + 2*x exactly; supply the intercept ourselves.
/// let x = array![[1.0, 0.0], [1.0, 1.0], [1.0, 2.0], [1.0, 3.0]];
/// let y = array![1.0, 3.0, 5.0, 7.0];
/// let fit = OlsFit::new(x, y).unwrap();
/// assert!((fit.coefficients()[0] - 1.0).abs() < 1e-9);
/// assert!((fit.coefficients()[1] - 2.0).abs() < 1e-9);
/// ```
#[derive(Debug, Clone)]
pub struct OlsFit {
    x: Array2<f64>,
    y: Array1<f64>,
    coefficients: Array1<f64>,
    fitted: Array1<f64>,
    residuals: Array1<f64>,
    leverage: Array1<f64>,
    xtx_inv: Array2<f64>,
    singular_values: Vec<f64>,
    intercept_col: Option<usize>,
    n: usize,
    p: usize,
    rss: f64,
    /// Residual variance estimate `RSS / (n - p)`.
    sigma2: f64,
}

impl OlsFit {
    /// Fit `y ~ X` by OLS, using the columns of `X` exactly as given.
    ///
    /// A constant column, if present, is auto-detected as the intercept. See the
    /// [type-level docs](OlsFit#intercept-convention) for the full convention.
    ///
    /// # Errors
    ///
    /// * [`RegressionError::EmptyInput`] if `X` or `y` is empty.
    /// * [`RegressionError::ShapeMismatch`] if `X.nrows() != y.len()`.
    /// * [`RegressionError::NoResidualDegreesOfFreedom`] if `n <= p`.
    /// * [`RegressionError::RankDeficient`] if the columns are collinear.
    pub fn new(x: Array2<f64>, y: Array1<f64>) -> Result<Self> {
        let intercept_col = detect_constant_column(&x);
        Self::build(x, y, intercept_col)
    }

    /// Fit `y ~ [1 | X]` by OLS, prepending a column of ones as the intercept.
    ///
    /// Errors are the same as [`OlsFit::new`], evaluated against the augmented
    /// design (so a design with `n == p_original + 1` will fail the
    /// residual-degrees-of-freedom check).
    pub fn with_intercept(x: Array2<f64>, y: Array1<f64>) -> Result<Self> {
        if x.nrows() == 0 {
            return Err(RegressionError::EmptyInput { what: "X" });
        }
        let n = x.nrows();
        let mut augmented = Array2::<f64>::ones((n, x.ncols() + 1));
        for j in 0..x.ncols() {
            augmented.column_mut(j + 1).assign(&x.column(j));
        }
        Self::build(augmented, y, Some(0))
    }

    fn build(x: Array2<f64>, y: Array1<f64>, intercept_col: Option<usize>) -> Result<Self> {
        let n = x.nrows();
        let p = x.ncols();
        if n == 0 || p == 0 {
            return Err(RegressionError::EmptyInput { what: "X" });
        }
        if y.is_empty() {
            return Err(RegressionError::EmptyInput { what: "y" });
        }
        if y.len() != n {
            return Err(RegressionError::ShapeMismatch {
                what: "y length vs X rows",
                expected: n,
                got: y.len(),
            });
        }
        if n <= p {
            return Err(RegressionError::NoResidualDegreesOfFreedom {
                n,
                p,
                df: n as isize - p as isize,
            });
        }

        let x_dm = dmatrix_from_rows(
            n,
            p,
            x.as_standard_layout().as_slice().expect("standard layout"),
        );
        let y_dv = dvector_from_slice(y.as_standard_layout().as_slice().expect("standard layout"));

        let qr = linalg::ols_via_qr(&x_dm, &y_dv)?;

        let coefficients = Array1::from_iter(qr.coef.iter().copied());
        let fitted = Array1::from_iter(qr.fitted.iter().copied());
        let residuals = &y - &fitted;
        let leverage = Array1::from_iter(qr.leverage.iter().copied());
        let xtx_inv = Array2::from_shape_fn((p, p), |(i, j)| qr.xtx_inv[(i, j)]);

        let rss: f64 = residuals.iter().map(|r| r * r).sum();
        let sigma2 = rss / (n - p) as f64;

        Ok(Self {
            x,
            y,
            coefficients,
            fitted,
            residuals,
            leverage,
            xtx_inv,
            singular_values: qr.singular_values,
            intercept_col,
            n,
            p,
            rss,
            sigma2,
        })
    }

    // ---- basic accessors --------------------------------------------------

    /// Number of observations `n`.
    pub fn n_observations(&self) -> usize {
        self.n
    }

    /// Number of model parameters `p` (design-matrix columns, intercept
    /// included if present).
    pub fn n_parameters(&self) -> usize {
        self.p
    }

    /// Whether the model includes an intercept (constant) column.
    pub fn has_intercept(&self) -> bool {
        self.intercept_col.is_some()
    }

    /// Index of the intercept column within the design matrix, if any.
    pub fn intercept_column(&self) -> Option<usize> {
        self.intercept_col
    }

    /// The design matrix `X` as fitted (with the intercept column if one was
    /// added or detected).
    pub fn design_matrix(&self) -> ArrayView2<'_, f64> {
        self.x.view()
    }

    /// The response vector `y`.
    pub fn response(&self) -> ArrayView1<'_, f64> {
        self.y.view()
    }

    /// Estimated coefficients, one per design-matrix column.
    pub fn coefficients(&self) -> ArrayView1<'_, f64> {
        self.coefficients.view()
    }

    /// Fitted values `ŷ = X β`.
    pub fn fitted_values(&self) -> ArrayView1<'_, f64> {
        self.fitted.view()
    }

    /// Raw residuals `e = y − ŷ`.
    pub fn residuals(&self) -> ArrayView1<'_, f64> {
        self.residuals.view()
    }

    /// Leverage vector `diag(H)`. Each entry lies in `[0, 1]` and the vector
    /// sums to `p` (the number of parameters) — a standard identity worth
    /// checking as a correctness probe.
    pub fn leverage(&self) -> ArrayView1<'_, f64> {
        self.leverage.view()
    }

    /// Residual sum of squares `Σ eᵢ²`.
    pub fn residual_sum_of_squares(&self) -> f64 {
        self.rss
    }

    /// Residual degrees of freedom `n − p`.
    pub fn df_residual(&self) -> f64 {
        (self.n - self.p) as f64
    }

    /// Model degrees of freedom: `p − 1` with an intercept, `p` without.
    pub fn df_model(&self) -> f64 {
        if self.has_intercept() {
            (self.p - 1) as f64
        } else {
            self.p as f64
        }
    }

    /// Unbiased residual variance estimate `s² = RSS / (n − p)`.
    pub fn residual_variance(&self) -> f64 {
        self.sigma2
    }

    /// Residual standard error `s = √(RSS / (n − p))`.
    pub fn residual_standard_error(&self) -> f64 {
        self.sigma2.sqrt()
    }

    /// Standard errors of the coefficients: `sⱼ = s · √((XᵀX)⁻¹ⱼⱼ)`.
    pub fn coefficient_standard_errors(&self) -> Array1<f64> {
        Array1::from_shape_fn(self.p, |j| (self.sigma2 * self.xtx_inv[(j, j)]).sqrt())
    }

    /// Singular values of the design matrix, in descending order.
    pub fn singular_values(&self) -> &[f64] {
        &self.singular_values
    }

    // ---- crate-internal helpers ------------------------------------------

    /// `R²` from regressing design column `j` on all other columns — the
    /// auxiliary regression VIF is built on. Returns `None` if `j` is the
    /// intercept column (VIF is undefined there).
    pub(crate) fn column_on_others_r2(&self, j: usize) -> Option<f64> {
        if self.intercept_col == Some(j) {
            return None;
        }
        let others: Vec<usize> = (0..self.p).filter(|&c| c != j).collect();
        let sub = self.x.select(ndarray::Axis(1), &others);
        let target = self.x.column(j).to_owned();

        let sub_dm = dmatrix_from_rows(
            self.n,
            others.len(),
            sub.as_standard_layout()
                .as_slice()
                .expect("standard layout"),
        );
        let target_dv = dvector_from_slice(target.as_slice().expect("contiguous"));
        // aux_r_squared returns Some(1.0) on a collinear auxiliary design →
        // caller maps that to an infinite VIF.
        linalg::aux_r_squared(&sub_dm, &target_dv)
    }
}

/// Detect the first constant column of `x` (all entries equal within a relative
/// tolerance), which is treated as an intercept.
fn detect_constant_column(x: &Array2<f64>) -> Option<usize> {
    for (j, col) in x.columns().into_iter().enumerate() {
        let first = col[0];
        let scale = first.abs().max(1.0);
        if col.iter().all(|&v| (v - first).abs() <= 1e-12 * scale) {
            return Some(j);
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use super::*;
    use ndarray::array;

    /// Hand-verified closed-form OLS on a tiny fixed dataset.
    ///
    /// With intercept, x = [0,1,2,3], y = [1,3,5,7]: the exact least-squares line
    /// is y = 1 + 2x (a perfect fit), so β = (1, 2), residuals are all zero.
    #[test]
    fn coefficients_match_closed_form() {
        let x = array![[1.0, 0.0], [1.0, 1.0], [1.0, 2.0], [1.0, 3.0]];
        let y = array![1.0, 3.0, 5.0, 7.0];
        let fit = OlsFit::new(x, y).unwrap();
        assert!((fit.coefficients()[0] - 1.0).abs() < 1e-9);
        assert!((fit.coefficients()[1] - 2.0).abs() < 1e-9);
        assert!(fit.residuals().iter().all(|&e| e.abs() < 1e-9));
    }

    /// Non-perfect fit, closed form checked by hand.
    ///
    /// x = [1,2,3,4,5], y = [1,2,1.3,3.75,2.25]. Standard OLS gives
    /// slope ≈ 0.425, intercept ≈ 0.785 (classic worked example).
    #[test]
    fn coefficients_match_worked_example() {
        let x = array![[1.0, 1.0], [1.0, 2.0], [1.0, 3.0], [1.0, 4.0], [1.0, 5.0]];
        let y = array![1.0, 2.0, 1.3, 3.75, 2.25];
        let fit = OlsFit::new(x, y).unwrap();
        assert!((fit.coefficients()[0] - 0.785).abs() < 1e-3);
        assert!((fit.coefficients()[1] - 0.425).abs() < 1e-3);
    }

    /// Leverage must sum to the number of parameters — a strong correctness
    /// probe independent of the coefficient values.
    #[test]
    fn leverage_sums_to_p() {
        let x = array![
            [1.0, 0.0, 2.0],
            [1.0, 1.0, 1.0],
            [1.0, 2.0, 4.0],
            [1.0, 3.0, 1.0],
            [1.0, 4.0, 5.0],
            [1.0, 5.0, 2.0],
        ];
        let y = array![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
        let fit = OlsFit::new(x, y).unwrap();
        let total: f64 = fit.leverage().sum();
        assert!((total - 3.0).abs() < 1e-9, "leverage sum = {total}");
        // Every leverage lies in [0, 1].
        assert!(fit
            .leverage()
            .iter()
            .all(|&h| (0.0..=1.0 + 1e-9).contains(&h)));
    }

    /// `with_intercept` prepends a ones column and detects it as the intercept.
    #[test]
    fn with_intercept_prepends_ones() {
        let x = array![[0.0], [1.0], [2.0], [3.0]];
        let y = array![1.0, 3.0, 5.0, 7.0];
        let fit = OlsFit::with_intercept(x, y).unwrap();
        assert_eq!(fit.n_parameters(), 2);
        assert_eq!(fit.intercept_column(), Some(0));
        assert!((fit.coefficients()[0] - 1.0).abs() < 1e-9);
        assert!((fit.coefficients()[1] - 2.0).abs() < 1e-9);
    }
}