python - 给定排列矩阵的 numpy 数组的每行(或列)的高效排列

Question

让我们假设，我有两个给定的 ndarray，其中矩阵 mapping 包含矩阵 mask 的行应该如何的信息排列。我们可以假设映射矩阵来自其他一些算法。

import numpy as np
T, K, F = 2, 3, 5
mask = np.random.randint(4, size=(T, K, F))
mapping = np.asarray([
    [0, 1, 2],
    [0, 1, 2],
    [2, 0, 1],
    [0, 1, 2],
    [1, 0, 2]
])

执行此操作的直接方法是应用 for 循环:

out = np.empty_like(mask)
for f in range(F):
    out[:, :, f] = mask[:, mapping[f, :], f]

这似乎相当有效，所以我查看了 Numpy advanced indexing并找到了这个解决方案:

out = mask[
    np.arange(T)[:, None, None],
    mapping.T[None, :, :],
    np.arange(F)[None, None, :]
]

安answer相关问题建议使用ogrid:

ogrid = np.ogrid[:T, :1, :F]
out = mask[
    ogrid[0],
    mapping.T[None, :, :],
    ogrid[2]
]

创建所有中间数组并正确广播它们似乎非常不舒服。那么执行所需的重新排序的最佳方法是什么？

时间信息:

为了提供有意义的计时信息，我使用了一些更接近我的应用程序的形状。随机排列只是为了简化示例。

T, K, F = 1000, 3, 257
mask = np.random.randint(4, size=(T, K, F))
mapping = np.stack([list(np.random.permutation(np.arange(3))) for _ in range(F)])

结果如下:

for loop:                 100 loops, best of 3: 8.4 ms per loop
three times broadcasting: 100 loops, best of 3: 8.37 ms per loop
ogrid:                    100 loops, best of 3: 8.33 ms per loop
swapaxis:                 100 loops, best of 3: 2.43 ms per loop
transpose:                100 loops, best of 3: 2.08 ms per loop

Answer 1

定义“最佳”是有争议的，但这是 advanced-indexing 的一种方法。 -

mask[:,mapping, np.arange(F)[:,None]].swapaxes(1,2)

<小时/>

另一种方法是转置映射，然后使用最后一个轴的范围数组，而不扩展到2D。 mapping 的每一行最后一个轴 (axis=-1) 决定元素沿倒数第二个轴 (axis=-2) 的顺序掩码。因此，我们需要在映射上进行转置。在第一种方法中，我们通过后者交换轴来实现这种转置行为。我可以保证这一点的效率。

因此，我们将实现如下所示 -

mask[:,mapping.T, np.arange(F)]