We recently released numbakit-ode, a package that leverages numba to speed up ODE integration. Basically, we rewrote SciPy code in a numba compatible manner and got a 10X performance boost in integration each step. Multiple steps are performed in tight loop resulting in an even better performance.
But we still have a large overhead due to compilation of the integrator stepper function that it would be nice to remove by caching the compilation results. We are not able to make that work and I have wrote this simple example to show the problem.
import numba as nb
@nb.njit
def func(t, y):
return y
@nb.njit(cache=True)
def stepper(f, a, b):
return 3 * f(a, b)
print(stepper(func, 4., 2.))
which results in
Traceback (most recent call last):
[...]
data_name = overloads[key]
KeyError: [...]
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
[...]
return pickle.dumps(obj, protocol=-1)
TypeError: cannot pickle 'weakref' object
My intuition (albeit without enough knowledge of the compilation process) is that this could work as compiling stepper only requires the signature of func, which should be cacheable. I wonder if this is just current limitation of Numba or there is a fundamental reason why this will not work.
@astrojuanlu happy to hear this!. On the weekend we will release version 0.3 which settles the public API. We are looking forward to seeing numbakit-ode being use.
Out of curiosity: @hgrecco are you finding performance issues like the one I described in https://github.com/numba/numba/issues/2952 ? Or are the objective functions you’re testing always the bottleneck?
@stuartarchibald Indeed, @astrojuanlu expanded the original scope of this post to include issue 2952 about dispatching (which I was unaware about the discussion in github).
I think that it makes sense as there are a few things that hit the performance of functions that take other functions as arguments, and therefore conspire against a wider adoption of numba in libraries that, like numbakit-ode, aim to provide generic algorithms.
Until now, I have identified:
Lack of cache for compiled functions that take other functions
Lack of function pointers, i.e. that only the signature matters
Lack of similar dispatching performance when arguments are functions.