Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
pola-rs
GitHub Repository: pola-rs/polars
Path: blob/main/crates/polars-utils/src/kahan_sum.rs
8445 views
1
use std::ops::{Add, AddAssign};
2
3
use num_traits::Num;
4
5
use crate::float::IsFloat;
6
7
#[derive(Debug, Clone)]
8
pub struct KahanSum<T> {
9
sum: T,
10
err: T,
11
}
12
13
impl<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
26
impl<T: Num> Default for KahanSum<T> {
27
fn default() -> Self {
28
KahanSum {
29
sum: T::zero(),
30
err: T::zero(),
31
}
32
}
33
}
34
35
impl<T: IsFloat + Num + AddAssign + Copy> AddAssign<T> for KahanSum<T> {
36
fn add_assign(&mut self, rhs: T) {
37
let y = rhs - self.err;
38
let new_sum = self.sum + y;
39
let new_err = (new_sum - self.sum) - y;
40
self.sum = new_sum;
41
if new_err.is_finite() {
42
// Ensure err stays finite so we don't introduce NaNs through Inf - Inf.
43
self.err = new_err;
44
}
45
}
46
}
47
48
impl<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
}
57
58