重新排序ndarray的第n维

Question

我想根据索引order列表对任意维度d的 ndarray 的第 n轴进行重新排序。 如果轴是最后一个，那么这个问题的解决方案（ 重新排序 numpy ndarray 的最后一个维度）就足够了。 然而，就我而言，轴通常不是第一个或最后一个，因此 Ellipsis 本身并不能解决问题。

这是我到目前为止提出的解决方案：

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

我不喜欢这个解决方案是我必须手动转置 ndarray。 我想知道是否存在 Ellipsis 运算符的泛化，以便我们可以考虑选定数量的轴，这样

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

将跳过第n个轴并重新排序第(n+1) 个轴。

这样的事情可能吗？

Answer 1

使用命令np.take_along_axis并将 output 分配给新变量。 请看下面的代码：

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 的评论也是一个有效的答案（可能更好，请给他们一个赞成票：）：

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)]