Optional numba support #66
Description
Add contrib.numba, and use a noop @jit if not available on relevant functions.
Additionally, the scipy.integrate.quad integrand can be sped up with cfunc, see https://numba.readthedocs.io/en/stable/user/cfunc.html#example
In places where scipy.special functions are used, some trickery is needed.
For example, this snippet is used to make scipy.special.erfi work within numba jitted functions for np.float64 input:
# lmo/contrib/_numba.py
def _overload_scipy_special_erfi():
_erfi_f8 = ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double)(
numba.extending.get_cython_function_address(
'scipy.special.cython_special',
'__pyx_fuse_1erfi',
),
)
@numba.extending.overload(scipy.special.erfi)
def numba_erfi(*args):
match args:
case (numba.types.Float(),):
def _numba_erfi(*args):
return _erfi_f8(*args)
return _numba_erfi
case _:
return None
def overload_scipy_special():
_overload_scipy_special_erfi()# lmo/pyproject.toml
[project.entry-points.numba_extensions]
init = "lmo.contrib.numba:overload_scipy_special"