Skip to main content

rill_ml/stats/
mean.rs

1//! Online mean using a numerically stable incremental update.
2//!
3//! Time complexity per update: `O(1)`. Space complexity: `O(1)`.
4
5use crate::error::{RillError, checked_increment, ensure_finite};
6use crate::traits::OnlineStatistic;
7
8/// Incremental mean computed with the delta method to minimize floating-point
9/// accumulation error.
10///
11/// # Examples
12///
13/// ```
14/// use rill_ml::stats::Mean;
15/// use rill_ml::OnlineStatistic;
16///
17/// let mut m = Mean::new();
18/// m.update(1.0).unwrap();
19/// m.update(2.0).unwrap();
20/// m.update(3.0).unwrap();
21/// assert_eq!(m.value(), 2.0);
22/// assert_eq!(m.samples_seen(), 3);
23/// ```
24#[derive(Debug, Clone, Default)]
25#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
26pub struct Mean {
27    count: u64,
28    mean: f64,
29}
30
31impl Mean {
32    /// Create a new empty mean accumulator.
33    pub const fn new() -> Self {
34        Self {
35            count: 0,
36            mean: 0.0,
37        }
38    }
39
40    /// Current mean, or `0.0` if no observations have been seen.
41    pub const fn value(&self) -> f64 {
42        self.mean
43    }
44
45    /// Number of observations seen so far.
46    pub const fn count(&self) -> u64 {
47        self.count
48    }
49}
50
51impl OnlineStatistic for Mean {
52    fn update(&mut self, value: f64) -> Result<(), RillError> {
53        ensure_finite("value", value)?;
54        let next_count = checked_increment(self.count, "mean sample")?;
55        let delta = value - self.mean;
56        ensure_finite("mean delta", delta)?;
57        let next_mean = self.mean + delta / next_count as f64;
58        ensure_finite("mean", next_mean)?;
59
60        self.count = next_count;
61        self.mean = next_mean;
62        Ok(())
63    }
64
65    fn samples_seen(&self) -> u64 {
66        self.count
67    }
68
69    fn reset(&mut self) {
70        self.count = 0;
71        self.mean = 0.0;
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78    use rand::SeedableRng;
79
80    #[test]
81    fn mean_of_simple_sequence() {
82        let mut m = Mean::new();
83        for x in [1.0, 2.0, 3.0, 4.0, 5.0] {
84            m.update(x).unwrap();
85        }
86        assert_eq!(m.value(), 3.0);
87        assert_eq!(m.count(), 5);
88    }
89
90    #[test]
91    fn mean_empty_is_zero() {
92        let m = Mean::new();
93        assert_eq!(m.value(), 0.0);
94        assert_eq!(m.count(), 0);
95    }
96
97    #[test]
98    fn mean_rejects_non_finite() {
99        let mut m = Mean::new();
100        assert!(m.update(f64::NAN).is_err());
101        assert!(m.update(f64::INFINITY).is_err());
102        assert_eq!(m.count(), 0);
103    }
104
105    #[test]
106    fn mean_rejects_overflow_without_mutating_state() {
107        let mut m = Mean::new();
108        m.update(f64::MAX).unwrap();
109        let before = m.clone();
110        assert!(m.update(-f64::MAX).is_err());
111        assert_eq!(m.count(), before.count());
112        assert_eq!(m.value(), before.value());
113    }
114
115    #[test]
116    fn mean_reset() {
117        let mut m = Mean::new();
118        m.update(10.0).unwrap();
119        m.update(20.0).unwrap();
120        m.reset();
121        assert_eq!(m.count(), 0);
122        assert_eq!(m.value(), 0.0);
123    }
124
125    #[test]
126    fn mean_matches_batch_formula() {
127        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(42);
128        let mut m = Mean::new();
129        let mut data = Vec::new();
130        for _ in 0..1000 {
131            let x = rand::Rng::gen_range(&mut rng, -100.0..100.0);
132            m.update(x).unwrap();
133            data.push(x);
134        }
135        let batch: f64 = data.iter().sum::<f64>() / data.len() as f64;
136        assert!((m.value() - batch).abs() < 1e-9);
137    }
138}