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]

在numpy中迭代任意维度

2 个答案: