使用 cython 或替代方法优化 np.searchsorted()答案

【问题标题】：Optimising np.searchsorted() with cython or alternative使用 cython 或替代方法优化 np.searchsorted()
【发布时间】：2022-01-12 04:59:39
【问题描述】：

我有以下功能。我试图优化 np.searchsorted() 函数，并在 Cython 中编写了我自己的函数，结果很好。但是后来意识到将前者简单地包装在 cdef cfunc() 中效果更好，而且我不了解 python、cython 或 numpy 的机制来知道为什么。编写此函数的最佳、最快/有效的方法是什么。在我的脚本中，这是在某些计算过程中被调用数千次的函数之一，我希望每次调用的成本尽可能低。另外，我不需要完全复制 np.searchsorted()，我只需要生成特定的输出 i = np.searchsorted(time_points > t, True)

案例一

%%cython --compile-args=-fopenmp --link-args=-fopenmp -a
cimport openmp
cimport cython
#cimport openmp
"""
%%cython -a
"""
from cpython cimport array
import array
import cython.parallel as cp
from cython.parallel import parallel, prange
import numpy as np
cimport numpy as np
import random
import pickle
import matplotlib.pyplot as pl

from timeit import default_timer as timer
cdef int ntime = 10**4
cdef float t = 2 #some variable in general 
time_points = np.arange(ntime, dtype=np.float64)
#print(time_points)
#search_time =[]
start_update = timer()
i = np.searchsorted(time_points > t, True)
print(i)
end_update = timer()
print(end_update-start_update)

案例 2

@cython.boundscheck(False)  # Deactivate bounds checking                                                                  
@cython.wraparound(False)   # Deactivate negative indexing.                                                               
@cython.cdivision(True)     # Deactivate division by 0 checking.
cdef mysearch(np.ndarray[np.float64_t, ndim=1] time_points, float val, int nt):
    cdef int idx
    cdef int return_val
    #cdef int total 
    if val >= nt or nt-1< val<nt:
        print("first")
        return nt
    else:
        for idx in range(nt):
            #print(idx)
            
            if time_points[idx] <= val:
                #return_val = np.int(time_points[idx])
                #break
                #print("sec")
                continue
            else:
                #print("third", idx)
                return_val = np.int(time_points[idx])
                #continue
                break
        
    return return_val
print(mysearch(time_points, t, ntime))
#i = np.searchsorted(time_points > 2, True)
#print(i)
end_update = timer()
b= end_update-start_update
print(b)

案例 3

start_update = timer()
cdef cfunc():
    i = np.searchsorted(time_points > t, True)
    print(i, np.searchsorted(time_points > t, True))
    return i
end_update = timer()
a=  end_update-start_update
print(a)

案例1、案例2、案例3的时间为0.00013096071779727936 5.245860666036606e-05 和 6.770715117454529e-07

我很高兴案例 2 比案例 1 运行得更好，但为什么案例 3 快这么多？发生了什么特别的事情，我怎样才能超越这些速度？

【问题讨论】：

np.searchsorted 已编译代码；我不希望手写的cython 替代品做得更好——除非你是一个优秀的c 程序员。当 Python 级别的迭代替换为 compile、c 代码迭代时，我们期望获得最佳的速度提升。
@hpaulj “我们期望将 Python 级别的迭代替换为编译、c 代码迭代时的最佳速度改进”是在案例 3 中将 np.searchsort 包装在 cfun 中时发生的情况吗？案例 3 和案例 1 有何不同？
你能在没有print 语句的情况下进行测试吗？我喜欢用timeit 重复测试（这在ipython 环境中特别容易）。
Case 3 实际上没有运行！它只定义函数（在编译时发生，在运行时是无操作的

标签： python numpy cython

【解决方案1】：

np.searchsorted(time_points > t, True) 的主要问题是time_points > t。实际上，虽然np.searchsorted 在O(log n) 时间内运行（n 的大小与time_points 相同），但表达式time_points > t 是在O(n) 时间内计算的，显然是瓶颈。事实上，您不需要这个子表达式：i = np.searchsorted(time_points, t, 'right') 应该可以正确完成工作。请注意，最后一种情况不会在两个计时器之间执行任何操作，正如@DavidW 在 cmets 中所解释的那样：它只是定义了一个 Cython 函数。

【讨论】：

谢谢。这个解释很有道理。