迭代numpy数组的第一个d轴

Question

我给了一个具有任意数量轴的数组,我想迭代,说出它们的第一个'd'.我该怎么做呢？

最初我以为我会创建一个包含我想要遍历的所有索引的数组,使用

i = np.indices(a.shape[:d])
indices = np.transpose(np.asarray([x.flatten() for x in i]))
for idx in indices:
    a[idx]

但显然我不能像这样索引数组,即使用另一个包含索引的数组.

Answer 1

你可以使用ndindex:

d = 2
a = np.random.random((2,3,4))
for i in np.ndindex(a.shape[:d]):
    print i, a[i]

输出:

(0, 0) [ 0.72730488  0.2349532   0.36569509  0.31244037]
(0, 1) [ 0.41738425  0.95999499  0.63935274  0.9403284 ]
(0, 2) [ 0.90690468  0.03741634  0.33483221  0.61093582]
(1, 0) [ 0.06716122  0.52632369  0.34441657  0.80678942]
(1, 1) [ 0.8612884   0.22792671  0.15628046  0.63269415]
(1, 2) [ 0.17770685  0.47955698  0.69038541  0.04838387]