diff options
| author | Srimukh Sripada <git@srimukh.com> | 2022-06-27 05:06:00 +0200 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2022-06-26 23:06:00 -0400 |
| commit | e67fe9f1a86013dbfd4a5ad2bd69aa6e7f60e1fb (patch) | |
| tree | c63192e3b1c5eeede2bd042f4531f44568256d5e /numpy/ma/tests/test_extras.py | |
| parent | b65f0b7b8ba7e80b65773e06aae22a8369678868 (diff) | |
| download | numpy-e67fe9f1a86013dbfd4a5ad2bd69aa6e7f60e1fb.tar.gz | |
BUG: Use `keepdims` during normalization in `np.average` and `np.ma.average` (#21851)
The keepdims flag needs to be applied during the calculation of the sum of
the weights in np.average and np.ma.average. Not passing it causes weights
to broadcast incorrectly.
Fixes #21850
Diffstat (limited to 'numpy/ma/tests/test_extras.py')
| -rw-r--r-- | numpy/ma/tests/test_extras.py | 9 |
1 files changed, 9 insertions, 0 deletions
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])) |