diff options
| -rw-r--r-- | numpy/lib/function_base.py | 2 | ||||
| -rw-r--r-- | numpy/lib/tests/test_function_base.py | 12 | ||||
| -rw-r--r-- | numpy/ma/extras.py | 2 | ||||
| -rw-r--r-- | numpy/ma/tests/test_extras.py | 9 |
4 files changed, 23 insertions, 2 deletions
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index 843e1b85a..d90c23bfe 100644 --- a/numpy/lib/function_base.py +++ b/numpy/lib/function_base.py @@ -542,7 +542,7 @@ def average(a, axis=None, weights=None, returned=False, *, wgt = np.broadcast_to(wgt, (a.ndim-1)*(1,) + wgt.shape) wgt = wgt.swapaxes(-1, axis) - scl = wgt.sum(axis=axis, dtype=result_dtype) + scl = wgt.sum(axis=axis, dtype=result_dtype, **keepdims_kw) if np.any(scl == 0.0): raise ZeroDivisionError( "Weights sum to zero, can't be normalized") diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py index 64318255b..9cc18a5e4 100644 --- a/numpy/lib/tests/test_function_base.py +++ b/numpy/lib/tests/test_function_base.py @@ -360,6 +360,18 @@ class TestAverage: assert_(np.average(y3, weights=w3).dtype == np.result_type(y3, w3)) + # test weights with `keepdims=False` and `keepdims=True` + x = np.array([2, 3, 4]).reshape(3, 1) + w = np.array([4, 5, 6]).reshape(3, 1) + + actual = np.average(x, weights=w, axis=1, keepdims=False) + desired = np.array([2., 3., 4.]) + assert_array_equal(actual, desired) + + actual = np.average(x, weights=w, axis=1, keepdims=True) + desired = np.array([[2.], [3.], [4.]]) + assert_array_equal(actual, desired) + def test_returned(self): y = np.array([[1, 2, 3], [4, 5, 6]]) diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py index d90831b9b..d2986012b 100644 --- a/numpy/ma/extras.py +++ b/numpy/ma/extras.py @@ -645,7 +645,7 @@ def average(a, axis=None, weights=None, returned=False, *, wgt = wgt*(~a.mask) wgt.mask |= a.mask - scl = wgt.sum(axis=axis, dtype=result_dtype) + scl = wgt.sum(axis=axis, dtype=result_dtype, **keepdims_kw) avg = np.multiply(a, wgt, dtype=result_dtype).sum(axis, **keepdims_kw) / scl diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py index 1827edd1f..04bf8cfc2 100644 --- a/numpy/ma/tests/test_extras.py +++ b/numpy/ma/tests/test_extras.py @@ -241,6 +241,15 @@ class TestAverage: a2dma = average(a2dm, axis=1) assert_equal(a2dma, [1.5, 4.0]) + def test_testAverage4(self): + # Test that `keepdims` works with average + x = np.array([2, 3, 4]).reshape(3, 1) + b = np.ma.array(x, mask=[[False], [False], [True]]) + w = np.array([4, 5, 6]).reshape(3, 1) + actual = average(b, weights=w, axis=1, keepdims=True) + desired = masked_array([[2.], [3.], [4.]], [[False], [False], [True]]) + assert_equal(actual, desired) + def test_onintegers_with_mask(self): # Test average on integers with mask a = average(array([1, 2])) |