volas_core/
stats.rs

1//! Pure NaN-skipping numeric algorithms over `f64` slices, shared by the Series
2//! / DataFrame numeric methods in the binding layer. They live here (not in the
3//! pyo3 binding) so they stay unit-tested and coverage-measured.
4//!
5//! NaN handling mirrors pandas's `skipna=True`: a NaN does not contribute to a
6//! running accumulator / reduction, and in the cumulative forms it is kept as
7//! NaN in place.
8//!
9//! The cumulative kernels are generic over [`Numeric`] so they preserve the
10//! input dtype (int stays int) and compute natively (no f64 round-trip); the
11//! reductions / moment kernels stay f64. A missing element (`NaN`, f64 only —
12//! `i64::is_missing` is always false) passes through unchanged.
13
14use std::cmp::Ordering;
15
16use crate::numeric::Numeric;
17
18/// Cumulative sum (pandas `cumsum`, skipna=True), dtype-preserving.
19pub fn cumsum<T: Numeric>(v: &[T]) -> Vec<T> {
20    let mut acc = T::ZERO;
21    v.iter()
22        .map(|&x| {
23            if x.is_missing() {
24                x
25            } else {
26                acc = acc.wrapping_add(x);
27                acc
28            }
29        })
30        .collect()
31}
32
33/// Cumulative maximum (pandas `cummax`, skipna=True), dtype-preserving.
34pub fn cummax<T: Numeric>(v: &[T]) -> Vec<T> {
35    cum_extreme(v, true)
36}
37
38/// Cumulative minimum (pandas `cummin`, skipna=True), dtype-preserving.
39pub fn cummin<T: Numeric>(v: &[T]) -> Vec<T> {
40    cum_extreme(v, false)
41}
42
43/// Shared running-extreme for `cummax` (`want_max`) / `cummin`. NaN is excluded
44/// before any comparison, so the partial order is total over what we compare.
45fn cum_extreme<T: Numeric>(v: &[T], want_max: bool) -> Vec<T> {
46    let mut acc: Option<T> = None;
47    v.iter()
48        .map(|&x| {
49            if x.is_missing() {
50                return x;
51            }
52            let next = match acc {
53                Some(a) if (x > a) != want_max => a,
54                _ => x,
55            };
56            acc = Some(next);
57            next
58        })
59        .collect()
60}
61
62/// Element-wise absolute value (pandas `abs`), dtype-preserving; a missing value
63/// passes through. Generic so the wrapping `abs` resolves to the [`Numeric`] impl
64/// (matching pandas int64 overflow: `abs(i64::MIN) == i64::MIN`).
65pub fn abs<T: Numeric>(v: &[T]) -> Vec<T> {
66    v.iter()
67        .map(|&x| if x.is_missing() { x } else { x.wrapping_abs() })
68        .collect()
69}
70
71/// Cumulative product (pandas `cumprod`, skipna=True), dtype-preserving.
72pub fn cumprod<T: Numeric>(v: &[T]) -> Vec<T> {
73    let mut acc = T::ONE;
74    v.iter()
75        .map(|&x| {
76            if x.is_missing() {
77                x
78            } else {
79                acc = acc.wrapping_mul(x);
80                acc
81            }
82        })
83        .collect()
84}
85
86/// Missing-skipping sum (`0` when empty / all-missing, matching pandas),
87/// dtype-preserving (i64 sums in i64, wrapping like pandas).
88pub fn sum<T: Numeric>(v: &[T]) -> T {
89    let mut acc = T::ZERO;
90    for &x in v {
91        if !x.is_missing() {
92            acc = acc.wrapping_add(x);
93        }
94    }
95    acc
96}
97
98/// Missing-skipping product (`1` when empty / all-missing, matching pandas),
99/// dtype-preserving.
100pub fn prod<T: Numeric>(v: &[T]) -> T {
101    let mut acc = T::ONE;
102    for &x in v {
103        if !x.is_missing() {
104            acc = acc.wrapping_mul(x);
105        }
106    }
107    acc
108}
109
110/// Missing-skipping minimum (`want_max = false`) / maximum, dtype-preserving;
111/// `None` when empty / all-missing. The `want_max` branch is hoisted out of the
112/// loop so each fold is a single tight comparison (as fast as a specialized
113/// min/max, and no intermediate allocation).
114pub fn extreme<T: Numeric>(v: &[T], want_max: bool) -> Option<T> {
115    let mut it = v.iter().copied().filter(|x| !x.is_missing());
116    let first = it.next()?;
117    Some(if want_max {
118        it.fold(first, |a, x| if x > a { x } else { a })
119    } else {
120        it.fold(first, |a, x| if x < a { x } else { a })
121    })
122}
123
124/// The finite (non-NaN) values, in order.
125fn non_nan(v: &[f64]) -> Vec<f64> {
126    v.iter().copied().filter(|x| !x.is_nan()).collect()
127}
128
129/// Standard error of the mean (pandas `sem`, ddof=1): `sample_std / sqrt(n)`.
130/// NaN with fewer than two finite values.
131pub fn sem(v: &[f64]) -> f64 {
132    let xs = non_nan(v);
133    let n = xs.len();
134    if n < 2 {
135        return f64::NAN;
136    }
137    let nf = n as f64;
138    let mean = xs.iter().sum::<f64>() / nf;
139    let var = xs.iter().map(|&x| (x - mean) * (x - mean)).sum::<f64>() / (nf - 1.0);
140    (var / nf).sqrt()
141}
142
143/// Adjusted Fisher-Pearson sample skewness (pandas `skew`). NaN with fewer than
144/// three finite values; 0.0 for a zero-variance sample.
145pub fn skew(v: &[f64]) -> f64 {
146    let xs = non_nan(v);
147    let n = xs.len();
148    if n < 3 {
149        return f64::NAN;
150    }
151    let nf = n as f64;
152    let mean = xs.iter().sum::<f64>() / nf;
153    let m2: f64 = xs.iter().map(|&x| (x - mean).powi(2)).sum();
154    let m3: f64 = xs.iter().map(|&x| (x - mean).powi(3)).sum();
155    if m2 == 0.0 {
156        return 0.0;
157    }
158    (nf * (nf - 1.0).sqrt() / (nf - 2.0)) * (m3 / m2.powf(1.5))
159}
160
161/// Bias-corrected excess kurtosis (pandas `kurt`, Fisher's definition). NaN with
162/// fewer than four finite values; 0.0 for a zero-variance sample.
163pub fn kurt(v: &[f64]) -> f64 {
164    let xs = non_nan(v);
165    let n = xs.len();
166    if n < 4 {
167        return f64::NAN;
168    }
169    let nf = n as f64;
170    let mean = xs.iter().sum::<f64>() / nf;
171    let m2: f64 = xs.iter().map(|&x| (x - mean).powi(2)).sum();
172    let m4: f64 = xs.iter().map(|&x| (x - mean).powi(4)).sum();
173    if m2 == 0.0 {
174        return 0.0;
175    }
176    let num = nf * (nf + 1.0) * (nf - 1.0) * m4;
177    let den = (nf - 2.0) * (nf - 3.0) * m2 * m2;
178    let adj = 3.0 * (nf - 1.0).powi(2) / ((nf - 2.0) * (nf - 3.0));
179    num / den - adj
180}
181
182/// Element-wise choose: `cond[i] ? a[i] : b[i]`. Backs `where` / `mask`. Generic
183/// so it preserves dtype (picks i64 natively). `cond` / `a` / `b` are equal length.
184pub fn select<T: Numeric>(cond: &[bool], a: &[T], b: &[T]) -> Vec<T> {
185    cond.iter()
186        .enumerate()
187        .map(|(i, &c)| if c { a[i] } else { b[i] })
188        .collect()
189}
190
191/// The aligned `(x, y)` pairs with neither value NaN (pandas pairwise NaN drop).
192fn pairs(x: &[f64], y: &[f64]) -> Vec<(f64, f64)> {
193    x.iter()
194        .zip(y)
195        .filter(|(a, b)| !a.is_nan() && !b.is_nan())
196        .map(|(&a, &b)| (a, b))
197        .collect()
198}
199
200/// Pairwise sample covariance, ddof=1 (pandas `cov`); NaN with fewer than two
201/// finite pairs. `x` / `y` are aligned positionally.
202pub fn cov(x: &[f64], y: &[f64]) -> f64 {
203    let p = pairs(x, y);
204    let n = p.len();
205    if n < 2 {
206        return f64::NAN;
207    }
208    let nf = n as f64;
209    let mx = p.iter().map(|t| t.0).sum::<f64>() / nf;
210    let my = p.iter().map(|t| t.1).sum::<f64>() / nf;
211    p.iter().map(|(a, b)| (a - mx) * (b - my)).sum::<f64>() / (nf - 1.0)
212}
213
214/// Pairwise Pearson correlation (pandas `corr`); NaN with fewer than two finite
215/// pairs or a zero-variance input. `x` / `y` are aligned positionally.
216pub fn corr(x: &[f64], y: &[f64]) -> f64 {
217    let p = pairs(x, y);
218    let n = p.len();
219    if n < 2 {
220        return f64::NAN;
221    }
222    let nf = n as f64;
223    let mx = p.iter().map(|t| t.0).sum::<f64>() / nf;
224    let my = p.iter().map(|t| t.1).sum::<f64>() / nf;
225    let (mut sxy, mut sxx, mut syy) = (0.0, 0.0, 0.0);
226    for (a, b) in &p {
227        let (dx, dy) = (a - mx, b - my);
228        sxy += dx * dy;
229        sxx += dx * dx;
230        syy += dy * dy;
231    }
232    if sxx == 0.0 || syy == 0.0 {
233        return f64::NAN;
234    }
235    // sqrt(sxx*syy) (not sqrt(sxx)*sqrt(syy)) so a column with itself is exactly
236    // 1.0; clamp to [-1, 1] against rounding (matching NumPy / pandas).
237    (sxy / (sxx * syy).sqrt()).clamp(-1.0, 1.0)
238}
239
240/// How tied values share ranks (pandas `rank(method=)`).
241#[derive(Clone, Copy, PartialEq, Eq)]
242pub enum RankMethod {
243    /// Mean of the tied positions (pandas default).
244    Average,
245    /// Lowest tied position.
246    Min,
247    /// Highest tied position.
248    Max,
249    /// Distinct, in order of appearance.
250    First,
251    /// Ties share one rank; the next group is +1 (no gaps).
252    Dense,
253}
254
255/// Rank the finite values (pandas `rank`, 1-based, `na_option='keep'` so NaN
256/// stays NaN). `pct` divides by the count of finite values (max dense rank for
257/// `Dense`).
258pub fn rank(v: &[f64], method: RankMethod, ascending: bool, pct: bool) -> Vec<f64> {
259    rank_by(
260        v.len(),
261        |i| !v[i].is_nan(),
262        |a, b| v[a].partial_cmp(&v[b]).unwrap(),
263        method,
264        ascending,
265        pct,
266    )
267}
268
269/// Order-based rank over `n` positions, driven by an `is_valid` predicate and a
270/// total `cmp` over present positions — so the one rank algorithm (tie handling,
271/// `pct`) serves every ordered dtype (numeric by value, `str` lexically,
272/// `datetime` by raw `i64`), not just `f64`. Invalid positions rank `NaN`; `cmp`
273/// is only ever called on positions that pass `is_valid`.
274pub fn rank_by(
275    n: usize,
276    is_valid: impl Fn(usize) -> bool,
277    cmp: impl Fn(usize, usize) -> Ordering,
278    method: RankMethod,
279    ascending: bool,
280    pct: bool,
281) -> Vec<f64> {
282    let mut idx: Vec<usize> = (0..n).filter(|&i| is_valid(i)).collect();
283    let cnt = idx.len();
284    // Sort surviving positions by value (stable, so `First` keeps input order).
285    idx.sort_by(|&a, &b| {
286        let o = cmp(a, b);
287        if ascending {
288            o
289        } else {
290            o.reverse()
291        }
292    });
293
294    let mut out = vec![f64::NAN; n];
295    let mut dense_rank = 0.0;
296    let mut i = 0;
297    while i < idx.len() {
298        // [i, j) is a run of tied values (equal regardless of sort direction).
299        let mut j = i + 1;
300        while j < idx.len() && cmp(idx[j], idx[i]) == Ordering::Equal {
301            j += 1;
302        }
303        dense_rank += 1.0;
304        for (k, &pos) in idx[i..j].iter().enumerate() {
305            out[pos] = match method {
306                RankMethod::Average => (i + j + 1) as f64 / 2.0, // mean of 1-based [i+1, j]
307                RankMethod::Min => (i + 1) as f64,
308                RankMethod::Max => j as f64,
309                RankMethod::First => (i + k + 1) as f64,
310                RankMethod::Dense => dense_rank,
311            };
312        }
313        i = j;
314    }
315
316    if pct && cnt > 0 {
317        let denom = if method == RankMethod::Dense {
318            dense_rank
319        } else {
320            cnt as f64
321        };
322        for x in out.iter_mut() {
323            if !x.is_nan() {
324                *x /= denom;
325            }
326        }
327    }
328    out
329}
330
331#[cfg(test)]
332mod tests {
333    use super::*;
334
335    fn nan() -> f64 {
336        f64::NAN
337    }
338
339    fn eq(a: &[f64], b: &[f64]) -> bool {
340        a.len() == b.len()
341            && a.iter()
342                .zip(b)
343                .all(|(&x, &y)| (x == y) || (x.is_nan() && y.is_nan()))
344    }
345
346    #[test]
347    fn cumulatives_skip_nan() {
348        assert!(eq(&cumsum(&[1.0, nan(), 2.0, 3.0]), &[1.0, nan(), 3.0, 6.0]));
349        assert!(eq(&cummax(&[1.0, nan(), 2.0, 4.0]), &[1.0, nan(), 2.0, 4.0]));
350        assert!(eq(&cummin(&[3.0, nan(), 1.0, 2.0]), &[3.0, nan(), 1.0, 1.0]));
351        assert!(eq(&cumprod(&[1.0, nan(), 2.0, 4.0]), &[1.0, nan(), 2.0, 8.0]));
352        // leading NaN: accumulator stays unset until the first finite value
353        assert!(eq(&cummax(&[nan(), 5.0, 3.0]), &[nan(), 5.0, 5.0]));
354        assert!(eq(&cumsum(&[]), &[]));
355    }
356
357    #[test]
358    fn reductions_skip_missing_and_preserve_dtype() {
359        // float: skip NaN; empty/all-NaN -> identity
360        assert_eq!(prod(&[1.0, nan(), 2.0, 3.0]), 6.0);
361        assert_eq!(prod(&[nan(), nan()]), 1.0);
362        assert_eq!(prod::<f64>(&[]), 1.0);
363        assert_eq!(sum(&[1.0, nan(), 2.0, 3.0]), 6.0);
364        assert_eq!(sum::<f64>(&[]), 0.0);
365        // int: native i64 (wrapping like pandas)
366        assert_eq!(sum(&[1_i64, 2, 3]), 6);
367        assert_eq!(prod(&[2_i64, 3, 4]), 24);
368        assert_eq!(sum(&[i64::MAX, 1]), i64::MIN); // wraps
369        // extreme: skip NaN, None when empty/all-missing
370        assert_eq!(extreme(&[3.0, nan(), 1.0, 2.0], false), Some(1.0));
371        assert_eq!(extreme(&[3.0, nan(), 1.0, 2.0], true), Some(3.0));
372        assert_eq!(extreme(&[1_i64, 5, 2], true), Some(5));
373        assert_eq!(extreme::<f64>(&[], false), None);
374        assert_eq!(extreme(&[nan(), nan()], true), None);
375    }
376
377    fn close(a: f64, b: f64) -> bool {
378        (a - b).abs() <= 1e-9 * a.abs().max(1.0)
379    }
380
381    #[test]
382    fn moments_match_pandas() {
383        // values cross-checked against pandas 3.0
384        let v = [2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0];
385        assert!(close(sem(&v), 0.755_928_946_018));
386        assert!(close(skew(&v), 0.818_487_553_4));
387        assert!(close(kurt(&v), 0.940_625));
388        // NaN is skipped before counting
389        assert!(close(sem(&[2.0, nan(), 4.0]), sem(&[2.0, 4.0])));
390        // too-few-values guards
391        assert!(sem(&[5.0]).is_nan());
392        assert!(skew(&[1.0, 2.0]).is_nan());
393        assert!(kurt(&[1.0, 2.0, 3.0]).is_nan());
394        // zero-variance samples -> 0.0
395        assert_eq!(skew(&[3.0, 3.0, 3.0]), 0.0);
396        assert_eq!(kurt(&[3.0, 3.0, 3.0, 3.0]), 0.0);
397    }
398
399    #[test]
400    fn rank_methods_and_pct() {
401        use RankMethod::*;
402        let v = [3.0, 1.0, 1.0, 2.0, nan()];
403        let r = |m, asc, pct| rank(&v, m, asc, pct);
404        assert!(eq(&r(Average, true, false), &[4.0, 1.5, 1.5, 3.0, nan()]));
405        assert!(eq(&r(Min, true, false), &[4.0, 1.0, 1.0, 3.0, nan()]));
406        assert!(eq(&r(Max, true, false), &[4.0, 2.0, 2.0, 3.0, nan()]));
407        assert!(eq(&r(First, true, false), &[4.0, 1.0, 2.0, 3.0, nan()]));
408        assert!(eq(&r(Dense, true, false), &[3.0, 1.0, 1.0, 2.0, nan()]));
409        // descending
410        assert!(eq(&r(Min, false, false), &[1.0, 3.0, 3.0, 2.0, nan()]));
411        // pct: average / count(=4); dense / max-dense(=3)
412        assert!(eq(&r(Average, true, true), &[1.0, 0.375, 0.375, 0.75, nan()]));
413        assert!(eq(&r(Dense, true, true), &[1.0, 1.0 / 3.0, 1.0 / 3.0, 2.0 / 3.0, nan()]));
414        // all-NaN: every rank is NaN, pct is a no-op
415        assert!(eq(&rank(&[nan(), nan()], Average, true, true), &[nan(), nan()]));
416    }
417
418    #[test]
419    fn corr_cov_pairwise() {
420        assert!(close(corr(&[1.0, 2.0, 3.0], &[1.0, 2.0, 3.0]), 1.0));
421        assert!(close(corr(&[1.0, 2.0, 3.0], &[3.0, 2.0, 1.0]), -1.0));
422        assert!(close(cov(&[1.0, 2.0, 3.0], &[1.0, 2.0, 3.0]), 1.0)); // == sample var
423        // NaN pairs are dropped (drops index 1 -> corr([1,3],[1,3]) == 1)
424        assert!(close(corr(&[1.0, nan(), 3.0], &[1.0, 5.0, 3.0]), 1.0));
425        // guards: fewer than two pairs / zero variance -> NaN
426        assert!(corr(&[1.0], &[2.0]).is_nan());
427        assert!(cov(&[1.0], &[2.0]).is_nan());
428        assert!(corr(&[1.0, 1.0, 1.0], &[1.0, 2.0, 3.0]).is_nan());
429    }
430
431    #[test]
432    fn select_picks_per_condition() {
433        let cond = [true, false, true];
434        let a = [1.0, 2.0, 3.0];
435        let b = [10.0, 20.0, 30.0];
436        assert_eq!(select(&cond, &a, &b), vec![1.0, 20.0, 3.0]);
437        // backs `mask` by swapping the branches
438        assert_eq!(select(&cond, &b, &a), vec![10.0, 2.0, 30.0]);
439    }
440}
volas_core/stats.rs

volas_core/
stats.rs
