fix(backport): Mean/WeightedMean accumulator addition fix (#537)

* fix: Mean accumulator math fix

* fix: also fix WeighedMean

* tests: add summing tests

Author: henryiii

include/bh_python/accumulators/mean.hpp

@@ -59,12 +59,20 @@ struct mean {
     }
 
     mean& operator+=(const mean& rhs) noexcept {
-        if(count != 0 || rhs.count != 0) {
-            const auto tmp = value * count + rhs.value * rhs.count;
-            count += rhs.count;
-            value = tmp / count;
-        }
+        if(rhs.count == 0)
+            return *this;
+
+        const auto mu1 = value;
+        const auto mu2 = rhs.value;
+        const auto n1  = count;
+        const auto n2  = rhs.count;
+
+        count += rhs.count;
+        value = (n1 * mu1 + n2 * mu2) / count;
         sum_of_deltas_squared += rhs.sum_of_deltas_squared;
+        sum_of_deltas_squared
+            += n1 * (value - mu1) * (value - mu1) + n2 * (value - mu2) * (value - mu2);
+
         return *this;
     }
 

include/bh_python/accumulators/weighted_mean.hpp

@@ -62,13 +62,23 @@ struct weighted_mean {
     }
 
     weighted_mean& operator+=(const weighted_mean& rhs) {
-        if(sum_of_weights != 0 || rhs.sum_of_weights != 0) {
-            const auto tmp = value * sum_of_weights + rhs.value * rhs.sum_of_weights;
-            sum_of_weights += rhs.sum_of_weights;
-            sum_of_weights_squared += rhs.sum_of_weights_squared;
-            value = tmp / sum_of_weights;
-        }
+        if(rhs.sum_of_weights == 0)
+            return *this;
+
+        const auto mu1 = value;
+        const auto mu2 = rhs.value;
+        const auto n1  = sum_of_weights;
+        const auto n2  = rhs.sum_of_weights;
+
+        sum_of_weights += rhs.sum_of_weights;
+        sum_of_weights_squared += rhs.sum_of_weights_squared;
+
+        value = (n1 * mu1 + n2 * mu2) / sum_of_weights;
+
         _sum_of_weighted_deltas_squared += rhs._sum_of_weighted_deltas_squared;
+        _sum_of_weighted_deltas_squared
+            += n1 * (value - mu1) * (value - mu1) + n2 * (value - mu2) * (value - mu2);
+
         return *this;
     }
 

include/bh_python/register_accumulator.hpp

@@ -53,6 +53,13 @@ py::class_<A> register_accumulator(py::module acc, Args&&... args) {
 
         .def(py::self *= double())
 
+        .def("__add__",
+             [](const A& self, const A& other) {
+                 A retval(self);
+                 retval += other;
+                 return retval;
+             })
+
         // The c++ name is replaced with the Python name here
         .def("__repr__",
              [](py::object self) {

tests/test_accumulators.py

@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+from pytest import approx
 
 import boost_histogram as bh
 
@@ -87,3 +88,47 @@ def test_mean():
     assert a == bh.accumulators.Mean(3, 2, 1)
 
     assert repr(a) == "Mean(count=3, value=2, variance=1)"
+
+
+def test_sum_mean():
+    a = bh.accumulators.Mean()
+    a.fill([1, 2, 3])
+
+    b = bh.accumulators.Mean()
+    b.fill([5, 6])
+
+    c = bh.accumulators.Mean()
+    c.fill([1, 2, 3, 5, 6])
+
+    ab = a + b
+    assert ab.value == approx(c.value)
+    assert ab.variance == approx(c.variance)
+    assert ab.count == approx(c.count)
+
+    a += b
+    assert a.value == approx(c.value)
+    assert a.variance == approx(c.variance)
+    assert a.count == approx(c.count)
+
+
+def test_sum_weighed_mean():
+    a = bh.accumulators.WeightedMean()
+    a.fill([1, 2, 3], weight=[2, 5, 3])
+
+    b = bh.accumulators.WeightedMean()
+    b.fill([5, 6], weight=[12, 17])
+
+    c = bh.accumulators.WeightedMean()
+    c.fill([1, 2, 3, 5, 6], weight=[2, 5, 3, 12, 17])
+
+    ab = a + b
+    assert ab.value == approx(c.value)
+    assert ab.variance == approx(c.variance)
+    assert ab.sum_of_weights == approx(c.sum_of_weights)
+    assert ab.sum_of_weights_squared == approx(c.sum_of_weights_squared)
+
+    a += b
+    assert a.value == approx(c.value)
+    assert a.variance == approx(c.variance)
+    assert a.sum_of_weights == approx(c.sum_of_weights)
+    assert a.sum_of_weights_squared == approx(c.sum_of_weights_squared)

tests/test_storage.py

@@ -2,6 +2,7 @@
 import numpy as np
 import pytest
 from numpy.testing import assert_array_equal
+from pytest import approx
 
 import boost_histogram as bh
 
@@ -216,3 +217,45 @@ def test_modify_weights_by_view():
 
     assert hist.view().value[0] == pytest.approx(1.5)
     assert hist.view().value[1] == pytest.approx(2)
+
+
+# Issue #531
+def test_summing_mean_storage():
+    np.random.seed(42)
+    values = np.random.normal(loc=1.3, scale=0.1, size=1000)
+    samples = np.random.normal(loc=1.3, scale=0.1, size=1000)
+
+    h1 = bh.Histogram(bh.axis.Regular(20, -1, 3), storage=bh.storage.Mean())
+    h1.fill(values, sample=samples)
+
+    h2 = bh.Histogram(bh.axis.Regular(1, -1, 3), storage=bh.storage.Mean())
+    h2.fill(values, sample=samples)
+
+    s1 = h1.sum()
+    s2 = h2.sum()
+
+    assert s1.value == approx(s2.value)
+    assert s1.count == approx(s2.count)
+    assert s1.variance == approx(s2.variance)
+
+
+# Issue #531
+def test_summing_weighted_mean_storage():
+    np.random.seed(42)
+    values = np.random.normal(loc=1.3, scale=0.1, size=1000)
+    samples = np.random.normal(loc=1.3, scale=0.1, size=1000)
+    weights = np.random.uniform(0.1, 5, size=1000)
+
+    h1 = bh.Histogram(bh.axis.Regular(20, -1, 3), storage=bh.storage.WeightedMean())
+    h1.fill(values, sample=samples, weight=weights)
+
+    h2 = bh.Histogram(bh.axis.Regular(1, -1, 3), storage=bh.storage.WeightedMean())
+    h2.fill(values, sample=samples, weight=weights)
+
+    s1 = h1.sum()
+    s2 = h2.sum()
+
+    assert s1.value == approx(s2.value)
+    assert s1.sum_of_weights == approx(s2.sum_of_weights)
+    assert s1.sum_of_weights_squared == approx(s2.sum_of_weights_squared)
+    assert s1.variance == approx(s2.variance)