Numba: support for scipy.special functions : feature request

[sahaskn]

If possible, it would be great if numba supports scipy special functions which are mostly universal functions and there is cython API for those functions as well.

[stuartarchibald]

Agree that this needs to be done. As an interim, it's trivial to create these as needed following the docs http://numba.pydata.org/numba-doc/latest/extending/high-level.html. Here is an example for scipy.special.beta.

from numba.extending import get_cython_function_address
from numba import vectorize, njit
import ctypes
import numpy as np

addr = get_cython_function_address("scipy.special.cython_special", "beta")
functype = ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double, ctypes.c_double)
beta_fn = functype(addr)


@vectorize('float64(float64, float64)')
def vec_beta(x, y):
    return beta_fn(x, y)

@njit
def beta_in_njit(x, y):
    return vec_beta(x, y)

d = np.linspace(0, 20, 50)

got = vec_beta(d, d[::-1])

got_njit_context = beta_in_njit(d, d[::-1])

from scipy.special import beta
expected = beta(d, d[::-1])

np.testing.assert_allclose(got, expected)
np.testing.assert_allclose(got_njit_context, expected)

[person142]

I also started a pure Numba port of special functions here:

https://github.com/person142/special

There is a (small) chance that what you want is there.

[sahaskn]

@stuartarchibald Thanks a lot for the example. :)
@person142 Thanks for pointing page.

[guoyang0123]

@stuartarchibald, Thank you very much for your example! Your example works very well for functions involving "double".
However, in my case, there is a special type "Dd_number_t" defined in the cython_special.pxd

ctypedef fused Dd_number_t:
double complex
double

The numba manual said "check the extension module’s __pyx_capi__ attribute".
I tried very hard. But can not get any clue. Could you please give me a example about how to define "CFUNCTYPE" or some hints?
Thank you very much!

[stuartarchibald]

@guoyang0123 the __pyx_capi__ attribute is an attribute of the Cython exported module. For example, the scipy.special.cython_special export:

In [9]: import scipy.special.cython_special as cysp                                                                             

In [10]: cysp                                                                                                                   
Out[10]: <module 'scipy.special.cython_special' from '<path>/lib/python3.7/site-packages/scipy/special/cython_special.cpython-37m-x86_64-linux-gnu.so'>

In [11]: cysp.__pyx_capi__                                                                                                      
Out[11]: 
{'agm': <capsule object "double (double, double, int __pyx_skip_dispatch)" at 0x7f2496c2a360>,
 'bdtrik': <capsule object "double (double, double, double, int __pyx_skip_dispatch)" at 0x7f2496c2a9c0>,
 'bdtrin': <capsule object "double (double, double, double, int __pyx_skip_dispatch)" at 0x7f2496c2a810>,
 'bei': <capsule object "double (double, int __pyx_skip_dispatch)" at 0x7f2494229210>,
 'beip': <capsule object "double (double, int __pyx_skip_dispatch)" at 0x7f2494229390>,
 'ber': <capsule object "double (double, int __pyx_skip_dispatch)" at 0x7f2495bd52a0>,
 'berp': <capsule object "double (double, int __pyx_skip_dispatch)" at 0x7f2495bd5330>,
... (This continues for pages, there's a lot of functions)

The fused types in e.g. scipy.special.cython_special will resolve into concrete implementations, for example, this is one way to bind to the float64 (double) variant of the airy function:

from numba.extending import get_cython_function_address
from numba import vectorize, njit
import ctypes
import numpy as np

_PTR = ctypes.POINTER
_dble = ctypes.c_double
_ptr_dble = _PTR(_dble)

# signature is:
# void airy(Dd_number_t x0, Dd_number_t *y0, Dd_number_t *y1, Dd_number_t *y2, Dd_number_t *y3) nogil
# bind to the real space variant of the function
addr = get_cython_function_address("scipy.special.cython_special", "__pyx_fuse_1airy")
functype = ctypes.CFUNCTYPE(None, _dble, _ptr_dble, _ptr_dble, _ptr_dble, _ptr_dble)
airy_float64_fn = functype(addr)

@njit
def numba_airy_float64(x):
    y0 = np.empty(1, dtype=np.float64)
    y1 = np.empty(1, dtype=np.float64)
    y2 = np.empty(1, dtype=np.float64)
    y3 = np.empty(1, dtype=np.float64)
    airy_float64_fn(x, y0.ctypes, y1.ctypes, y2.ctypes, y3.ctypes)
    return y0[0], y1[0], y2[0], y3[0]

x = -15

got = numba_airy_float64(x)

from scipy.special import airy
expected = airy(x)

np.testing.assert_allclose(got, expected)

[guoyang0123]

Dear @stuartarchibald,
Thank you for your help! My code works! It is good for me to learn deeper about numba.
I found numba is more efficient than C code connected by pybind11. I will stick on numba!
Yang

[stuartarchibald]

@guoyang0123 no problem. Glad it's working for you!

[max9111]

@stuartarchibald
Can you also add an example dealing with complex double and complex double *?
eg. for void airy(double complex, double complex *, double complex *, double complex *, double complex *)

Thanks a lot!

[stuartarchibald]

@max9111 yes, but be advised that because ctypes does not support C complex types that this is a) a bit of a hack b) may not be portable c) makes a load of assumptions, however, it'll probably be ok.

from numba.extending import get_cython_function_address
from numba import vectorize, njit
import ctypes
import numpy as np

_PTR = ctypes.POINTER
_dble = ctypes.c_double
_ptr_dble = _PTR(_dble)

addr = get_cython_function_address("scipy.special.cython_special", "__pyx_fuse_0airy")
# pass the complex value on the stack as two adjacent double values
functype = ctypes.CFUNCTYPE(None, _dble, _dble, _ptr_dble, _ptr_dble, _ptr_dble, _ptr_dble)
airy_complex128_fn = functype(addr)

@njit
def numba_airy_complex128(x):
    y0 = np.empty(2, dtype=np.float64)
    y1 = np.empty(2, dtype=np.float64)
    y2 = np.empty(2, dtype=np.float64)
    y3 = np.empty(2, dtype=np.float64)
    airy_complex128_fn(np.real(x), np.imag(x), y0.ctypes, y1.ctypes, y2.ctypes, y3.ctypes)
    def ascmplx(z):
        return np.complex(z[0] + 1j * z[1])
    return ascmplx(y0), ascmplx(y1), ascmplx(y2), ascmplx(y3)

x = -2 + 3j

got = numba_airy_complex128(x)

from scipy.special import airy
expected = airy(x)

np.testing.assert_allclose(got, expected)

[mpetroff]

Would this work for functions that return a single complex value, e.g., scipy.special.wofz?

I can get the real part without issue by using a double as the return type, but my efforts to get the full complex value have not been successful. Using a pointer, I get TypeError: cannot convert native float64* to Python object, and using ctypes.c_double * 2 I get Unsupported ctypes type: <class '__main__.c_double_Array_2'>.

[guoyang0123]

Hi @mpetroff ,
please follow stuartarchibald's reply on Jan 08. It works. I tried a few functions already.

[stuartarchibald]

@mpetroff I don't think this is possible right now. The returned complex number is too wide for any of the supported types and it is necessary for it to be declared as the returned item to match complex return behaviour in C99. In practice this is done however the compiler sees fit, but typically through putting the real/imag parts into xmm* registers on return and them loading from those in the caller as needed.

I can get the real part without issue by using a double as the return type

I'd expect this to be the case but purely down to luck of code generation. It so happens that the ctypes call places the real part in xmm0 which is then read on return as an assumed double value.

[max9111]

@mpetroff It is always a possibility to write a cython cdef function which exposes a simpler interface.

Writing a Cyhon wrapper

Special.pyx

cimport scipy.special.cython_special
cimport numpy as np

cdef api wofz(double in1_real,double in1_imag,double *out_real,double *out_imag):
  cdef double complex z
  z.real=in1_real
  z.imag=in1_imag

  cdef double complex out=scipy.special.cython_special.wofz(z)

  out_real[0]=out.real
  out_imag[0]=out.imag

Setup.py

from distutils.core import setup, Extension
from Cython.Build import cythonize
import numpy

setup(
     ext_modules=cythonize("special.pyx"),
     include_dirs=[numpy.get_include()]
)  

Wrapper

from numba.extending import get_cython_function_address
from numba import vectorize, njit
import ctypes
import numpy as np
import numba as nb

_PTR = ctypes.POINTER
_dble = ctypes.c_double
_ptr_dble = _PTR(_dble)

addr = get_cython_function_address("special", "wofz")
functype = ctypes.CFUNCTYPE(None, _dble, _dble, _ptr_dble, _ptr_dble)
wofz_fn = functype(addr)

#It looks like the Numba wrapper takes about 85% of the total runtime
@nb.njit
def numba_wofz_complex(x):
    out_real = np.empty(1,dtype=np.float64)
    out_imag = np.empty(1,dtype=np.float64)
    wofz_fn(np.real(x), np.imag(x), out_real.ctypes, out_imag.ctypes)

    return np.complex(out_real[0] + 1j * out_imag[0])


from scipy.special import wofz

val=np.complex(15+2j)
expected = wofz(val)
got = numba_wofz_complex(val)

np.testing.assert_allclose(got, expected)

@stuartarchibald By commenting out the line wofz_fn(np.real(x), np.imag(x), out_real.ctypes, out_imag.ctypes) I observed that this line (the actual function) takes only about 15% of the total runtime. Is there a faster way to wrap this function?

[mpetroff]

@max9111 Thanks! I had started writing something with CFFI and libcerf (same Faddeeva function implementation as SciPy), but your Cython code is better, since it uses SciPy directly and since it can be included directly in a notebook using Cython's cell magic.

[stuartarchibald]

I raised the issue of returning complex values from C-ABI functions at yesterday's core developer meeting. The outcome was that it's probably possible for Numba to be able to generate a C-ABI compliant wrapper (as clang and GCC are probably doing something sufficiently similar) to handle the case where a complex value is returned from a ctypes bound function. https://github.com/numba/numba/issues/3860 tracks.

[joernweissenborn]

I stumbled upon this issue today which is kind of roadblock for introducing numba to my project.

Is there any update or even an ETA for this feature?

If not, I see some workarounds. I didn't fully understood them yet, can somebody explain it a little more in detail? I need the scipy.special.erf and scipy.special.erfcx functions.

[ickc]

The key to find the right function is __pyx_capi__.

e.g.

>>> {key: value for key, value in scipy.special.cython_special.__pyx_capi__.items() if 'erf' in key}
{'__pyx_fuse_0erf': <capsule object "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)" at 0x7fc717fd5780>,
 '__pyx_fuse_1erf': <capsule object "double (double, int __pyx_skip_dispatch)" at 0x7fc717fd57b0>,
 '__pyx_fuse_0erfc': <capsule object "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)" at 0x7fc717fd57e0>,
 '__pyx_fuse_1erfc': <capsule object "double (double, int __pyx_skip_dispatch)" at 0x7fc717fd5810>,
 '__pyx_fuse_0erfcx': <capsule object "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)" at 0x7fc717fd5840>,
 '__pyx_fuse_1erfcx': <capsule object "double (double, int __pyx_skip_dispatch)" at 0x7fc717fd5870>,
 '__pyx_fuse_0erfi': <capsule object "__pyx_t_double_complex (__pyx_t_double_complex, int __pyx_skip_dispatch)" at 0x7fc717fd58a0>,
 '__pyx_fuse_1erfi': <capsule object "double (double, int __pyx_skip_dispatch)" at 0x7fc717fd58d0>}

Then you can choose the one corresponds to the signature you need.

[joernweissenborn]

Ok, so this gives me the info to construct the native call which numba cannot generate at the moment, if I understand correctly?

[ickc]

Yes, following the airy example above together with the function you find out then you can construct a jittable function providing the same result.

[joernweissenborn]

Ok, thanks for the clarification.

Just as a comment: This functions are crucial for real science. Numba really needs to support this out of the box.

[person142]

Numba really needs to support this out of the box.

Hm I don’t know that the burden should be entirely placed on Numba here because it’s unreasonable to expect it to reimplement the API of every major scientific Python package. At a certain point the community needs to figure out a better interop story.

[person142]

More concretely: I’m very interested in ideas as to how we could make scipy.special or scipy.special.cython_special more compatible with Numba. (I wrote cython_special, and remain interested in making it more broadly useful.)

[ickc]

Something like numba_special that has those wrap in a jit function already so that it can be used in numba jitted functions?

[seibert]

Making Numba able to compile SciPy math functions has been on our radar for a while (see #2109, @stuartarchibald even came up with the project name of "Scumba" 😄 ), but we've been trying to solidify our core functionality and NumPy support first. We would prefer to do it as a Numba extension (that we would be happy to host under the numba org).

On the topic of minimizing code, we have wanted for a long time to figure out how to call Cython functions from Numba as easily as we can call CFFI and ctypes wrapped functions. If that would allow us to more directly use interfaces from SciPy and not have to reimplement stuff, that would be great.

[astrojuanlu]

(Arrived here from SciPy-Dev)

FWIW, this is an experiment I did to wrap directly the underlying C code with CFFI and call it from numba:

https://github.com/poliastro/pycephes/tree/v0.1.0

(Explanatory article translated to English)

@horta then followed a similar approach in https://github.com/limix/hcephes.

[joernweissenborn]

Following the advises of @stuartarchibald and @ickc I solved my issues. A short question: how much likely is it that this solution will break with future updates, i.e. if scipy people make special more canonical?

I really got nice speedups in my case. Without parallel it's on par with cython, with double as fast in my use-case. So I can throw out the cython dependency and have much easier maintenance. Thanks to Numba team!

[seibert]

Not sure what the durability of this solution will be, but it has worked for some time now, and Cython extensions depend on this interface, so it is not likely to go away. We're having a discussion on the SciPy mailing list now to see if there is a more official way we could migrate toward.

[stuartarchibald]

Followers of this ticket may also be interested in scumba discussion here: https://github.com/numba/numba/issues/4292

[stuartarchibald]

The Numba extension project for SciPy: https://github.com/numba/numba-scipy is now the place to request for Numba support of SciPy functionality.