iterating trên một mảng NumPy

Question

Có một lựa chọn ít tiết như sau:

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

tôi đến với điều này:

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

nào tiết kiệm một thụt đầu dòng, nhưng vẫn còn khá xấu xí.

Tôi hy vọng kiếm cái gì đó trông giống như 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?

Answer 1

Tôi nghĩ bạn đang tìm kiếm ndenumerate.

>>> 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 đang lo lắng về hiệu suất mà bạn có thể tối ưu hóa một chút nữa bằng cách nhìn vào thực hiện ndenumerate, mà làm 2 điều, chuyển sang một mảng và vòng lặp. Nếu bạn biết bạn có một mảng, bạn có thể gọi thuộc tính .coords của trình lặp vòng phẳng.

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

Answer 2

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)] 

Answer 3

thấy 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 = 3 sẽ không hoạt động, sử dụng y *= 0 và y += 3 để thay thế.