python - 从 numba jitted 函数调用非 jitted 函数

Question

我的代码是这样的:

@jit(nopython=True)
def sum_fn(arg1, arg2, ...argn):
    .....
    for i in xrange(len(arg2)): #For each bin
        l, p = fn1(arg1, arg2...argn) 
        res = res + fn2(arg1, arg2, arg3) 
    return res

@jit(nopython=True)
def fn1(a1, a2,...an):
    ....
    return r1, r2

def fn2(l_lk, l_pk, l_lvals):
        f_i = interp1d(l_lk, l_pk,kind='linear') #Scipy.interpolate.interp1d
        ftmp = fn3(f_i,l_lk,l_pk)
        return 10**ftmp(l_lvals)

调用 fn2 似乎会出错，因为它没有被导入到 numba 编译过程中，这与 fn1 不同，后者是 jitted。不幸的是，我不能 jit fn2 因为它使用了 Numba 无法识别的 scipy.interpolate.interp1d 函数。我该如何解决这个问题？

Answer 1

您不能从使用 nopython=True 进行 jitted 的函数中调用非 jitted 函数。前者根据定义是一个 python 函数。您唯一真正的选择是将您自己的 interp1d 版本编写为 jitted 函数(同样使用 nopython=True)。有时，通过查看原始 scipy 或 numpy 函数的源代码，可以很简单地删除所需的功能。不幸的是，有时这很困难。

以下库可能会有帮助:

https://github.com/EconForge/interpolation.py