diff options
| -rw-r--r-- | doc/release/upcoming_changes/20020.new_function.rst | 4 | ||||
| -rw-r--r-- | doc/source/reference/routines.ma.rst | 1 | ||||
| -rw-r--r-- | numpy/ma/__init__.pyi | 1 | ||||
| -rw-r--r-- | numpy/ma/extras.py | 80 | ||||
| -rw-r--r-- | numpy/ma/extras.pyi | 1 | ||||
| -rw-r--r-- | numpy/ma/tests/test_extras.py | 40 |
6 files changed, 120 insertions, 7 deletions
diff --git a/doc/release/upcoming_changes/20020.new_function.rst b/doc/release/upcoming_changes/20020.new_function.rst new file mode 100644 index 000000000..0f310ceac --- /dev/null +++ b/doc/release/upcoming_changes/20020.new_function.rst @@ -0,0 +1,4 @@ +`ndenumerate` specialization for masked arrays +---------------------------------------------- +The masked array module now provides the `numpy.ma.ndenumerate` function, +an alternative to `numpy.ndenumerate` that skips masked values by default. diff --git a/doc/source/reference/routines.ma.rst b/doc/source/reference/routines.ma.rst index 5404c43d8..1de5c1c02 100644 --- a/doc/source/reference/routines.ma.rst +++ b/doc/source/reference/routines.ma.rst @@ -190,6 +190,7 @@ Finding masked data .. autosummary:: :toctree: generated/ + ma.ndenumerate ma.flatnotmasked_contiguous ma.flatnotmasked_edges ma.notmasked_contiguous diff --git a/numpy/ma/__init__.pyi b/numpy/ma/__init__.pyi index 04368b6c4..7f5cb56a8 100644 --- a/numpy/ma/__init__.pyi +++ b/numpy/ma/__init__.pyi @@ -216,6 +216,7 @@ from numpy.ma.extras import ( masked_all_like as masked_all_like, median as median, mr_ as mr_, + ndenumerate as ndenumerate, notmasked_contiguous as notmasked_contiguous, notmasked_edges as notmasked_edges, polyfit as polyfit, diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py index b72b2d2cb..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', ] @@ -1552,6 +1552,74 @@ 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. diff --git a/numpy/ma/extras.pyi b/numpy/ma/extras.pyi index 05ad87210..56228b927 100644 --- a/numpy/ma/extras.pyi +++ b/numpy/ma/extras.pyi @@ -74,6 +74,7 @@ class mr_class(MAxisConcatenator): mr_: mr_class +def ndenumerate(a, compressed=...): ... def flatnotmasked_edges(a): ... def notmasked_edges(a, axis=...): ... def flatnotmasked_contiguous(a): ... diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py index 01a47bef8..05403344b 100644 --- a/numpy/ma/tests/test_extras.py +++ b/numpy/ma/tests/test_extras.py @@ -28,7 +28,7 @@ from numpy.ma.extras import ( ediff1d, apply_over_axes, apply_along_axis, compress_nd, compress_rowcols, mask_rowcols, clump_masked, clump_unmasked, flatnotmasked_contiguous, notmasked_contiguous, notmasked_edges, masked_all, masked_all_like, isin, - diagflat, stack, vstack + diagflat, ndenumerate, stack, vstack ) @@ -1671,6 +1671,44 @@ class TestShapeBase: assert_equal(b.mask.shape, b.data.shape) +class TestNDEnumerate: + + def test_ndenumerate_nomasked(self): + ordinary = np.ndarray(6).reshape((1, 3, 2)) + empty_mask = np.zeros_like(ordinary, dtype=bool) + with_mask = masked_array(ordinary, mask=empty_mask) + assert_equal(list(np.ndenumerate(ordinary)), + list(ndenumerate(ordinary))) + assert_equal(list(ndenumerate(ordinary)), + list(ndenumerate(with_mask))) + assert_equal(list(ndenumerate(with_mask)), + list(ndenumerate(with_mask, compressed=False))) + + def test_ndenumerate_allmasked(self): + a = masked_all(()) + b = masked_all((100,)) + c = masked_all((2, 3, 4)) + assert_equal(list(ndenumerate(a)), []) + assert_equal(list(ndenumerate(b)), []) + assert_equal(list(ndenumerate(b, compressed=False)), + list(zip(np.ndindex((100,)), 100 * [masked]))) + assert_equal(list(ndenumerate(c)), []) + assert_equal(list(ndenumerate(c, compressed=False)), + list(zip(np.ndindex((2, 3, 4)), 2 * 3 * 4 * [masked]))) + + def test_ndenumerate_mixedmasked(self): + a = masked_array(np.arange(12).reshape((3, 4)), + mask=[[1, 1, 1, 1], + [1, 1, 0, 1], + [0, 0, 0, 0]]) + items = [((1, 2), 6), + ((2, 0), 8), ((2, 1), 9), ((2, 2), 10), ((2, 3), 11)] + assert_equal(list(ndenumerate(a)), items) + assert_equal(len(list(ndenumerate(a, compressed=False))), a.size) + for coordinate, value in ndenumerate(a, compressed=False): + assert_equal(a[coordinate], value) + + class TestStack: def test_stack_1d(self): |