索引搜索：以准确性换取绩效

Question

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

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

我愿意牺牲一些准确性来获得更好的性能。也就是说，我可以使用指向容差为 +-0.2 的“关闭”元素的索引，而不是 绝对最近 元素（这就是我现在所做的）

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

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