numbakit-anjit: leveraging numba to speed up ODE integration

numbakit-anjit (nbkanjit) to assist Numba intensive project by providing anjit, an annotation aware numba jit decorator.

It runs in Python 3.7+ depending on Numba. It is licensed under BSD.

It is extremely easy and natural to use:

>>> from numba import types as nt
>>> import nbkanjit
>>> @nbkanjit.anjit
... def func(x: nt.float64, y: nt.float64) -> nt.float64:
...     return x + y

You can also use Python types:

>>> @nbkanjit.anjit
... def func(x: float, y: float) -> float:
...     return x + y

which are mapped to numba types.

You can use

>>> from nbkanjit import Function as F_
>>> @nbkanjit.anjit
... def func1(x: int, y: float) -> float:
...     return x + y
>>> def func2(x: int, y: F_(funct1)._return) -> float:
...     return x + y

You can also use the annotation of any argument. For example, F_(func).x in this case is equivalent to int. Or even the full function F_(func) that will return FunctionType(float64(int, float64))

It also provides a manager to encapsulate (and reuse different parameters)

>>> import nbkanjit
>>> jm = nbkanjit.JitManager(cache=True)
>>> @jm.anjit
... def func(x: float, y:float) -> nt.float64:
...     return x + y

even to be applied in to the standard numba njit.

>>> jm = nbkanjit.JitManager(cache=True)
>>> @jm.njit
... def func(x, y):
...     return x + y

And you can teach the manager new tricks:

>>> jm.mapping["array1d"] = nt.float64[:]

by mapping any python object into a numba type.

Quick Installation

To install numbakit-anjit, simply (soon):

$ pip install numbakit-anjit

or utilizing conda, with the conda-forge channel (soon):

$ conda install -c conda-forge numbakit-anjit

and then simply enjoy it!

Why

Numba njit is awesome. Simple to use, produces the appropriate machine code once that the function is called. As the Numba docs says:

But numba also has an eager mode:

This can produce slightly faster code as the compiler does not need to infer the types. It also provides type check at definition time ensuring correctness. In numba intensive projects, this can be an useful trait. Finally, eager compilation is currently required to have two functions with the same signature to be arguments of a third one, without needing to recompile this last one in each case.

While developing numbakit-ode I was missing two things:

1. That eager compilation make use of function annotations
2. A global manager object to manipulate in one place numba jit options.

So, numbakit-anjit was born.

numbakit-anjit is maintained by a community. See AUTHORS for a complete list.

To review an ordered list of notable changes for each version of a project, see CHANGES