alximo_corr/
lib.rs

1use anyhow::{ensure, Result};
2use maths::{corrcoef, cov, squared_sum};
3use ndarray::{Array1, Array2, Axis};
4use statrs::distribution::{Beta, ContinuousCDF};
5
6mod maths;
7
8pub fn pearsonr(x: Array1<f64>, y: Array1<f64>, alternative: Option<&str>) -> Result<(f64, f64)> {
9    ensure!(x.len() == x.len(), "x and y have different lenghts. To calculate the Pearson R coefficient, x and y should have the same length");
10    let n = x.len() as f64;
11
12    let (sum_x, sum_y) = (x.sum(), y.sum());
13
14    let (squared_x_sum, squared_y_sum) = (squared_sum(&x), squared_sum(&y));
15
16    let cross_product_sum = (x * y).sum();
17
18    let numinator = n * cross_product_sum - (sum_x * sum_y);
19    let denominator =
20        ((n * squared_x_sum - sum_x.powi(2)) * (n * squared_y_sum - sum_y.powi(2))).sqrt();
21
22    let mut r = numinator / denominator;
23
24    r = r.min(1.).max(-1.);
25
26    let alternative = alternative.unwrap_or("two-sided");
27
28    let p_val = pearson_significance(r, n, alternative);
29
30    Ok((r, p_val))
31}
32
33fn pearson_significance(r: f64, n: f64, alternative: &str) -> f64 {
34    let ab = (n / 2.) - 1.;
35    let beta = Beta::new(ab, ab).unwrap();
36
37    match alternative {
38        "two-sided" => beta.sf(r.abs()) * 2.,
39        "less" => beta.cdf(r),
40        "greater" => beta.sf(r),
41        _ => panic!(),
42    }
43}
44
45pub fn corrcoef_1d(
46    x: Array1<f64>,
47    y: Option<Array1<f64>>,
48    bias: bool,
49    ddof: Option<usize>,
50    rowvar: bool,
51) -> Result<Array2<f64>> {
52    let c = cov_1d(x, y, rowvar, bias, ddof)?;
53    corrcoef(c)
54}
55
56pub fn corrcoef_2d(
57    x: Array2<f64>,
58    y: Option<Array2<f64>>,
59    bias: bool,
60    ddof: Option<usize>,
61    rowvar: bool,
62) -> Result<Array2<f64>> {
63    let c = cov_2d(x, y, rowvar, bias, ddof)?;
64    corrcoef(c)
65}
66
67pub fn cov_1d(
68    m: Array1<f64>,
69    y: Option<Array1<f64>>,
70    rowvar: bool,
71    bias: bool,
72    ddof: Option<usize>,
73) -> Result<Array2<f64>> {
74    let m = m.insert_axis(Axis(0));
75
76    let y = y.map(|y| y.insert_axis(Axis(0)));
77    cov(m, y, rowvar, bias, ddof)
78}
79
80pub fn cov_2d(
81    m: Array2<f64>,
82    y: Option<Array2<f64>>,
83    rowvar: bool,
84    bias: bool,
85    ddof: Option<usize>,
86) -> Result<Array2<f64>> {
87    cov(m, y, rowvar, bias, ddof)
88}
89
90#[cfg(test)]
91mod tests {
92    use ndarray::{Array, Array2};
93    use ndarray_rand::{rand_distr::Uniform, RandomExt};
94    use statrs::assert_almost_eq;
95
96    use crate::{corrcoef_2d, pearsonr};
97
98    #[test]
99    fn test_pearson_corr() {
100        let x = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
101        let y = Array::from_vec(vec![2.0, 4.0, 6.0, 8.0, 10.0, 12.0]);
102
103        let (r, p) = pearsonr(x, y, None).unwrap();
104        println!("Pearson correlation coefficient: {}, Probability {}", r, p);
105        assert_almost_eq!(r, 1., 1e-14);
106        assert_almost_eq!(p, 0., 1e-14);
107    }
108
109    #[test]
110    fn test_corrcoef() {
111        let x_arr = Array2::random((3, 3), Uniform::new(-40., 200.));
112        let r = corrcoef_2d(x_arr, None, false, None, true).unwrap();
113        println!("{:?}", r);
114
115        let x_arr = Array2::random((3, 3), Uniform::new(-40., 200.));
116        let y_arr = Array2::random((3, 3), Uniform::new(-40., 200.));
117        let r = corrcoef_2d(x_arr, Some(y_arr), false, None, true).unwrap();
118        println!("{:?}", r);
119    }
120}
alximo_corr/lib.rs

alximo_corr/
lib.rs
