fix: accumulate central moments relative to the first observation

src/statistics/online.rs

@@ -15,15 +15,26 @@ use num_traits::Float as _;
 /// - `2` + variance
 /// - `3` + skewness
 /// - above does not presently implement further moments
+///
+/// Moments are accumulated for `x - offset`, where `offset` is the first value
+/// pushed. Central moments are invariant under that shift, and it is what makes
+/// the accumulator usable on data with a large offset: Welford's `mean +=
+/// delta / n` update cannot represent a small increment against a large running
+/// mean, so `1e12 + U(0, 1)` came out with `2.5e-4` relative error in the
+/// variance. Referring everything to the first observation keeps the magnitudes
+/// small and brings that to `5e-15` at no cost (statrs-dev/statrs#376).
 pub struct OnlineMoments<const ORDER: usize> {
     pub count: u64,
+    /// The first value pushed; `m` holds the moments of `x - offset`.
+    offset: f64,
     m: [f64; ORDER],
 }
 
 impl<const ORDER: usize> Default for OnlineMoments<ORDER> {
     fn default() -> Self {
         Self {
             count: 0,
+            offset: 0.0,
             m: [0.0; ORDER],
         }
     }
@@ -35,7 +46,7 @@ impl OnlineMoments<2> {
         if self.count == 0 {
             None
         } else {
-            Some(self.m[0])
+            Some(self.offset + self.m[0])
         }
     }
 
@@ -78,7 +89,7 @@ impl OnlineMoments<3> {
         if self.count == 0 {
             None
         } else {
-            Some(self.m[0])
+            Some(self.offset + self.m[0])
         }
     }
 
@@ -132,6 +143,60 @@ impl OnlineMoments<3> {
     }
 }
 
+impl<const ORDER: usize> OnlineMoments<ORDER> {
+    /// Merges two accumulators as if all observations had been pushed into
+    /// one, using the pairwise update of Chan, Golub & LeVeque (extended to
+    /// the third moment by Pébay, 2008).
+    ///
+    /// Besides combining accumulators built in parallel, folding into several
+    /// accumulators and merging at the end is also slightly *better*
+    /// conditioned than one long Welford chain, since each chain's rounding
+    /// errors accumulate over fewer updates.
+    ///
+    /// ```
+    /// use statrs::statistics::OnlineVariance;
+    /// use statrs::statistics::Accumulate;
+    /// let a = [1.0_f64, 2.0].iter().copied().fold(OnlineVariance::default(), OnlineVariance::push);
+    /// let b = [3.0_f64, 4.0].iter().copied().fold(OnlineVariance::default(), OnlineVariance::push);
+    /// let all = [1.0_f64, 2.0, 3.0, 4.0].iter().copied().fold(OnlineVariance::default(), OnlineVariance::push);
+    /// assert_eq!(a.merge(b).variance(), all.variance());
+    /// ```
+    pub fn merge(self, other: Self) -> Self {
+        if other.count == 0 {
+            return self;
+        }
+        if self.count == 0 {
+            return other;
+        }
+        let na = self.count as f64;
+        let nb = other.count as f64;
+        let n = na + nb;
+        // The two accumulators generally have different offsets, so re-express
+        // `other`'s mean in `self`'s frame. Grouping the two differences
+        // separately keeps this accurate when the offsets are close, which is
+        // the common case (both are data values).
+        let delta = (other.offset - self.offset) + (other.m[0] - self.m[0]);
+
+        let mut m = [0.0; ORDER];
+        m[0] = self.m[0] + delta * nb / n;
+        if let (Some(&m2a), Some(&m2b)) = (self.m.get(1), other.m.get(1)) {
+            if let (Some(&m3a), Some(&m3b)) = (self.m.get(2), other.m.get(2)) {
+                m[2] = m3a
+                    + m3b
+                    + delta * delta * delta * na * nb * (na - nb) / (n * n)
+                    + 3.0 * delta * (na * m2b - nb * m2a) / n;
+            }
+            m[1] = m2a + m2b + delta * delta * na * nb / n;
+        }
+
+        Self {
+            count: self.count + other.count,
+            offset: self.offset,
+            m,
+        }
+    }
+}
+
 /// Single-pass mean accumulator (alias of [`OnlineMoments<2>`]).
 pub type OnlineMean = OnlineMoments<2>;
 
@@ -151,8 +216,13 @@ impl<const ORDER: usize> crate::statistics::Accumulate for OnlineMoments<ORDER>
     ///     .fold(OnlineVariance::default(), OnlineVariance::push);
     /// ```
     fn push(mut self, x: f64) -> Self {
+        if self.count == 0 {
+            self.offset = x;
+        }
         self.count += 1;
         let n = self.count as f64;
+        // work relative to the first observation; see the type-level docs
+        let x = x - self.offset;
 
         // Welford / Pebay (2008) central moment update. Update order: M3
         // before M2 before mean; each step uses the previous observation's
@@ -208,6 +278,71 @@ mod tests {
         prec::assert_abs_diff_eq!(s.population_std_dev().unwrap(), 2.0);
     }
 
+    /// Welford's `mean += delta / n` cannot represent a small increment against
+    /// a large running mean, so `1e12 + U(0, 1)` used to come out with `2.5e-4`
+    /// relative error in the variance (statrs-dev/statrs#376). Accumulating
+    /// relative to the first observation keeps the magnitudes small.
+    ///
+    /// The offsets and the step are powers of two and the values need only 51
+    /// bits, so every sample is exactly representable at every offset and the
+    /// reference variance is exact - otherwise quantisation at `2^40` would
+    /// swamp what is being measured.
+    #[test]
+    fn variance_is_accurate_for_data_with_a_large_offset() {
+        const STEP: f64 = 1.0 / 1024.0; // 2^-10
+        const PERIOD: usize = 1024;
+        const REPEATS: usize = 100;
+        let n = (PERIOD * REPEATS) as f64;
+        // population variance of {0, STEP, ..., 1023 * STEP}
+        let pop = ((PERIOD * PERIOD - 1) as f64 / 12.0) * STEP * STEP;
+        let expected_variance = pop * n / (n - 1.0);
+        let expected_mean_offset = (PERIOD - 1) as f64 / 2.0 * STEP;
+
+        for exp in [0i32, 10, 20, 30, 40] {
+            let offset = f64::powi(2.0, exp);
+            // Folded straight off the iterator rather than collected, so the
+            // test also builds under `no_std`, where there is no `Vec`.
+            let s = (0..PERIOD * REPEATS)
+                .map(|i| offset + (i % PERIOD) as f64 * STEP)
+                .fold(OnlineMoments::<2>::default(), OnlineMoments::push);
+            prec::assert_relative_eq!(
+                s.variance().unwrap(),
+                expected_variance,
+                epsilon = 0.0,
+                max_relative = 1e-13
+            );
+            prec::assert_relative_eq!(
+                s.mean().unwrap(),
+                offset + expected_mean_offset,
+                epsilon = 0.0,
+                // the accumulated error lives in the shifted mean, so it is
+                // largest relative to the total when the offset is small
+                max_relative = 1e-14
+            );
+        }
+    }
+
+    /// The shift is per-accumulator, so `merge` has to reconcile two different
+    /// offsets; check it still matches a single chain on offset data.
+    #[test]
+    fn merge_reconciles_different_offsets() {
+        // Closures returning fresh iterators, so the same data can be folded
+        // three ways without collecting it (there is no `Vec` under `no_std`).
+        let a = || (0..500).map(|i| 1e12 + i as f64 * 1e-3);
+        let b = || (0..500).map(|i| 1e12 + 5.0 + i as f64 * 1e-3);
+        let ma = a().fold(OnlineMoments::<2>::default(), OnlineMoments::push);
+        let mb = b().fold(OnlineMoments::<2>::default(), OnlineMoments::push);
+        let mw = a()
+            .chain(b())
+            .fold(OnlineMoments::<2>::default(), OnlineMoments::push);
+        prec::assert_relative_eq!(
+            ma.merge(mb).variance().unwrap(),
+            mw.variance().unwrap(),
+            epsilon = 0.0,
+            max_relative = 1e-12
+        );
+    }
+
     #[test]
     fn nan_propagates() {
         let s = [1.0_f64, f64::NAN]
@@ -230,6 +365,59 @@ mod tests {
         prec::assert_abs_diff_eq!(s.skewness().unwrap(), 0.65625);
     }
 
+    #[test]
+    fn merge_matches_single_accumulator() {
+        let data = [3.0_f64, -1.0, 4.0, 1.0, -5.0, 9.0, 2.0, 6.0, -3.0];
+        for split in 0..=data.len() {
+            let (lo, hi) = data.split_at(split);
+            let a = lo
+                .iter()
+                .copied()
+                .fold(OnlineMoments::<3>::default(), OnlineMoments::push);
+            let b = hi
+                .iter()
+                .copied()
+                .fold(OnlineMoments::<3>::default(), OnlineMoments::push);
+            let merged = a.merge(b);
+            let whole = data
+                .iter()
+                .copied()
+                .fold(OnlineMoments::<3>::default(), OnlineMoments::push);
+            assert_eq!(merged.count, whole.count);
+            prec::assert_relative_eq!(
+                merged.mean().unwrap(),
+                whole.mean().unwrap(),
+                max_relative = 1e-14
+            );
+            prec::assert_relative_eq!(
+                merged.variance().unwrap(),
+                whole.variance().unwrap(),
+                max_relative = 1e-13
+            );
+            prec::assert_relative_eq!(
+                merged.skewness().unwrap(),
+                whole.skewness().unwrap(),
+                max_relative = 1e-12
+            );
+        }
+    }
+
+    #[test]
+    fn merge_with_empty_is_identity() {
+        let a = [1.0_f64, 2.0, 3.0]
+            .iter()
+            .copied()
+            .fold(OnlineMoments::<2>::default(), OnlineMoments::push);
+        let empty = OnlineMoments::<2>::default();
+        assert_eq!(a.merge(empty).variance(), Some(1.0));
+        let a = [1.0_f64, 2.0, 3.0]
+            .iter()
+            .copied()
+            .fold(OnlineMoments::<2>::default(), OnlineMoments::push);
+        let empty = OnlineMoments::<2>::default();
+        assert_eq!(empty.merge(a).variance(), Some(1.0));
+    }
+
     #[test]
     fn order_3_mean_and_variance_match_order_2() {
         let data = [2.0_f64, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0];