BUG: fix `ma.average` not working well with nan weights

If a weight is nan it behaves erratically. See the tests added.

2 files changed, 27 insertions, 1 deletions

diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
index 38bf1f0e8..56da5a7e5 100644
--- a/numpy/ma/extras.py
+++ b/numpy/ma/extras.py
@@ -592,7 +592,7 @@ def average(a, axis=None, weights=None, returned=False):
         avg = a.mean(axis)
         scl = avg.dtype.type(a.count(axis))
     else:
-        wgt = np.asanyarray(weights)
+        wgt = asarray(weights)
 
         if issubclass(a.dtype.type, (np.integer, np.bool_)):
             result_dtype = np.result_type(a.dtype, wgt.dtype, 'f8')
@@ -618,6 +618,7 @@ def average(a, axis=None, weights=None, returned=False):
 
         if m is not nomask:
             wgt = wgt*(~a.mask)
+            wgt.mask |= a.mask
 
         scl = wgt.sum(axis=axis, dtype=result_dtype)
         avg = np.multiply(a, wgt, dtype=result_dtype).sum(axis)/scl
diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py
index e735b9bc7..d30dfd92f 100644
--- a/numpy/ma/tests/test_extras.py
+++ b/numpy/ma/tests/test_extras.py
@@ -309,6 +309,31 @@ class TestAverage:
         expected_masked = np.array([6.0, 5.576271186440678, 6.576271186440678])
         assert_almost_equal(avg_masked, expected_masked)
 
+        # weights should be masked if needed
+        # depending on the array mask. This is to avoid summing
+        # masked nan or other values that are not cancelled by a zero
+        a = np.ma.array([1.0,   2.0,   3.0,  4.0],
+                   mask=[False, False, True, True])
+        avg_unmasked = average(a, weights=[1, 1, 1, np.nan])
+
+        assert_almost_equal(avg_unmasked, 1.5)
+
+        a = np.ma.array([
+            [1.0, 2.0, 3.0, 4.0],
+            [5.0, 6.0, 7.0, 8.0],
+            [9.0, 1.0, 2.0, 3.0],
+        ], mask=[
+            [False, True, True, False],
+            [True, False, True, True],
+            [True, False, True, False],
+        ])
+
+        avg_masked = np.ma.average(a, weights=[1, np.nan, 1], axis=0)
+        avg_expected = np.ma.array([1.0, np.nan, np.nan, 3.5],
+                              mask=[False, True, True, False])
+
+        assert_almost_equal(avg_masked, avg_expected)
+        assert_equal(avg_masked.mask, avg_expected.mask)
 
 class TestConcatenator:
     # Tests for mr_, the equivalent of r_ for masked arrays.