遍历 ndarray 的切片答案

【问题标题】：iterate over slices of an ndarray遍历 ndarray 的切片
【发布时间】：2013-06-25 14:54:07
【问题描述】：

假设我有一个 3D numpy.array，例如尺寸为 x y z，有没有办法沿特定轴迭代切片？比如：

for layer in data.slices(dim=2):
    # do something with layer

编辑：澄清一下，这个例子是一个 dim=3 数组，即 shape=(len_x, len_y, len_z)。 Elazar 和等效 kamjagin 的解决方案有效，但不是那么通用 - 您必须手动构造 [:, :, i]，这意味着您需要知道维度，并且代码不够通用，无法处理任意维度的数组。您可以使用 [..., :] 之类的东西来填充缺失的维度，但您仍然需要自己构建它。

对不起，应该更清楚一点，例子有点太简单了！

【问题讨论】：

查看stackoverflow.com/questions/1589706/…。 swapaxes 方法是最快的，但最不清晰。 kamjagin 的方法可以通过构造一个元组传递到括号中来概括（即data[tuple(slice(None), slice(None), i)] 与data[:,:,i] 相同）。
@AFoglia 我不认为这个问题的时间安排非常相关。使用swapaxes 或rollaxis 将花费更多时间来设置循环，但实际迭代更快，请参阅添加到我的答案中的时间。在您非常小的示例中，设置在实际迭代中占主导地位。我不同意可读性，但我可能太习惯于 numpy 而没有注意到它。
@AFoglia numpy 有一个巧妙的技巧，可以将切片（元组）与它们要索引的对象分开构建：numpy.s_[:,:,i] 相当于 tuple(slice(None), slice(None), i)。
（请注意numpy.s_ 有一个属性maketuple，它决定使用单个切片是否会产生一个元组[即slice(start, end, step) 与(slice(start, end, step),)] . 默认为False。)
@JAB 是的，但是如果他想要针对任意秩数组的任意维度的通用解决方案，使用(slice(None),)*(idim)+(idx,) 以编程方式创建一个元组比使用: 字符更容易。你不能这样做data[(:,)*(idim)+(idx,)]。

标签： python numpy slice

【解决方案1】：

迭代第一个维度非常容易，见下文。要迭代其他维度，将该维度滚动到前面并执行相同操作：

>>> data = np.arange(24).reshape(2, 3, 4)
>>> for dim_0_slice in data: # the first dimension is easy
...     print dim_0_slice
... 
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
[[12 13 14 15]
 [16 17 18 19]
 [20 21 22 23]]
>>> for dim_1_slice in np.rollaxis(data, 1): # for the others, roll it to the front
...     print dim_1_slice
... 
[[ 0  1  2  3]
 [12 13 14 15]]
[[ 4  5  6  7]
 [16 17 18 19]]
[[ 8  9 10 11]
 [20 21 22 23]]
>>> for dim_2_slice in np.rollaxis(data, 2):
...     print dim_2_slice
... 
[[ 0  4  8]
 [12 16 20]]
[[ 1  5  9]
 [13 17 21]]
[[ 2  6 10]
 [14 18 22]]
[[ 3  7 11]
 [15 19 23]]

编辑一些时间，比较大型数组的不同方法：

In [7]: a = np.arange(200*100*300).reshape(200, 100, 300)

In [8]: %timeit for j in xrange(100): a[:, j]
10000 loops, best of 3: 60.2 us per loop

In [9]: %timeit for j in xrange(100): a[:, j, :]
10000 loops, best of 3: 82.8 us per loop

In [10]: %timeit for j in np.rollaxis(a, 1): j
10000 loops, best of 3: 28.2 us per loop

In [11]: %timeit for j in np.swapaxes(a, 0, 1): j
10000 loops, best of 3: 26.7 us per loop

【讨论】：

哦，太好了，我不知道rollaxis()。
这里也一样，确实很不错！
从 Numpy 1.11 开始改用np.moveaxis( data, k, 0 )。

【解决方案2】：

这可能比这更优雅地解决，但如果您事先知道 dim（例如 2），一种方法是：

for i in range(data.shape[dim]):
    layer = data[:,:,i]

或者如果 dim=0

for i in range(data.shape[dim]):
    layer = data[i,:,:]

等等

【讨论】：

+1 这就是想法，但您需要动态构建索引元组，例如idx = (slice(None),)*dim + (i,) + (slice(None),)*(2-dim)，然后是 layer = data[idx]。将维度滚动到设定位置（第一个是最方便的）然后总是具有相同的迭代机制，请参阅我的答案。
啊，太好了。滚动解决方案绝对更优雅。

【解决方案3】：

类似的东西？

>>> data = np.array([[1,2,3],[4,5,6],[7,8,9]])
>>> for layer in [data[:,i] for i in range(3)]:
...     print layer
... 
[1 4 7]
[2 5 8]
[3 6 9]

【讨论】：

【解决方案4】：

如果我错了，请纠正我，但在我看来，您的 3D 阵列看起来像：

>>> my_array.shape
    (3,N)

其中 N 是数组的大小。所以如果你想迭代一个维度，你可以这样做：

>>> for item in my_array[1,:]:

这将在第二个维度上进行迭代。

【讨论】：

啊抱歉，应该更清楚 - 这是一个 dim=3 数组，即 shape=(nx, ny, nz)

【解决方案5】：

我认为最初的问题模棱两可，标题也是如此：

对所有 k 迭代 x[k,:,:,...] 会产生与 x.shape[0] 一样多的项目，因此我将其称为迭代维度；
相比之下，在我看来，迭代 slice 意味着迭代，例如超过x[:,i,j,k,...] 对于所有 i,j,k。例如，遍历 ndarray 中的所有列。

虽然这不是 OP 在澄清后提出的问题，但寻找上述第二种情况的解决方案的人可能会发现以下代码很有用：

from itertools import product

def iterslice(x,axis=0):
    sub = [ range(s) for s in x.shape ]
    sub[axis] = (slice(None),)
    for p in product(*sub):
        yield x[p]

【讨论】：