From 315b57aca8d5aa4cfb817a773c8822c5fbb19292 Mon Sep 17 00:00:00 2001
From: Doron Behar <doron.behar@gmail.com>
Date: Mon, 21 Oct 2024 11:49:23 +0300
Subject: [PATCH] unumpy.average: init

Ref: https://github.com/lmfit/uncertainties/issues/38
---
 CHANGES.rst                  |  1 +
 doc/numpy_guide.rst          | 16 ++++++++++++++++
 tests/test_unumpy.py         |  8 ++++++++
 uncertainties/unumpy/core.py | 20 ++++++++++++++++++++
 4 files changed, 45 insertions(+)

diff --git a/CHANGES.rst b/CHANGES.rst
index c45ef266..9058854f 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -7,6 +7,7 @@ Unreleased
 Fixes:
 
 - fix `readthedocs` configuration so that the build passes (#254)
+- Add ``unumpy.average`` to calculate uncertainties aware average (#38)
 
 3.2.2   2024-July-08
 -----------------------
diff --git a/doc/numpy_guide.rst b/doc/numpy_guide.rst
index f350765f..4242fe73 100644
--- a/doc/numpy_guide.rst
+++ b/doc/numpy_guide.rst
@@ -144,6 +144,22 @@ functions is available in the documentation for  :mod:`uncertainties.umath`.
 .. index::
    pair: testing and operations (in arrays); NaN
 
+Uncertainties aware average
+---------------------------
+
+If you have measured a certain value multiple times, with a different uncertainty every measurement. Averaging over the results in a manner aware of the different uncertainties, is not trivial. The function ``unumpy.average()`` does that:
+
+>>> measurements = numpy.array([2.1, 2.0, 2.05, 2.08, 2.02])
+>>> stds = numpy.array([0.05, 0.03, 0.04, 0.06, 0.05])
+>>> arr = unumpy.uarray(measurements, stds)
+>>> unumpy.average(arr)
+2.03606+/-0.00018
+
+Note how that function gives a value different from numpy's ``mean`` function:
+
+>>> numpy.mean(arr)
+2.050+/-0.021
+
 NaN testing and NaN-aware operations
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
diff --git a/tests/test_unumpy.py b/tests/test_unumpy.py
index 783fd336..5a4d1664 100644
--- a/tests/test_unumpy.py
+++ b/tests/test_unumpy.py
@@ -1,3 +1,5 @@
+import pytest
+
 try:
     import numpy
 except ImportError:
@@ -300,3 +302,9 @@ def test_array_comparisons():
     # For matrices, 1D arrays are converted to 2D arrays:
     mat = unumpy.umatrix([1, 2], [1, 4])
     assert numpy.all((mat == [mat[0, 0], 4]) == [True, False])
+
+
+def test_average_type_check():
+    arr = numpy.array(["bla"])
+    with pytest.raises(ValueError):
+        unumpy.average(arr)
diff --git a/uncertainties/unumpy/core.py b/uncertainties/unumpy/core.py
index 763f57d9..3b337393 100644
--- a/uncertainties/unumpy/core.py
+++ b/uncertainties/unumpy/core.py
@@ -30,10 +30,30 @@
     # Utilities:
     "nominal_values",
     "std_devs",
+    "average",
     # Classes:
     "matrix",
 ]
 
+
+def average(arr, axis=None):
+    """
+    Return a weighted averaged along with a weighted mean over a certain axis.
+
+    By default, operates on all axes of the given array.
+    """
+    if arr.dtype != numpy.dtype("object") or (
+        not isinstance(arr.flat[0], uncert_core.Variable)
+    ):
+        raise ValueError(
+            "unumpy.average is meant to operate only upon numpy arrays of ufloats."
+        )
+    weights = std_devs(arr) ** -2
+    weights_sum = weights.sum(axis=axis)
+    weighted_mean = (nominal_values(arr) * weights).sum(axis=axis) / weights_sum
+    return uarray(weighted_mean, weights_sum**-1 / 2)
+
+
 ###############################################################################
 # Utilities: