重新排序ndarray的第n维

Question

I want to reorder the nth axis of a ndarray of abitrary dimension d according to a list of indices order . If the axis was the last one, the solution from this question ( Reordering last dimension of a numpy ndarray ) would be enough. In my case, however, the axis is not, in general, the first or the last one, so that Ellipsis alone does not solve the issue.

This is the solution I have come up so far:

axes_list = list(range(d))
axes_list[0], axes_list[i] = axes_list[i], axes_list[0]

ndarr = np.transpose(ndarr, axes=axes_list)[order,...]   # Switches the axis with the first and reorders
ndarr = np.transpose(ndarr, axes=axes_list)              # Switches the axes back

What I don't like about this solution is that I have to manually transpose the ndarray. I was wondering if there is a generalization of the Ellipsis operator such that we can account for a chosen number of axes, such that

ndarr[GenEllipsis(n),order,...]

would skip the n first axes and would reorder the (n+1)th one.

Is such a thing possible?

Answer 1

Use the command np.take_along_axis and assign the output to a new variable. See the code below:

arr = np.random.randn(10,3,23,42,3)
ax = 3 #change this to your 'random' axis
order = np.random.permutation(list(range(arr.shape[ax])))
#order needs to have the same number of dims as arr
order = np.expand_dims(order,tuple(i for i in range(len(arr.shape)) if i != ax)) 
shuff_arr = np.take_along_axis(arr,order,ax)

@hpaulj's comment is also a valid answer (and likely better, please give them an upvote:):

arr = np.random.randn(10,3,23,42,3)
ax = 3 #change this to your 'random' axis
order = np.random.permutation(list(range(arr.shape[ax])))
#set the correct dim to 'order'
alist = [slice(None)]*len(arr.shape)
alist[ax] = order
shuff_arr = arr[tuple(alist)]