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

Question

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

以及大小为N * M的输出矩阵B

还有一些大小为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倍。

最后，请注意，此解决方案不仅避免了额外的依赖性和潜在的安装梦night，而且变得更简单，更短，更快捷：

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

Answer 2

对这些类型的性能至关重要的问题，一个非常好的解决方案是使其简单易用，并利用其中一种高性能软件包。最简单的可能是Numba，它提供了jit装饰器，该装饰器可编译数组并将大量代码循环到优化的LLVM。下面是一个完整的示例：

from time import time
import numpy as np
from numba import jit

# Function doing the computation
def normal(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

# The same exact function, but with the Numba jit decorator
@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

# Create sample data
def create_sample_data(I, J, N, M):
    A = np.random.random((I, J))
    B = np.empty((N, M))
    mapping = np.asarray(np.stack((
        np.random.random((N, M))*I,
        np.random.random((N, M))*J,
        ), axis=2), dtype=int)
    return A, B, mapping
A, B, mapping = create_sample_data(500, 600, 700, 800)

# Run normally
t0 = time()
B = normal(A, B, mapping)
t1 = time()
print('normal took', t1 - t0, 'seconds')

# Run using Numba.
# First we should run the function with smaller arrays,
# just to compile the code.
fast(*create_sample_data(5, 6, 7, 8))
# Now, run with real data
t0 = time()
B = fast(A, B, mapping)
t1 = time()
print('fast took', t1 - t0, 'seconds')

这使用您自己的循环解决方案，该解决方案在本质上使用标准Python慢​​，但是在使用Numba时与C一样快。在我的机器上，normal函数在0.270秒内执行，而fast函数在0.00248秒内执行。也就是说，Numba几乎免费为我们提供了 109倍加速（！）。

请注意，fast Numba函数被调用两次，首先用小的输入数组，然后才用实数据。这是一个经常被忽略的关键步骤。没有它，您会发现性能提升不如第一次调用用于编译代码那样好。在此初始调用中，输入数组的类型和尺寸应该相同，但是每个尺寸的大小并不重要。

我在函数外创建B并将其作为参数传递（“填充值”）。您也可以在函数内部分配B，Numba不在乎。

获取Numba的最简单方法是通过Anaconda发行版。

Answer 3

一种选择是使用numba，它通常可以大大简化这种简单的算法代码。

import numpy as np
from numba import njit

I, J = 5000, 5000
N, M = 3000, 3000

A = np.random.randint(0, 10, [I, J])
B = np.random.randint(0, 10, [N, M])
mapping = np.dstack([np.random.randint(0, I - 1, (N, M)),
                     np.random.randint(0, J - 1, (N, M))])

B0 = B.copy()

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

new = njit(orig)

为我们提供匹配结果：

In [313]: Bold = B0.copy()

In [314]: orig(A, Bold, mapping)

In [315]: Bnew = B0.copy()

In [316]: new(A, Bnew, mapping)

In [317]: (Bold == Bnew).all()
Out[317]: True

并且速度更快：

In [320]: %time orig(A, B0.copy(), mapping)
Wall time: 6.11 s

In [321]: %time new(A, B0.copy(), mapping)
Wall time: 257 ms

在第一次调用后仍需要进行jit工作时速度更快：

In [322]: %time new(A, B0.copy(), mapping)
Wall time: 171 ms

In [323]: %time new(A, B0.copy(), mapping)
Wall time: 163 ms

将添加两行代码的效果提高了30倍。

Answer 4

您可以做的最直接的优化是删除本机python循环并使用花式numpy索引。您已经有了执行此操作的数组：

import numpy as np

A = np.random.rand(2000,3000)
B = np.empty((2500,3500)) # just for shape, really

# this is the same as your original, but with random indices
mapping = np.stack([np.random.randint(0, A.shape[0] - 1, B.shape),
                    np.random.randint(0, A.shape[1] - 1, B.shape)],
                   axis=-1)

# your loopy original
def loopy(A, B, mapping):
    B = B.copy()
    for i in range(B.shape[0]):
        for j in range(B.shape[1]):
            B[i, j] = A[mapping[i, j, 0], mapping[i, j, 1]]
    return B

# vectorization with fancy indexing
def fancy(A, mapping):
    return A[mapping[...,0], mapping[...,1]]

请注意，fancy高级索引功能不需要B数组的预分配，因为通过索引操作构造了一个新数组。

fancy indexing版本有一些变化，可能的效率略高：首先将mapping的最后一个维度放在第一位，这样两个索引数组都是连续的内存块。从我的时序测试中可以看出，在上述设置中，这恰好慢一些。无论如何：

mapping_T = mapping.transpose(2, 0, 1).copy() # but it's actually `mapping` without axis=-1 kwarg
# has shape (2, N, M)

def fancy_T(A, mapping_T):
    return A[tuple(mapping_T)]

与Paul Panzer noted in a comment一样，仅在.transpose上调用mapping不会创建副本，而是使用跨步技巧来实现转置。为了获得一个连续的数组（这是优化的重点），我们需要强制创建副本。

我在ipython中得到以下计时：

# loopy(A, B, mapping)
6.63 s ± 141 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
# fancy(A, mapping)
250 ms ± 3.79 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
# fancy_T(A, mapping_T)
277 ms ± 1.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

说实话，我不明白为什么原始的数组顺序比转置的数组更快，但是确实如此。