More pythonic way to iterate through volume in arbitrary axis?

Question

I have a function which takes a 3D numpy array (which we'll call a volume), and converts it into a list of 2D slices. I want the user to be able to specify the axis through which it is sliced. I am managing this with the code below, but the triple if statement doesn't seem like the most elegant way to do this. I would appreciate people's thoughts on if this could be achieved in a better way.

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)

Answer 1

Simply use np.moveaxis -

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

Best part is that it's a view into the input and hence virtually free on runtime. Let's verify the view-part -

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

Alternatively, use np.rollaxis that does the same -

np.rollaxis(volume,axis,0)