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)