fix: accumulate central moments relative to the first observation

src/statistics/online.rs

@@ -11,19 +11,27 @@ use num_traits::Float as _;
 /// Single-pass accumulator for central moments via Welford's online algorithm.
 ///
 /// `ORDER` controls which moments are tracked:
-/// - `1` mean
-/// - `2` + variance
-/// - `3` + skewness
-/// - above does not presently implement further moments
+/// - `1`: count + mean
+/// - `2`: `1` + variance
+/// - `3`: `2` + skewness
+/// - values above: not presently implemented
+///
+/// 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.
 pub struct OnlineMoments<const ORDER: usize> {
     pub count: u64,
+    // `m` holds the moments of `x - offset`. Welford's update can become
+    // insensitive reducing unscaled data. See statrs-dev/statrs#376.
+    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 +43,7 @@ impl OnlineMoments<2> {
         if self.count == 0 {
             None
         } else {
-            Some(self.m[0])
+            Some(self.offset + self.m[0])
         }
     }
 
@@ -78,7 +86,7 @@ impl OnlineMoments<3> {
         if self.count == 0 {
             None
         } else {
-            Some(self.m[0])
+            Some(self.offset + self.m[0])
         }
     }
 
@@ -132,6 +140,95 @@ 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).
+    /// In addition to API after computing on parallel streams, merging
+    /// at the end is also slightly *better* conditioned than one long Welford
+    /// chain.
+    ///
+    /// ```
+    /// 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());
+    /// ```
+    ///
+    /// # Precision
+    /// Consider breaking apart large streams and merging for precision.
+    /// Recursively merging in a binary tree accumulates roughly O(log N)
+    /// rounding error against O(N) for a single chain, the same effect as
+    /// pairwise vs. naive summation.
+    ///
+    /// ```
+    /// use statrs::statistics::{OnlineVariance, Accumulate};
+    /// use approx::assert_relative_eq;
+    ///
+    /// // Repeating a block leaves its variance unchanged, so this stream has
+    /// // an exactly known variance to check against: mean 5.0, M2 = 32 per
+    /// // block.
+    /// let block = [2.0_f64, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0];
+    /// let blocks = 1 << 16;
+    /// let exact_variance = (32.0 * blocks as f64) / (8.0 * blocks as f64 - 1.0);
+    ///
+    /// let chained: OnlineVariance = (0..blocks)
+    ///     .flat_map(|_| block.iter().copied())
+    ///     .fold(Default::default(), Accumulate::push);
+    ///
+    /// let stream = |count| {
+    ///     (0..count)
+    ///         .flat_map(|_| block.iter().copied())
+    ///         .fold(OnlineVariance::default(), OnlineVariance::push)
+    /// };
+    /// let per_stream = blocks / 4;
+    /// let merged = stream(per_stream)
+    ///     .merge(stream(per_stream))
+    ///     .merge(stream(per_stream).merge(stream(per_stream)));
+    ///
+    /// let chained_err = (chained.variance().unwrap() - exact_variance).abs();
+    /// let merged_err = (merged.variance().unwrap() - exact_variance).abs();
+    /// assert!(merged_err < chained_err);
+    /// assert_relative_eq!(merged.variance().unwrap(), exact_variance, max_relative = 1e-13);
+    /// ```
+    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>;
 
@@ -150,9 +247,18 @@ impl<const ORDER: usize> crate::statistics::Accumulate for OnlineMoments<ORDER>
     /// let s = [1.0_f64, 2.0, 3.0].iter().copied()
     ///     .fold(OnlineVariance::default(), OnlineVariance::push);
     /// ```
+    ///
+    /// # Precision
+    /// Sensitive to data ordering, especially with regard to scale of initial item.
+    /// If consuming very large streams, see [`merge`](OnlineMoments::merge)
     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 +314,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 +401,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];