python之在任意轴上迭代体积的更多 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) 

请您参考如下方法:

只需使用 np.moveaxis -

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

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

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

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

np.rollaxis(volume,axis,0)