Python& Numpy:使用multi_index迭代特定的轴?

时间:2013-12-29 00:56:16

标签: python loops numpy axes multi-index

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

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中实现它?

2 个答案:

答案 0 :(得分:3)

我们numpy.ndindex()

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

答案 1 :(得分:1)

据我所知,你主要想要的是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))