Use CFUNCTYPE with complex return value for scipy.special.jv + GPU

Numba Support: How do I do ...?
1

Hi,

I tried to follow this answer to call scipy.special.jv from numba: support for scipy.special functions : feature request · Issue #3086 · numba/numba · GitHub
I’m aware that this function is available through numba-scipy but I need a version supporting a complex argument as it is possible in scipy but not in numba-scipy. I ended up with something like this:

from numba.extending import get_cython_function_address
from numba import 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_0jv")
functype = ctypes.CFUNCTYPE(_dble, _dble, _dble, _dble)
jv_fn = functype(addr)


@njit
def numba_jv(v, z):
    real, imag = jv_fn(v, np.real(z), np.imag(z))
    return real + 1j * imag


x = -2.0 + 3.0j

numba_res = numba_jv(1.0, x)

Obviously this cannot work as the return value is complex and not double (first argument of CFUNCTYPE). But I don’t know how to define a complex return value there as it should be according to scipy docs: /doc/scipy/reference/special.cython_special.html#module-scipy.special.cython_special (sorry, I can’t include a link as a new user).

Optimally I would like to run this on a GPU (CUDA) using numba which isn’t supported by numba-scipy at all. Would there be a way to do this for this special case?

Thanks in advance!
Jan