diff options
| -rw-r--r-- | doc/release/upcoming_changes/21485.new_feature.rst | 5 | ||||
| -rw-r--r-- | numpy/lib/function_base.py | 34 | ||||
| -rw-r--r-- | numpy/lib/tests/test_function_base.py | 28 | ||||
| -rw-r--r-- | numpy/ma/extras.py | 39 | ||||
| -rw-r--r-- | numpy/ma/tests/test_extras.py | 27 |
5 files changed, 122 insertions, 11 deletions
diff --git a/doc/release/upcoming_changes/21485.new_feature.rst b/doc/release/upcoming_changes/21485.new_feature.rst new file mode 100644 index 000000000..99fd5e92d --- /dev/null +++ b/doc/release/upcoming_changes/21485.new_feature.rst @@ -0,0 +1,5 @@ +``keepdims`` parameter for ``average`` +-------------------------------------- +The parameter ``keepdims`` was added to the functions `numpy.average` +and `numpy.ma.average`. The parameter has the same meaning as it +does in reduction functions such as `numpy.sum` or `numpy.mean`. diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index d611dd225..b8ae9a470 100644 --- a/numpy/lib/function_base.py +++ b/numpy/lib/function_base.py @@ -388,12 +388,14 @@ def iterable(y): return True -def _average_dispatcher(a, axis=None, weights=None, returned=None): +def _average_dispatcher(a, axis=None, weights=None, returned=None, *, + keepdims=None): return (a, weights) @array_function_dispatch(_average_dispatcher) -def average(a, axis=None, weights=None, returned=False): +def average(a, axis=None, weights=None, returned=False, *, + keepdims=np._NoValue): """ Compute the weighted average along the specified axis. @@ -428,6 +430,14 @@ def average(a, axis=None, weights=None, returned=False): is returned, otherwise only the average is returned. If `weights=None`, `sum_of_weights` is equivalent to the number of elements over which the average is taken. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the original `a`. + *Note:* `keepdims` will not work with instances of `numpy.matrix` + or other classes whose methods do not support `keepdims`. + + .. versionadded:: 1.23.0 Returns ------- @@ -471,7 +481,7 @@ def average(a, axis=None, weights=None, returned=False): >>> np.average(np.arange(1, 11), weights=np.arange(10, 0, -1)) 4.0 - >>> data = np.arange(6).reshape((3,2)) + >>> data = np.arange(6).reshape((3, 2)) >>> data array([[0, 1], [2, 3], @@ -488,11 +498,24 @@ def average(a, axis=None, weights=None, returned=False): >>> avg = np.average(a, weights=w) >>> print(avg.dtype) complex256 + + With ``keepdims=True``, the following result has shape (3, 1). + + >>> np.average(data, axis=1, keepdims=True) + array([[0.5], + [2.5], + [4.5]]) """ a = np.asanyarray(a) + if keepdims is np._NoValue: + # Don't pass on the keepdims argument if one wasn't given. + keepdims_kw = {} + else: + keepdims_kw = {'keepdims': keepdims} + if weights is None: - avg = a.mean(axis) + avg = a.mean(axis, **keepdims_kw) scl = avg.dtype.type(a.size/avg.size) else: wgt = np.asanyarray(weights) @@ -524,7 +547,8 @@ def average(a, axis=None, weights=None, returned=False): raise ZeroDivisionError( "Weights sum to zero, can't be normalized") - avg = np.multiply(a, wgt, dtype=result_dtype).sum(axis)/scl + avg = np.multiply(a, wgt, + dtype=result_dtype).sum(axis, **keepdims_kw) / scl if returned: if scl.shape != avg.shape: diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py index 874754a64..bdcbef91d 100644 --- a/numpy/lib/tests/test_function_base.py +++ b/numpy/lib/tests/test_function_base.py @@ -305,6 +305,29 @@ class TestAverage: assert_almost_equal(y5.mean(0), average(y5, 0)) assert_almost_equal(y5.mean(1), average(y5, 1)) + @pytest.mark.parametrize( + 'x, axis, expected_avg, weights, expected_wavg, expected_wsum', + [([1, 2, 3], None, [2.0], [3, 4, 1], [1.75], [8.0]), + ([[1, 2, 5], [1, 6, 11]], 0, [[1.0, 4.0, 8.0]], + [1, 3], [[1.0, 5.0, 9.5]], [[4, 4, 4]])], + ) + def test_basic_keepdims(self, x, axis, expected_avg, + weights, expected_wavg, expected_wsum): + avg = np.average(x, axis=axis, keepdims=True) + assert avg.shape == np.shape(expected_avg) + assert_array_equal(avg, expected_avg) + + wavg = np.average(x, axis=axis, weights=weights, keepdims=True) + assert wavg.shape == np.shape(expected_wavg) + assert_array_equal(wavg, expected_wavg) + + wavg, wsum = np.average(x, axis=axis, weights=weights, returned=True, + keepdims=True) + assert wavg.shape == np.shape(expected_wavg) + assert_array_equal(wavg, expected_wavg) + assert wsum.shape == np.shape(expected_wsum) + assert_array_equal(wsum, expected_wsum) + def test_weights(self): y = np.arange(10) w = np.arange(10) @@ -1242,11 +1265,11 @@ class TestTrimZeros: res = trim_zeros(arr) assert_array_equal(arr, res) - def test_list_to_list(self): res = trim_zeros(self.a.tolist()) assert isinstance(res, list) + class TestExtins: def test_basic(self): @@ -1759,6 +1782,7 @@ class TestLeaks: finally: gc.enable() + class TestDigitize: def test_forward(self): @@ -2339,6 +2363,7 @@ class Test_I0: with pytest.raises(TypeError, match="i0 not supported for complex values"): res = i0(a) + class TestKaiser: def test_simple(self): @@ -3474,6 +3499,7 @@ class TestQuantile: assert np.isscalar(actual) assert_equal(np.quantile(a, 0.5), np.nan) + class TestLerp: @hypothesis.given(t0=st.floats(allow_nan=False, allow_infinity=False, min_value=0, max_value=1), diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py index 048d94bb7..b72b2d2cb 100644 --- a/numpy/ma/extras.py +++ b/numpy/ma/extras.py @@ -475,6 +475,7 @@ def apply_over_axes(func, a, axes): "an array of the correct shape") return val + if apply_over_axes.__doc__ is not None: apply_over_axes.__doc__ = np.apply_over_axes.__doc__[ :np.apply_over_axes.__doc__.find('Notes')].rstrip() + \ @@ -524,7 +525,8 @@ if apply_over_axes.__doc__ is not None: """ -def average(a, axis=None, weights=None, returned=False): +def average(a, axis=None, weights=None, returned=False, *, + keepdims=np._NoValue): """ Return the weighted average of array over the given axis. @@ -550,6 +552,14 @@ def average(a, axis=None, weights=None, returned=False): Flag indicating whether a tuple ``(result, sum of weights)`` should be returned as output (True), or just the result (False). Default is False. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the original `a`. + *Note:* `keepdims` will not work with instances of `numpy.matrix` + or other classes whose methods do not support `keepdims`. + + .. versionadded:: 1.23.0 Returns ------- @@ -582,14 +592,29 @@ def average(a, axis=None, weights=None, returned=False): mask=[False, False], fill_value=1e+20) + With ``keepdims=True``, the following result has shape (3, 1). + + >>> np.ma.average(x, axis=1, keepdims=True) + masked_array( + data=[[0.5], + [2.5], + [4.5]], + mask=False, + fill_value=1e+20) """ a = asarray(a) m = getmask(a) # inspired by 'average' in numpy/lib/function_base.py + if keepdims is np._NoValue: + # Don't pass on the keepdims argument if one wasn't given. + keepdims_kw = {} + else: + keepdims_kw = {'keepdims': keepdims} + if weights is None: - avg = a.mean(axis) + avg = a.mean(axis, **keepdims_kw) scl = avg.dtype.type(a.count(axis)) else: wgt = asarray(weights) @@ -621,7 +646,8 @@ def average(a, axis=None, weights=None, returned=False): wgt.mask |= a.mask scl = wgt.sum(axis=axis, dtype=result_dtype) - avg = np.multiply(a, wgt, dtype=result_dtype).sum(axis)/scl + avg = np.multiply(a, wgt, + dtype=result_dtype).sum(axis, **keepdims_kw) / scl if returned: if scl.shape != avg.shape: @@ -713,6 +739,7 @@ def median(a, axis=None, out=None, overwrite_input=False, keepdims=False): else: return r + def _median(a, axis=None, out=None, overwrite_input=False): # when an unmasked NaN is present return it, so we need to sort the NaN # values behind the mask @@ -840,6 +867,7 @@ def compress_nd(x, axis=None): data = data[(slice(None),)*ax + (~m.any(axis=axes),)] return data + def compress_rowcols(x, axis=None): """ Suppress the rows and/or columns of a 2-D array that contain @@ -912,6 +940,7 @@ def compress_rows(a): raise NotImplementedError("compress_rows works for 2D arrays only.") return compress_rowcols(a, 0) + def compress_cols(a): """ Suppress whole columns of a 2-D array that contain masked values. @@ -929,6 +958,7 @@ def compress_cols(a): raise NotImplementedError("compress_cols works for 2D arrays only.") return compress_rowcols(a, 1) + def mask_rows(a, axis=np._NoValue): """ Mask rows of a 2D array that contain masked values. @@ -979,6 +1009,7 @@ def mask_rows(a, axis=np._NoValue): "will raise TypeError", DeprecationWarning, stacklevel=2) return mask_rowcols(a, 0) + def mask_cols(a, axis=np._NoValue): """ Mask columns of a 2D array that contain masked values. @@ -1516,6 +1547,7 @@ class mr_class(MAxisConcatenator): mr_ = mr_class() + #####-------------------------------------------------------------------------- #---- Find unmasked data --- #####-------------------------------------------------------------------------- @@ -1682,6 +1714,7 @@ def flatnotmasked_contiguous(a): i += n return result + def notmasked_contiguous(a, axis=None): """ Find contiguous unmasked data in a masked array along the given axis. diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py index d30dfd92f..01a47bef8 100644 --- a/numpy/ma/tests/test_extras.py +++ b/numpy/ma/tests/test_extras.py @@ -75,7 +75,7 @@ class TestGeneric: assert_equal(len(masked_arr['b']['c']), 1) assert_equal(masked_arr['b']['c'].shape, (1, 1)) assert_equal(masked_arr['b']['c']._fill_value.shape, ()) - + def test_masked_all_with_object(self): # same as above except that the array is not nested my_dtype = np.dtype([('b', (object, (1,)))]) @@ -292,6 +292,29 @@ class TestAverage: assert_almost_equal(wav1.real, expected1.real) assert_almost_equal(wav1.imag, expected1.imag) + @pytest.mark.parametrize( + 'x, axis, expected_avg, weights, expected_wavg, expected_wsum', + [([1, 2, 3], None, [2.0], [3, 4, 1], [1.75], [8.0]), + ([[1, 2, 5], [1, 6, 11]], 0, [[1.0, 4.0, 8.0]], + [1, 3], [[1.0, 5.0, 9.5]], [[4, 4, 4]])], + ) + def test_basic_keepdims(self, x, axis, expected_avg, + weights, expected_wavg, expected_wsum): + avg = np.ma.average(x, axis=axis, keepdims=True) + assert avg.shape == np.shape(expected_avg) + assert_array_equal(avg, expected_avg) + + wavg = np.ma.average(x, axis=axis, weights=weights, keepdims=True) + assert wavg.shape == np.shape(expected_wavg) + assert_array_equal(wavg, expected_wavg) + + wavg, wsum = np.ma.average(x, axis=axis, weights=weights, + returned=True, keepdims=True) + assert wavg.shape == np.shape(expected_wavg) + assert_array_equal(wavg, expected_wavg) + assert wsum.shape == np.shape(expected_wsum) + assert_array_equal(wsum, expected_wsum) + def test_masked_weights(self): # Test with masked weights. # (Regression test for https://github.com/numpy/numpy/issues/10438) @@ -335,6 +358,7 @@ class TestAverage: 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. @@ -1642,7 +1666,6 @@ class TestShapeBase: assert_equal(a.mask.shape, a.shape) assert_equal(a.data.shape, a.shape) - b = diagflat(1.0) assert_equal(b.shape, (1, 1)) assert_equal(b.mask.shape, b.data.shape) |