python 和 NumPy : Iterating over specific axes with multi_index?

Question

我有一个有五个轴的数组:

colors = numpy.zeros(3, 3, 3, 6, 3))

我想像 this link 中的第二个示例一样使用 multi_index 对其进行迭代，但我不想遍历整个 5 个维度，而是想遍历前三个维度。一种 Pythonic 的方式(不涉及 Numpy)是这样的:

indexes = itertools.product(range(3), repeat=3)
for coordinates in indexes:
  colors[coordinates]

我如何在纯 Numpy 中实现它？

Answer 1

我们numpy.ndindex():

for idx in np.ndindex(*colors.shape[:3]):
    data = colors[coordinates]