polars_utils/
kahan_sum.rs

1use std::ops::{Add, AddAssign};
2
3use num_traits::Float;
4
5#[derive(Debug, Clone)]
6pub struct KahanSum<T: Float> {
7    sum: T,
8    err: T,
9}
10
11impl<T: Float> KahanSum<T> {
12    pub fn new(v: T) -> Self {
13        KahanSum {
14            sum: v,
15            err: T::zero(),
16        }
17    }
18
19    pub fn sum(&self) -> T {
20        self.sum
21    }
22}
23
24impl<T: Float> Default for KahanSum<T> {
25    fn default() -> Self {
26        KahanSum {
27            sum: T::zero(),
28            err: T::zero(),
29        }
30    }
31}
32
33impl<T: Float + AddAssign> AddAssign<T> for KahanSum<T> {
34    fn add_assign(&mut self, rhs: T) {
35        if rhs.is_finite() {
36            let y = rhs - self.err;
37            let new_sum = self.sum + y;
38            self.err = (new_sum - self.sum) - y;
39            self.sum = new_sum;
40        } else {
41            self.sum += rhs
42        }
43    }
44}
45
46impl<T: Float + AddAssign> Add<T> for KahanSum<T> {
47    type Output = Self;
48
49    fn add(self, rhs: T) -> Self::Output {
50        let mut rv = self;
51        rv += rhs;
52        rv
53    }
54}