Using Scipy routines with Numba

Question

I was writing a program in which Scipy CubicSpline routine is used at certain points, because of the use of Scipy routine I cannot use Numba @jit on my whole program.

I recently came across the @overload feature and I was wondering if it could be used in this way,

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())

kindly correct me if my understanding of the @overload feature is wrong and what is the possible solution for using such Scipy libraries with Numba.

Answer 1

Wrapping compiled functions using ctypes (Numba)

Especially for more complex functions, reimplementing everything in numba-compileable Python code can be quite a lot of work, and sometimes slower. The following answer will be on calling C-like functions directly from a shared object or dynamic library.

Compiling the fortran routines

This example will show a way to do this on windows, but it should be straight forward on other os. For a portable interface defining an ISO_C_BINDING is highly recommendable. In this answer I will try it without an interface.

dll.def

EXPORTS
   SPLEV @1

Compilation

ifort /dll dll.def splev.f fpbspl.f /O3 /fast

Calling this function directly from Numba

• Have a look what is expected by the Fortran routine
• Check every input in the wrapper (datatype, contiguousness). You just provide some pointers to the fortran function. There is no additional safety check.

Wrapper

The following code shows two ways on how to call this functions. In Numba it isn't directly possible to pass scalar by reference. You can either allocate an array on the heap (slow for small functions), or use an intrinsic to use stack arrays.

import numba as nb
import numpy as np
import ctypes

lib = ctypes.cdll.LoadLibrary("splev.dll")

dble_p=ctypes.POINTER(ctypes.c_double)
int_p =ctypes.POINTER(ctypes.c_longlong)

SPLEV=lib.SPLEV
SPLEV.restype =  ctypes.c_void_p
SPLEV.argtypes = (dble_p,int_p,dble_p,int_p,dble_p,dble_p,int_p,int_p,int_p)

from numba import types
from numba.extending import intrinsic
from numba.core import cgutils

@intrinsic
def val_to_ptr(typingctx, data):
    def impl(context, builder, signature, args):
        ptr = cgutils.alloca_once_value(builder,args[0])
        return ptr
    sig = types.CPointer(nb.typeof(data).instance_type)(nb.typeof(data).instance_type)
    return sig, impl

@intrinsic
def ptr_to_val(typingctx, data):
    def impl(context, builder, signature, args):
        val = builder.load(args[0])
        return val
    sig = data.dtype(types.CPointer(data.dtype))
    return sig, impl

#with intrinsics, temporary arrays are allocated on stack
#faster but much more relevant for functions with very low runtime
@nb.njit()
def splev_wrapped(x, coeff,e):
    #There are just pointers passed to the fortran function.
    #The arrays have to be contiguous!
    t=np.ascontiguousarray(coeff[0])
    x=np.ascontiguousarray(x)
    
    c=coeff[1]
    k=coeff[2]
    
    y=np.empty(x.shape[0],dtype=np.float64)
    
    n_arr=val_to_ptr(nb.int64(t.shape[0]))
    k_arr=val_to_ptr(nb.int64(k))
    m_arr=val_to_ptr(nb.int64(x.shape[0]))
    e_arr=val_to_ptr(nb.int64(e))
    ier_arr=val_to_ptr(nb.int64(0))
    
    SPLEV(t.ctypes,n_arr,c.ctypes,k_arr,x.ctypes,
        y.ctypes,m_arr,e_arr,ier_arr)
    return y, ptr_to_val(ier_arr)

#without using intrinsics
@nb.njit()
def splev_wrapped_2(x, coeff,e):
    #There are just pointers passed to the fortran function.
    #The arrays have to be contiguous!
    t=np.ascontiguousarray(coeff[0])
    x=np.ascontiguousarray(x)
    
    c=coeff[1]
    k=coeff[2]
    y=np.empty(x.shape[0],dtype=np.float64)
    
    n_arr = np.empty(1,  dtype=np.int64)
    k_arr = np.empty(1,  dtype=np.int64)
    m_arr = np.empty(1,  dtype=np.int64)
    e_arr = np.empty(1,  dtype=np.int64)
    ier_arr = np.zeros(1,  dtype=np.int64)
    
    n_arr[0]=t.shape[0]
    k_arr[0]=k
    m_arr[0]=x.shape[0]
    e_arr[0]=e
    
    SPLEV(t.ctypes,n_arr.ctypes,c.ctypes,k_arr.ctypes,x.ctypes,
        y.ctypes,m_arr.ctypes,e_arr.ctypes,ier_arr.ctypes)
    return y, ier_arr[0]

Answer 2

This is a repost of my solution, posted on numba discourse https://numba.discourse.group/t/call-scipy-splev-routine-in-numba-jitted-function/1122/7.

I had originally gone ahead with @max9111 suggestion of using objmode. It gave a temporary fix. But, since the code was performance critical, I finally ended up writing a numba version of scipy's 'interpolate.splev' subroutine for the spline interpolation.

import numpy as np
import numba
from scipy import interpolate
import matplotlib.pyplot as plt
import time

# Custom wrap of scipy's splrep
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.
    
    """
    
    # Check if x 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]
    
    # Calculate knots & coefficients (cubic spline by default)
    t,c,k = interpolate.splrep(x,y, k=k) 
    
    return (t,c,k,equi_spaced,dx) 

# Numba accelerated implementation of scipy's splev
@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):
        
       # fetch a new x-value arg
       arg = x[i]
       
       # search for knot interval t[l] <= arg <= t[l+1]
       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
       
       # evaluate the non-zero b-splines at arg.    
       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

######################### Testing and comparison #############################

# Generate a data set for interpolation
x, dx = np.linspace(10,100,200, retstep=True)
y = np.sin(x)

# Calculate the cubic spline spline coeff's
coeff_1 = interpolate.splrep(x,y, k=3)  # scipy's splrep
coeff_2 = custom_splrep(x,y, k=3)       # Custom wrap of scipy's splrep

# Generate data for interpolation and randomize
x2 = np.linspace(0,110,10000) 
np.random.shuffle(x2)

# Interpolate
y2 = interpolate.splev(x2, coeff_1) # scipy's splev
y3 = numba_splev(x2, coeff_2)       # Numba accelerated implementation of scipy's splev

# Plot data
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")
# Calculation time evaluation for scipy splev
t1 = time.time()
for n in range(0,10000):
  y2 = interpolate.splev(x2, coeff_1)
print("scipy splev", time.time() - t1)

# Calculation time evaluation for numba splev
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")
# Generate data for interpolation without randomize
x2 = np.linspace(0,110,10000) 

# Calculation time evaluation for scipy splev
t1 = time.time()
for n in range(0,10000):
  y2 = interpolate.splev(x2, coeff_1)
print("scipy splev", time.time() - t1)

# Calculation time evaluation for numba splev
t1 = time.time()
for n in range(0,10000):
  y2 = numba_splev(x2, coeff_2)
print("numba splev",time.time() - t1)

The above code is optimised for faster knot search if the knots are equispaced. On my corei7 machine, if the interpolation is done at random values, numba version is faster,

Scipy’s splev = 0.896s Numba splev = 0.375s

If the interpolation is not done at random values scipy’s version is faster,

Scipy’s splev = 0.281s Numba splev = 0.375s

Ref : https://github.com/scipy/scipy/tree/v1.7.1/scipy/interpolate/fitpack , https://github.com/dbstein/fast_splines

Answer 3

You would either need to fallback to object-mode (locally, like @max9111 suggested), or implement the CubicSpline function yourself in Numba.

For as far as I understand, the overload decorator "only" makes the compiler aware that it can use a Numba-compatible implementation if it encounters the overloaded function. It doesn't magically convert the function the be Numba compatible.

There is a package which expose some Scipy functionality to Numba, but that seems early days and only contains some scipy.special functions so far.

https://github.com/numba/numba-scipy