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}