Python & Numpy：使用 multi_index 迭代特定轴？答案

【问题标题】：Python & Numpy: Iterating over specific axes with multi_index?Python & Numpy：使用 multi_index 迭代特定轴？
【发布时间】：2014-01-16 05:56:03
【问题描述】：

我有一个包含五个轴的数组：

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

我想使用multi_index 对其进行迭代，就像在this link 的第二个示例中一样，但我不想迭代整个5 个维度，而是要迭代前三个维度。一种 Pythonic 的方式（不涉及 Numpy）将是这样的：

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

如何在纯 Numpy 中实现这一点？

【问题讨论】：

标签： python loops numpy axes multi-index

【解决方案1】：

我们numpy.ndindex():

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

【讨论】：

【解决方案2】：

据我了解，您主要想要的是 itertools.product() 的 numpy 替代品。 numpy 中最接近的类似物是 numpy.indices()。如果我们稍微修改问题中的代码示例，以便向我们展示纯粹在 numpy 中工作时需要能够重现的输出：

indexes = itertools.product(range(3), repeat=3)
for coordinates in indexes:
    print(coordinates)

我们得到以下结果：

(0, 0, 0)
(0, 0, 1)
(0, 0, 2)
(0, 1, 0)
(0, 1, 1)
(0, 1, 2)
(0, 2, 0)
(0, 2, 1)
(0, 2, 2)
(1, 0, 0)
(1, 0, 1)
(1, 0, 2)
(1, 1, 0)
(1, 1, 1)
(1, 1, 2)
(1, 2, 0)
(1, 2, 1)
(1, 2, 2)
(2, 0, 0)
(2, 0, 1)
(2, 0, 2)
(2, 1, 0)
(2, 1, 1)
(2, 1, 2)
(2, 2, 0)
(2, 2, 1)
(2, 2, 2)

以下代码示例将使用 numpy.indices() 而不是 itertools.product() 准确地逐行重现此结果：

import numpy

a, b, c = numpy.indices((3,3,3))
indexes = numpy.transpose(numpy.asarray([a.flatten(), b.flatten(), c.flatten()]))
for coordinates in indexes:
    print(tuple(coordinates))

【讨论】：