在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 成为您希望循环的轴。 m_swapped 与 m 相同，除了 axis=1 轴与最后一个 (axis=-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]