问将Scipy例程与Numba一起使用

EN

from numba.extending import overload
from numba import jit
from scipy.interpolate import CubicSpline
import numpy as np

x = np.arange(10)
y = np.sin(x)
xs = np.arange(-0.5, 9.6, 0.1)

def Spline_interp(xs,x,y):
    cs = CubicSpline(x, y)
    ds = cs(xs)
    return ds

@overload(Spline_interp)
def jit_Spline_interp(xs,x,y):
   ds = Spline_interp(xs,x,y)
   def jit_Spline_interp_impl(xs,x, y):
       return ds
   return jit_Spline_interp_impl

@jit(nopython=True)
def main():

    # other codes compatible with @njit

    ds = Spline_interp(xs,x,y)

    # other codes compatible with @njit

    return ds

print(main())