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.
You can see the project here.
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.