Merge pull request numpy#20020 from joukewitteveen/ma-ndenumerate

ENH: add ndenumerate specialization for masked arrays

Author: mattip

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.

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

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,

numpy/ma/extras.py

@@ -10,12 +10,12 @@
 """
 __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 @@ def __init__(self):
 #---- 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.

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): ...

numpy/ma/tests/test_extras.py

@@ -28,7 +28,7 @@
     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 @@ def test_shape_scalar(self):
         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):