2

There are many ways to iterate over a 2D array. Examples below.

Most of these approaches, while functional, strike me as difficult to read. Some are very memory-expensive, and most do not generalize well to N-dimensions.

Is there a more readable way, preferably one that is computationally efficient and generalizes to ND nicely, to iterate over all coordinate sets in an ndarray?

arr = np.arange(100).reshape([10,10])
x,y = np.indices(arr.shape)

for i,j in zip(x.flat,y.flat):
    dosomething(arr[i,j])

for i,j in np.nditer(np.indices(arr.shape).tolist()):
    dosomething(arr[i,j])

for i in xrange(arr.shape[0]):
    for j in xrange(arr.shape[1]):
        dosomething(arr[i,j])

for i,j in itertools.product(range(arr.shape[0], range.shape[1])):
    dosomething(arr[i,j])

# on further thought, maybe this one is OK?
for ind in xrange(arr.size):
    i,j = np.unravel_index(ind, arr.shape)
    dosomething(arr[i,j])

for i,j in itertools.product(*map(xrange, arr.shape)):
    dosomething(arr[i,j])

(the latter from Pythonic way of iterating over 3D array)

The question I really wanted an answer to was "how do I get the x,y indices of an array?" The answer is:

for i,j in (np.unravel_index(ind,arr.shape) for ind in xrange(arr.size)):
    dosomething(arr[i,j])

(np.unravel_index(ind,arr.shape) for ind in xrange(arr.size)) is a fairly readable and efficient generator.

But, for the question asked in the title, the other (linked) answers are better (np.nditer, np.enumerate)

Community
  • 1
  • 1
keflavich
  • 18,278
  • 20
  • 86
  • 118
  • After posting, I noticed this... http://stackoverflow.com/questions/6967463/iterating-over-a-numpy-array, which may be the answer. – keflavich Apr 20 '13 at 23:40
  • I was just writing up `ndenumerate` as a solution; it's what I do, on those rare occasions I need this. – DSM Apr 20 '13 at 23:41
  • Go ahead and put it up... considering I googled this question for about 10 minutes without finding an answer, I think the search terms for this style of question are not yet saturated. – keflavich Apr 20 '13 at 23:54
  • We can just close it as a duplicate. Your question will still be indexed so it can help to point others in the right direction. – DSM Apr 20 '13 at 23:57
  • That's the right approach. – keflavich Apr 20 '13 at 23:58
  • Do you even need the indices? For example nditer you could just as well plug in the array, and that is that (if you do not care about order, and only need views to the elements). – seberg Apr 21 '13 at 10:30

1 Answers1

1

From your examples it seems that you want to iterate your arrays in C-contiguous style. If you only care about the elements, and not the indices, the following can be used:

for e in arr.flat:
    dosomething(e)

Of course, it is faster then all alternatives. However, if you want to some funky things with indices, this method cannot be used (use ndenumerate -- your linked SO answer -- for that). This can be used with N-dims.

mg007
  • 2,888
  • 24
  • 29