遍历 ndarray 的切片

Question

假设我有一个三维的numpy数组，例如它的维度为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]，这意味着你需要知道维数，并且代码不能处理任意维度的数组。你可以使用类似于[..., :]的东西来填充缺少的维度，但是你仍然需要自己构建它。

抱歉，应该更清楚一些，示例有点过于简单了！

- lost

---

请查看https://dev59.com/I3I-5IYBdhLWcg3w99oH。swapaxes方法是最快的，但最不清晰的。kamjagin的方法可以通过构造一个元组传递到方括号中来进行泛化（即`data[tuple(slice(None), slice(None), i)]与data[:,:,i]`相同）。 - AFoglia

@AFoglia 我认为那个问题的时间不是很相关。使用swapaxes或rollaxis会在设置循环时花费更多时间，但实际迭代速度更快，请参见我答案中添加的时间。在您非常小的示例中，设置占主导地位而不是实际迭代。我不同意可读性，但我可能太习惯于numpy而没有注意到它。 - Jaime

numpy有一个很棒的技巧，可以将切片（元组）从它们要索引的对象中分开构建：numpy.s_[:,:,i]等同于tuple(slice(None), slice(None), i)。 - JAB

（请注意，numpy.s_具有一个属性maketuple，用于确定是否使用单个切片将导致元组或不是[i.e.slice(start，end，step) vs. (slice(start，end，step),)]。默认值为False。） - JAB

@JAB жҳҜзҡ„пјҢдҪҶеҰӮжһңд»–жғіиҰҒдёҖдёӘд»»ж„Ҹз»ҙеәҰе’Ңд»»ж„Ҹ秩数组зҡ„йҖҡз”Ёи§Јжі•пјҢйҖҡиҝҮзј–зЁӢдҪҝз”Ё(slice(None),)*(idim)+(idx,)еҲӣе»әе…ғз»„жҜ”дҪҝз”Ё:еӯ—з¬Ұжӣҙе®№жҳ“гҖӮдҪ дёҚиғҪиҝҷж ·еҒҡdata[(:,)*(idim)+(idx,)]гҖӮ - AFoglia

@AFoglia 但是你可以使用 data[np.index_exp[:]*idim + np.index_exp[idx]]，或者 data[np.index_exp[:]*idim + (np.s_[idx],)]，甚至是 s=np.index_exp 然后跟随 data[s[:]*idim + s[idx]]。(我个人最喜欢最后一个。另外，我已经忘记了 numpy.index_exp 是默认设置为 maketuple 的 IndexExpression 实例。) 虽然在您的具体情况中，data[..., idx] 可能是最明智的选择。 - JAB

Answer 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

- Jaime

---

自从Numpy 1.11版本以后，请使用np.moveaxis(data, k, 0)。 - Jonathan H

Answer 2

这个问题可能有更加优雅的解决方案，但是如果你事先知道维度（例如2），一种解决方法是：

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

或者如果dim=0

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

etc.

- kamjagin

---

+1 这就是想法，但你需要动态构建索引元组，例如 idx = (slice(None),)*dim + (i,) + (slice(None),)*(2-dim)，然后 layer = data[idx]。将维度滚动到固定位置（第一个最方便），然后始终具有相同的迭代机制会更加方便，详见我的答案。 - Jaime

啊，太好了。使用滚动解决方案肯定更加优美。 - kamjagin

Answer 3

我认为原问题存在歧义，标题也同样如此：

• 对于所有k迭代x[k,:,:,...]会产生与x.shape[0]一样多的项，因此我称这种迭代为 维度;
• 相反，对于一个 切片 进行迭代，在我看来意味着例如迭代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]

- Jonathan H

Answer 4

像这样的东西？

>>> 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]

- Elazar

Answer 5

如果我错了，请纠正我，但是在我看来，你的3D数组将会长成这样：

>>> my_array.shape
    (3,N)

N 是你的数组大小。所以，如果你想迭代一个维度，你只需要执行：

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

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

- C.J

---

啊，抱歉，我应该表达得更清楚-这是一个三维数组，即shape=(nx, ny, nz)。 - lost