这是我在 numba 论坛上发布的解决方案的转载
https://numba.discourse.group/t/call-scipy-splev-routine-in-numba-jitted-function/1122/7。
起初,我采用了 @max9111 建议的 objmode。它提供了一个临时的解决方案。但由于代码对性能要求很高,最终我编写了一个 numba 版本的 scipy 的 'interpolate.splev' 子程序进行样条插值。
import numpy as np
import numba
from scipy import interpolate
import matplotlib.pyplot as plt
import time
def custom_splrep(x, y, k=3):
"""
Custom wrap of scipy's splrep for calculating spline coefficients,
which also check if the data is equispaced.
"""
x_diff = np.diff(x)
equi_spaced = all(np.round(x_diff,5) == np.round(x_diff[0],5))
dx = x_diff[0]
t,c,k = interpolate.splrep(x,y, k=k)
return (t,c,k,equi_spaced,dx)
@numba.njit(cache=True)
def numba_splev(x, coeff):
"""
Custom implementation of scipy's splev for spline interpolation,
with additional section for faster search of knot interval, if knots are equispaced.
Spline is extrapolated from the end spans for points not in the support.
"""
t,c,k, equi_spaced, dx = coeff
t0 = t[0]
n = t.size
m = x.size
k1 = k+1
k2 = k1+1
nk1 = n - k1
l = k1
l1 = l+1
y = np.zeros(m)
h = np.zeros(20)
hh = np.zeros(19)
for i in range(m):
arg = x[i]
if(equi_spaced):
l = int((arg-t0)/dx) + k
l = min(max(l, k1), nk1)
else:
while not ((arg >= t[l-1]) or (l1 == k2)):
l1 = l
l = l-1
while not ((arg < t[l1-1]) or (l == nk1)):
l = l1
l1 = l+1
h[:] = 0.0
hh[:] = 0.0
h[0] = 1.0
for j in range(k):
for ll in range(j+1):
hh[ll] = h[ll]
h[0] = 0.0
for ll in range(j+1):
li = l + ll
lj = li - j - 1
if(t[li] != t[lj]):
f = hh[ll]/(t[li]-t[lj])
h[ll] += f*(t[li]-arg)
h[ll+1] = f*(arg-t[lj])
else:
h[ll+1] = 0.0
break
sp = 0.0
ll = l - 1 - k1
for j in range(k1):
ll += 1
sp += c[ll]*h[j]
y[i] = sp
return y
x, dx = np.linspace(10,100,200, retstep=True)
y = np.sin(x)
coeff_1 = interpolate.splrep(x,y, k=3)
coeff_2 = custom_splrep(x,y, k=3)
x2 = np.linspace(0,110,10000)
np.random.shuffle(x2)
y2 = interpolate.splev(x2, coeff_1)
y3 = numba_splev(x2, coeff_2)
plt.plot(x,y,'--', linewidth=1.0,color='green', label='data')
plt.plot(x2,y2,'o',color='blue', markersize=2.0, label='scipy splev')
plt.plot(x2,y3,'.',color='red', markersize=1.0, label='numba splev')
plt.legend()
plt.show()
print("\nTime for random interpolations")
t1 = time.time()
for n in range(0,10000):
y2 = interpolate.splev(x2, coeff_1)
print("scipy splev", time.time() - t1)
t1 = time.time()
for n in range(0,10000):
y2 = numba_splev(x2, coeff_2)
print("numba splev",time.time() - t1)
print("\nTime for non random interpolations")
x2 = np.linspace(0,110,10000)
t1 = time.time()
for n in range(0,10000):
y2 = interpolate.splev(x2, coeff_1)
print("scipy splev", time.time() - t1)
t1 = time.time()
for n in range(0,10000):
y2 = numba_splev(x2, coeff_2)
print("numba splev",time.time() - t1)
如果节点间距相等,则上述代码针对更快的节点搜索进行了优化。
在我的corei7机器上,如果在随机值处进行插值,则numba版本更快,
Scipy的splev = 0.896秒
Numba的splev = 0.375秒
如果不在随机值处进行插值,则Scipy的版本更快,
Scipy的splev = 0.281秒
Numba的splev = 0.375秒
参考:https://github.com/scipy/scipy/tree/v1.7.1/scipy/interpolate/fitpack,
https://github.com/dbstein/fast_splines