numpy - 使用一些预先计算的地图有效地将值从矩阵复制到矩阵

Question

我有一个大小为I*J的输入矩阵A.

并且输出矩阵B的大小为N*M.

还有一些预先计算出的大小为N*M*2的地图,它决定了B中的每个坐标,它们在A中进行协调.地图没有我可以使用的特定规则或线性.只是一张似乎随机的地图.

矩阵非常大(~5000*~3000)因此创建映射矩阵是不可能的(5000*3000*5000*3000)

我设法使用简单的地图和循环:

for i in range(N):
    for j in range(M):
        B[i, j] = A[mapping[i, j, 0], mapping[i, j, 1]]

我设法使用索引来做到这一点:

B[coords_y, coords_x] = A[some_mapping[:, 0], some_mapping[:, 1]]
# Where coords_x, coords_y are defined as all of the coordinates:
# [[0,0],[0,1]..[0,M-1],[1,0],[1,1]...[N-1,M-1]]

这样做效果更好,但仍然有点慢.

我有无限的时间来计算映射或任何其他效用计算.但是在这些预先计算之后,这种映射应该尽可能快地进行.

目前,我看到的唯一其他选择只是在C中重新实现这一点或更快...

(只是为了说清楚是否有人好奇,我用一些编码创建了一些其他形状不同的图像.但是它的'映射非常复杂,并不是简单或线性可以使用的东西)

Answer 1

如果你有无限的时间进行预计算，你可以通过使用平面索引来获得轻微的加速：

map_f = np.ravel_multi_index((*np.moveaxis(mapping, 2, 0),), A.shape)

然后简单地做：

A.ravel()[map_f]

请注意，这种加速是在我们从花式索引中获得的大幅加速之上的。例如：

>>> A = np.random.random((5000, 3000))
>>> mapping = np.random.randint(0, 15000, (5000, 3000, 2)) % [5000, 3000]
>>> 
>>> map_f = np.ravel_multi_index((*np.moveaxis(mapping, 2, 0),), A.shape)
>>> 
>>> np.all(A.ravel()[map_f] == A[mapping[..., 0], mapping[..., 1]])
True
>>> 
>>> timeit('A[mapping[:, :, 0], mappping[:, :, 1]]', globals=globals(), number=10)
4.101239089999581
>>> timeit('A.ravel()[map_f]', globals=globals(), number=10)
2.7831342950012186

如果我们与原始的循环代码进行比较，加速率将更像是 ~40 倍。

最后，请注意，此解决方案不仅避免了额外的依赖项和潜在的 numba 安装噩梦，而且更简单、更短、更快：

numba:
precomp:    132.957 ms
main        238.359 ms

flat indexing:
precomp:     76.223 ms
main:       219.910 ms

代码：

import numpy as np
from numba import jit

@jit
def fast(A, B, mapping):
    N, M = B.shape
    for i in range(N):
        for j in range(M):
            B[i, j] = A[mapping[i, j, 0], mapping[i, j, 1]]
    return B

from timeit import timeit

A = np.random.random((5000, 3000))
mapping = np.random.randint(0, 15000, (5000, 3000, 2)) % [5000, 3000]

a = np.random.random((5, 3))
m = np.random.randint(0, 15, (5, 3, 2)) % [5, 3]

print('numba:')
print(f"precomp: {timeit('b = fast(a, np.empty_like(a), m)', globals=globals(), number=1)*1000:10.3f} ms")
print(f"main     {timeit('B = fast(A, np.empty_like(A), mapping)', globals=globals(), number=10)*100:10.3f} ms")

print('\nflat indexing:')
print(f"precomp: {timeit('map_f = np.ravel_multi_index((*np.moveaxis(mapping, 2, 0),), A.shape)', globals=globals(), number=10)*100:10.3f} ms")
map_f = np.ravel_multi_index((*np.moveaxis(mapping, 2, 0),), A.shape)
print(f"main:    {timeit('B = A.ravel()[map_f]', globals=globals(), number=10)*100:10.3f} ms")