plotters_statistical/stats/
correlation.rs

1//! Pearson and Spearman correlation, and correlation matrices.
2
3use super::StatsError;
4
5/// Which correlation coefficient to compute.
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum CorrelationMethod {
8    /// Pearson product-moment correlation (linear).
9    Pearson,
10    /// Spearman rank correlation (monotonic) — Pearson on the ranks.
11    Spearman,
12}
13
14/// Pearson correlation of two equal-length samples.
15///
16/// # Errors
17/// * [`StatsError::LengthMismatch`] if the inputs differ in length.
18/// * [`StatsError::EmptyInput`] if empty.
19/// * [`StatsError::ZeroVariance`] if either sample is constant (correlation
20///   undefined rather than a silent `NaN`).
21pub fn pearson(x: &[f64], y: &[f64]) -> Result<f64, StatsError> {
22    if x.len() != y.len() {
23        return Err(StatsError::LengthMismatch {
24            scores: x.len(),
25            labels: y.len(),
26        });
27    }
28    if x.is_empty() {
29        return Err(StatsError::EmptyInput);
30    }
31    let n = x.len() as f64;
32    let mx = x.iter().sum::<f64>() / n;
33    let my = y.iter().sum::<f64>() / n;
34    let mut sxy = 0.0;
35    let mut sxx = 0.0;
36    let mut syy = 0.0;
37    for (&xi, &yi) in x.iter().zip(y) {
38        let dx = xi - mx;
39        let dy = yi - my;
40        sxy += dx * dy;
41        sxx += dx * dx;
42        syy += dy * dy;
43    }
44    if sxx <= 0.0 || syy <= 0.0 {
45        return Err(StatsError::ZeroVariance);
46    }
47    Ok((sxy / (sxx.sqrt() * syy.sqrt())).clamp(-1.0, 1.0))
48}
49
50/// Fractional ranks of `data` (average ranks for ties), 1-based. Used by
51/// Spearman correlation.
52pub fn rank(data: &[f64]) -> Vec<f64> {
53    let n = data.len();
54    let mut idx: Vec<usize> = (0..n).collect();
55    idx.sort_by(|&a, &b| {
56        data[a]
57            .partial_cmp(&data[b])
58            .unwrap_or(std::cmp::Ordering::Equal)
59    });
60    let mut ranks = vec![0.0; n];
61    let mut i = 0;
62    while i < n {
63        let mut j = i + 1;
64        // Extend over a run of equal values.
65        while j < n && data[idx[j]] == data[idx[i]] {
66            j += 1;
67        }
68        // Average of the 1-based ranks i+1..=j for the tied group.
69        let avg = ((i + 1 + j) as f64) / 2.0;
70        for &k in &idx[i..j] {
71            ranks[k] = avg;
72        }
73        i = j;
74    }
75    ranks
76}
77
78/// Spearman rank correlation — Pearson correlation of the ranks.
79pub fn spearman(x: &[f64], y: &[f64]) -> Result<f64, StatsError> {
80    if x.len() != y.len() {
81        return Err(StatsError::LengthMismatch {
82            scores: x.len(),
83            labels: y.len(),
84        });
85    }
86    if x.is_empty() {
87        return Err(StatsError::EmptyInput);
88    }
89    pearson(&rank(x), &rank(y))
90}
91
92/// Correlation matrix of `columns` (each inner slice is one variable's values,
93/// all the same length). The diagonal is exactly `1.0`; any pair with a
94/// constant column yields `f64::NAN` in that cell (so a chart can render it as a
95/// distinct "undefined" color) rather than erroring out the whole matrix.
96///
97/// # Errors
98/// * [`StatsError::EmptyInput`] if there are no columns.
99/// * [`StatsError::LengthMismatch`] if the columns are not all the same length.
100pub fn correlation_matrix(
101    columns: &[Vec<f64>],
102    method: CorrelationMethod,
103) -> Result<Vec<Vec<f64>>, StatsError> {
104    if columns.is_empty() {
105        return Err(StatsError::EmptyInput);
106    }
107    let len = columns[0].len();
108    if columns.iter().any(|c| c.len() != len) {
109        return Err(StatsError::LengthMismatch {
110            scores: len,
111            labels: columns
112                .iter()
113                .map(|c| c.len())
114                .find(|&l| l != len)
115                .unwrap_or(len),
116        });
117    }
118    let n = columns.len();
119    let corr = |a: &[f64], b: &[f64]| match method {
120        CorrelationMethod::Pearson => pearson(a, b),
121        CorrelationMethod::Spearman => spearman(a, b),
122    };
123    let mut m = vec![vec![0.0; n]; n];
124    for i in 0..n {
125        m[i][i] = 1.0;
126        for j in (i + 1)..n {
127            let v = corr(&columns[i], &columns[j]).unwrap_or(f64::NAN);
128            m[i][j] = v;
129            m[j][i] = v;
130        }
131    }
132    Ok(m)
133}
plotters_statistical/stats/correlation.rs

plotters_statistical/stats/
correlation.rs
