polars_utils/
kahan_sum.rs

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