重新排序ndarray的第n维

Question

我想根据索引列表 order 对任意维度 d 的 ndarray 的 nth 轴重新排序。如果轴是最后一个，那么这个问题（Reordering last dimension of a numpy ndarray）的解决方案就足够了。然而，就我而言，轴通常不是第一个或最后一个，因此仅靠省略号并不能解决问题。

这是我迄今为止提出的解决方案：

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 并将输出分配给一个新变量。请看下面的代码：

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