diff options
Diffstat (limited to 'numpy/ma/extras.py')
| -rw-r--r-- | numpy/ma/extras.py | 119 |
1 files changed, 110 insertions, 9 deletions
diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py index 048d94bb7..641f4746f 100644 --- a/numpy/ma/extras.py +++ b/numpy/ma/extras.py @@ -10,12 +10,12 @@ A collection of utilities for `numpy.ma`. """ __all__ = [ 'apply_along_axis', 'apply_over_axes', 'atleast_1d', 'atleast_2d', - 'atleast_3d', 'average', 'clump_masked', 'clump_unmasked', - 'column_stack', 'compress_cols', 'compress_nd', 'compress_rowcols', - 'compress_rows', 'count_masked', 'corrcoef', 'cov', 'diagflat', 'dot', - 'dstack', 'ediff1d', 'flatnotmasked_contiguous', 'flatnotmasked_edges', - 'hsplit', 'hstack', 'isin', 'in1d', 'intersect1d', 'mask_cols', 'mask_rowcols', - 'mask_rows', 'masked_all', 'masked_all_like', 'median', 'mr_', + 'atleast_3d', 'average', 'clump_masked', 'clump_unmasked', 'column_stack', + 'compress_cols', 'compress_nd', 'compress_rowcols', 'compress_rows', + 'count_masked', 'corrcoef', 'cov', 'diagflat', 'dot', 'dstack', 'ediff1d', + 'flatnotmasked_contiguous', 'flatnotmasked_edges', 'hsplit', 'hstack', + 'isin', 'in1d', 'intersect1d', 'mask_cols', 'mask_rowcols', 'mask_rows', + 'masked_all', 'masked_all_like', 'median', 'mr_', 'ndenumerate', 'notmasked_contiguous', 'notmasked_edges', 'polyfit', 'row_stack', 'setdiff1d', 'setxor1d', 'stack', 'unique', 'union1d', 'vander', 'vstack', ] @@ -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,10 +1547,79 @@ class mr_class(MAxisConcatenator): mr_ = mr_class() + #####-------------------------------------------------------------------------- #---- Find unmasked data --- #####-------------------------------------------------------------------------- +def ndenumerate(a, compressed=True): + """ + Multidimensional index iterator. + + Return an iterator yielding pairs of array coordinates and values, + skipping elements that are masked. With `compressed=False`, + `ma.masked` is yielded as the value of masked elements. This + behavior differs from that of `numpy.ndenumerate`, which yields the + value of the underlying data array. + + Notes + ----- + .. versionadded:: 1.23.0 + + Parameters + ---------- + a : array_like + An array with (possibly) masked elements. + compressed : bool, optional + If True (default), masked elements are skipped. + + See Also + -------- + numpy.ndenumerate : Equivalent function ignoring any mask. + + Examples + -------- + >>> a = np.ma.arange(9).reshape((3, 3)) + >>> a[1, 0] = np.ma.masked + >>> a[1, 2] = np.ma.masked + >>> a[2, 1] = np.ma.masked + >>> a + masked_array( + data=[[0, 1, 2], + [--, 4, --], + [6, --, 8]], + mask=[[False, False, False], + [ True, False, True], + [False, True, False]], + fill_value=999999) + >>> for index, x in np.ma.ndenumerate(a): + ... print(index, x) + (0, 0) 0 + (0, 1) 1 + (0, 2) 2 + (1, 1) 4 + (2, 0) 6 + (2, 2) 8 + + >>> for index, x in np.ma.ndenumerate(a, compressed=False): + ... print(index, x) + (0, 0) 0 + (0, 1) 1 + (0, 2) 2 + (1, 0) -- + (1, 1) 4 + (1, 2) -- + (2, 0) 6 + (2, 1) -- + (2, 2) 8 + """ + for it, mask in zip(np.ndenumerate(a), getmaskarray(a).flat): + if not mask: + yield it + elif not compressed: + yield it[0], masked + + def flatnotmasked_edges(a): """ Find the indices of the first and last unmasked values. @@ -1682,6 +1782,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. |