Lặp đi lặp lại trên một mảng numpy

Có một sự thay thế ít dài dòng hơn cho điều này:

for x in xrange(array.shape[0]):
    for y in xrange(array.shape[1]):
        do_stuff(x, y)

Tôi nghĩ ra cái này:

for x, y in itertools.product(map(xrange, array.shape)):
    do_stuff(x, y)

Mà lưu một vết lõm, nhưng vẫn còn khá xấu xí.

Tôi hy vọng cho một cái gì đó trông giống như mã giả này:

for x, y in array.indices:
    do_stuff(x, y)

Có bất cứ điều gì như thế tồn tại?

Tôi nghĩ rằng bạn đang tìm kiếm ndenum Cả .

>>> a =numpy.array([[1,2],[3,4],[5,6]])
>>> for (x,y), value in numpy.ndenumerate(a):
...  print x,y
... 
0 0
0 1
1 0
1 1
2 0
2 1

Về hiệu suất. Nó chậm hơn một chút so với việc hiểu danh sách.

X = np.zeros((100, 100, 100))

%timeit list([((i,j,k), X[i,j,k]) for i in range(X.shape[0]) for j in range(X.shape[1]) for k in range(X.shape[2])])
1 loop, best of 3: 376 ms per loop

%timeit list(np.ndenumerate(X))
1 loop, best of 3: 570 ms per loop

Nếu bạn lo lắng về hiệu suất, bạn có thể tối ưu hóa thêm một chút bằng cách xem xét việc thực hiện ndenumerate, trong đó có 2 điều, chuyển đổi thành một mảng và lặp. Nếu bạn biết bạn có một mảng, bạn có thể gọi .coordsthuộc tính của trình lặp phẳng.

a = X.flat
%timeit list([(a.coords, x) for x in a.flat])
1 loop, best of 3: 305 ms per loop

Nếu bạn chỉ cần các chỉ số, bạn có thể thử numpy.ndindex:

>>> a = numpy.arange(9).reshape(3, 3)
>>> [(x, y) for x, y in numpy.ndindex(a.shape)]
[(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2), (2, 0), (2, 1), (2, 2)]

xem nditer

import numpy as np
Y = np.array([3,4,5,6])
for y in np.nditer(Y, op_flags=['readwrite']):
    y += 3

Y == np.array([6, 7, 8, 9])

y = 3sẽ không hoạt động, sử dụng y *= 0và y += 3thay vào đó.