Merge pull request numpy#21485 from WarrenWeckesser/average-keepdims

ENH: Add 'keepdims' to 'average()' and 'ma.average()'.

Author: seberg

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`.

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:

numpy/lib/tests/test_function_base.py

@@ -305,6 +305,29 @@ def test_basic(self):
         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 @@ def test_no_trim(self):
         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 @@ def test_frompyfunc_leaks(self, name, incr):
         finally:
             gc.enable()
 
+
 class TestDigitize:
 
     def test_forward(self):
@@ -2339,6 +2363,7 @@ def test_complex(self):
         with pytest.raises(TypeError, match="i0 not supported for complex values"):
             res = i0(a)
 
+
 class TestKaiser:
 
     def test_simple(self):
@@ -3474,6 +3499,7 @@ def test_quantile_scalar_nan(self):
         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),

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 @@ def apply_over_axes(func, a, axes):
     """
 
 
-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 @@ def __init__(self):
 
 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.

numpy/ma/tests/test_extras.py

@@ -75,7 +75,7 @@ def test_masked_all_with_object_nested(self):
         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 @@ def test_complex(self):
         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 @@ def test_masked_weights(self):
         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 @@ def test_shape_scalar(self):
             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)