在numpy中迭代任意维度

Question

我有一个多维的numpy数组，我需要迭代给定的维度。 问题是，直到运行时我才知道哪个维度。 换句话说，给定一个数组m，我可能想要

m[:,:,:,i] for i in xrange(n)

或者我想要的

m[:,:,i,:] for i in xrange(n)

等等

我想在numpy中必须有一个简单的功能来写这个，但我无法弄清楚它是什么/它可能被称为什么。 有什么想法吗？

Answer 1

有很多方法可以做到这一点。 您可以使用切片列表构建正确的索引，或者可能改变m的步幅。 但是，最简单的方法可能是使用np.swapaxes ：

import numpy as np
m=np.arange(24).reshape(2,3,4)
print(m.shape)
# (2, 3, 4)

让axis成为您想要循环的轴。 除了axis=1轴与最后一个（ axis=-1 ）轴交换外， m_swapped与m相同。

axis=1
m_swapped=m.swapaxes(axis,-1)
print(m_swapped.shape)
# (2, 4, 3)

现在你可以遍历最后一个轴：

for i in xrange(m_swapped.shape[-1]):
    assert np.all(m[:,i,:] == m_swapped[...,i])

请注意， m_swapped是m的视图，而不是副本。 改变m_swapped将改变m 。

m_swapped[1,2,0]=100
print(m)
assert(m[1,0,2]==100)

Answer 2

您可以使用slice(None)代替: 。 例如，

from numpy import *

d = 2  # the dimension to iterate

x = arange(5*5*5).reshape((5,5,5))
s = slice(None)  # :

for i in range(5):
    slicer = [s]*3  # [:, :, :]
    slicer[d] = i   # [:, :, i]
    print x[slicer] # x[:, :, i]