python - 索引搜索 : trade accuracy for performance

Question

我有一个简单的两行代码块，它根据在另一个数组中找到的最接近的元素将值添加到一个数组。由于它深埋在 MCMC 中，它被执行了数百万次，我需要它尽可能高效。

下面的代码可以工作，而且很容易解释。基本上:数组 arr2[0](用于在 arr0 中查找最近元素的数组)包含 (10., 25.)< 范围内的值。目前，我使用 np.searchsorted() 为 arr2[0] 中的每个元素寻找 arr0 中的绝对最接近 元素| ，利用 arr0 已经排序的事实。

我愿意牺牲一些准确性来换取更好的性能。也就是说，我可以接受一个指向容差为 +-0.2 的“close”元素的索引，而不是 absolute closest 元素(这就是我现在做)

这能做到吗？更重要的是:是否可以做到这一点并提高代码的性能？

import numpy as np

# Random initial data with the actual shapes used by my code.
Nmax = 1000000
arr0 = np.linspace(5., 30., Nmax)
D = np.random.randint(2, 4)
arr1 = np.random.uniform(-3., 3., (D, Nmax))
arr2 = np.random.uniform(10., 25., (10, 1500))

# Can these two lines be made faster?
# Indexes of elements in 'arr0' closest to the elements in 'arr2[0]'
closest_idxs = np.searchsorted(arr0, arr2[0])
# Add elements from 'arr1' to the first dimensions of 'arr2', according
# to the indexes found above.
arr_final = arr2[:arr1.shape[0]] + arr1[:, closest_idxs]

Answer 1

对于具有给定容差值的近似匹配，我们可以使用它来将第一个 arg 减少到 searchsorted 并因此进行优化，就像这样 -

tol = 0.2 # tolerance value
s = int(np.round(tol/(arr0[1]-arr0[0])))
i = np.searchsorted(arr0[::s], arr2[0])
i -= (arr0[i*s]-arr2[0])>tol/2
closest_idxs_out = i*s

给定设置的时间 -

In [123]: %%timeit
     ...: closest_idxs = np.searchsorted(arr0, arr2[0])
     ...: arr_final = arr2[:arr1.shape[0]] + arr1[:, closest_idxs]
1000 loops, best of 3: 641 µs per loop

In [125]: %%timeit
     ...: tol = 0.2 # tolerance value
     ...: s = int(np.round(tol/(arr0[1]-arr0[0])))
     ...: i = np.searchsorted(arr0[::s], arr2[0])
     ...: i -= (arr0[i*s]-arr2[0])>tol/2
     ...: closest_idxs_out = i*s
10000 loops, best of 3: 63.2 µs per loop