首页 > 解决方案 > 在任意轴上迭代体积的更 Pythonic 方式?

问题描述

我有一个函数,它接受一个 3D numpy 数组(我们将其称为卷),并将其转换为 2D 切片列表。我希望用户能够指定对其进行切片的轴。我用下面的代码管理这个,但是三重 if 语句似乎不是最优雅的方法。我会感谢人们对是否可以以更好的方式实现这一点的想法。

axis = 0 # Can be set to 0, 1, or 2 

volume = np.ones((100, 100, 100))

n_slices = volume.shape[axis]

slices = []

for i in range(n_slices):

    if axis == 0:
        my_slice = volume[i, :, :]
    elif axis == 1:
        my_slice = volume[:, i, :]
    elif axis == 2:
        my_slice = volume[:, :, i]

    slices.append(my_slice)

标签: pythonarraysnumpyvolumenibabel

解决方案


只需使用np.moveaxis-

slices_ar = np.moveaxis(volume,axis,0)

最好的部分是它是一个输入视图,因此在运行时几乎是免费的。让我们验证view-part-

In [83]: np.shares_memory(volume, np.moveaxis(volume,axis,0))
Out[83]: True

或者,使用np.rollaxis它做同样的事情 -

np.rollaxis(volume,axis,0)

推荐阅读