在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是要遍历所有的轴.除了轴与last()轴交换外,m_swapped它是相同的.maxis=1axis=-1

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]