How to fix numba-scipy to work with scipy.special

I am trying to write a simulation that involves very large matrix multiplications and for loops. To speed up the process, I thought about using numba-scipy. After installing the package with conda, I tried running the following code

import scipy.special as sp
from numba import jit
@jit(nopython=True)
def Grh_gtr(J2,t1,t2):
t = t1-t2
if t == 0:
    return -1j/2
else:
    return (sp.jv(1,J2*t)-1j*sp.struve(1,J2*t))/(2j*J2*t)

I got instead, the following error

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<ufunc 'jv'>) found for signature:
 
 >>> jv(Literal[int](1), int64)

I have also tried installing the numba-special package using pip, but it gives me the following error

ERROR: Failed building wheel for numba-special
  Running setup.py clean for numba-special
Failed to build numba-special
ERROR: ERROR: Failed to build installable wheels for some pyproject.toml based projects (numba-special)

Any help is greatly appreciated.

Solution

Assuming you have successfully installed numba-scipy.

The first error complains about a type mismatch.

According to the documentation, the supported signature for jv is as follows.

Supported signature(s): float64(float64,float64)

On the other hand, your use of jv is interpreted as follows.

 >>> jv(Literal[int](1), int64)

The first argument is interpreted as a literal int, and the second argument is interpreted as an int64 variable. So an easy fix is to cast both to float.

@jit(nopython=True)
def Grh_gtr(J2, t1, t2):
    t = t1 - t2
    if t == 0:
        return -1j / 2
    return (sp.jv(float(1), float(J2 * t)) - 1j * sp.struve(float(1), float(J2 * t))) / (2j * J2 * t)

Note that casting from int to float is cheap, but it's not free. It might be better to minimize it by making it a variable.

import scipy.special as sp
from numba import jit

@jit(nopython=True)
def Grh_gtr(J2, t1, t2):
    t = t1 - t2
    if t == 0:
        return -1j / 2
    one = float(1.0)
    j2t = float(J2 * t)
    return (sp.jv(one, j2t) - 1j * sp.struve(one, j2t)) / (2j * j2t)


print(Grh_gtr(1.0, 2.0, 1.0))  # Also better to call with the float arguments if possible.