plotters-statistical 0.1.0

//! Pearson and Spearman correlation, and correlation matrices.

use super::StatsError;

/// Which correlation coefficient to compute.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CorrelationMethod {
    /// Pearson product-moment correlation (linear).
    Pearson,
    /// Spearman rank correlation (monotonic) — Pearson on the ranks.
    Spearman,
}

/// Pearson correlation of two equal-length samples.
///
/// # Errors
/// * [`StatsError::LengthMismatch`] if the inputs differ in length.
/// * [`StatsError::EmptyInput`] if empty.
/// * [`StatsError::ZeroVariance`] if either sample is constant (correlation
///   undefined rather than a silent `NaN`).
pub fn pearson(x: &[f64], y: &[f64]) -> Result<f64, StatsError> {
    if x.len() != y.len() {
        return Err(StatsError::LengthMismatch {
            scores: x.len(),
            labels: y.len(),
        });
    }
    if x.is_empty() {
        return Err(StatsError::EmptyInput);
    }
    let n = x.len() as f64;
    let mx = x.iter().sum::<f64>() / n;
    let my = y.iter().sum::<f64>() / n;
    let mut sxy = 0.0;
    let mut sxx = 0.0;
    let mut syy = 0.0;
    for (&xi, &yi) in x.iter().zip(y) {
        let dx = xi - mx;
        let dy = yi - my;
        sxy += dx * dy;
        sxx += dx * dx;
        syy += dy * dy;
    }
    if sxx <= 0.0 || syy <= 0.0 {
        return Err(StatsError::ZeroVariance);
    }
    Ok((sxy / (sxx.sqrt() * syy.sqrt())).clamp(-1.0, 1.0))
}

/// Fractional ranks of `data` (average ranks for ties), 1-based. Used by
/// Spearman correlation.
pub fn rank(data: &[f64]) -> Vec<f64> {
    let n = data.len();
    let mut idx: Vec<usize> = (0..n).collect();
    idx.sort_by(|&a, &b| {
        data[a]
            .partial_cmp(&data[b])
            .unwrap_or(std::cmp::Ordering::Equal)
    });
    let mut ranks = vec![0.0; n];
    let mut i = 0;
    while i < n {
        let mut j = i + 1;
        // Extend over a run of equal values.
        while j < n && data[idx[j]] == data[idx[i]] {
            j += 1;
        }
        // Average of the 1-based ranks i+1..=j for the tied group.
        let avg = ((i + 1 + j) as f64) / 2.0;
        for &k in &idx[i..j] {
            ranks[k] = avg;
        }
        i = j;
    }
    ranks
}

/// Spearman rank correlation — Pearson correlation of the ranks.
pub fn spearman(x: &[f64], y: &[f64]) -> Result<f64, StatsError> {
    if x.len() != y.len() {
        return Err(StatsError::LengthMismatch {
            scores: x.len(),
            labels: y.len(),
        });
    }
    if x.is_empty() {
        return Err(StatsError::EmptyInput);
    }
    pearson(&rank(x), &rank(y))
}

/// Correlation matrix of `columns` (each inner slice is one variable's values,
/// all the same length). The diagonal is exactly `1.0`; any pair with a
/// constant column yields `f64::NAN` in that cell (so a chart can render it as a
/// distinct "undefined" color) rather than erroring out the whole matrix.
///
/// # Errors
/// * [`StatsError::EmptyInput`] if there are no columns.
/// * [`StatsError::LengthMismatch`] if the columns are not all the same length.
pub fn correlation_matrix(
    columns: &[Vec<f64>],
    method: CorrelationMethod,
) -> Result<Vec<Vec<f64>>, StatsError> {
    if columns.is_empty() {
        return Err(StatsError::EmptyInput);
    }
    let len = columns[0].len();
    if columns.iter().any(|c| c.len() != len) {
        return Err(StatsError::LengthMismatch {
            scores: len,
            labels: columns
                .iter()
                .map(|c| c.len())
                .find(|&l| l != len)
                .unwrap_or(len),
        });
    }
    let n = columns.len();
    let corr = |a: &[f64], b: &[f64]| match method {
        CorrelationMethod::Pearson => pearson(a, b),
        CorrelationMethod::Spearman => spearman(a, b),
    };
    let mut m = vec![vec![0.0; n]; n];
    for i in 0..n {
        m[i][i] = 1.0;
        for j in (i + 1)..n {
            let v = corr(&columns[i], &columns[j]).unwrap_or(f64::NAN);
            m[i][j] = v;
            m[j][i] = v;
        }
    }
    Ok(m)
}