polars_utils/
kahan_sum.rs1use 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}