numpy: ndenumerate for masked arrays?

Question

Is there a way to enumerate over the non-masked locations of a masked numpy ndarray (e.g. in the way that ndenumerate does it for regular ndarrays, but omitting all the masked entries)?

EDIT: to be more precise: the enumeration should not only skip over the masked entries, but also show the indices of the non-masked ones in the original array. E.g. if the first five elements of a 1-d array are masked, and the next one has an unmasked value of 3, then the enumeration should start with something like ((5,), 3), ....

Thanks!

PS: note that, although it is possible to apply ndenumerate to a masked ndarray, the resulting enumeration does not discriminate between its masked and normal entries. In fact, ndenumerate not only does not filter out the masked entries from the enumeration, but it doesn't even replace the enumerated values with the masked constant. Therefore, one can't adapt ndenumerate for this task by just wrapping ndenumerate with a suitable filter.

Answer 1

You can access only valid entries using inverse of a mask as an index:

>>> import numpy as np
>>> import numpy.ma as ma
>>> x = np.array([11, 22, -1, 44])
>>> m_arr = ma.masked_array(x, mask=[0, 0, 1, 0])
>>> for index, i in np.ndenumerate(m_arr[~m_arr.mask]): 
        print index, i
(0,) 11
(1,) 22
(2,) 44

See this for details.

The enumeration over only valid entries with indices from the original array:

>>> for (index, val), m in zip(np.ndenumerate(m_arr), m_arr.mask):
      if not m:
        print index, val 
(0,) 11
(1,) 22
(3,) 44

Answer 2

How about:

import numpy as np
import itertools

def maenumerate(marr):
    mask = ~marr.mask.ravel()
    for i, m in itertools.izip(np.ndenumerate(marr), mask):
        if m: yield i

N = 12
a = np.arange(N).reshape(2, 2, 3)+10

b = np.ma.array(a, mask = (a%5 == 0))
for i, val in maenumerate(b):
    print i, val

which yields

(0, 0, 1) 11
(0, 0, 2) 12
(0, 1, 0) 13
(0, 1, 1) 14
(1, 0, 0) 16
(1, 0, 1) 17
(1, 0, 2) 18
(1, 1, 0) 19
(1, 1, 2) 21