numbakit-anjit 0.2

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 and manager object to handle Jit configuration.

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.

And a way to register a signature as a template (tmpl):

>>> import nbkanjit
>>> jm = nbkanjit.JitManager()
>>> jm.register("nice", nt.float64((nt.float64, nt.float64)))

and then use it for non-annotated function by explicitly name:

>>> @jm.njit_tmpl("nice")
... def other_func(x, y):
...     return x + y

or using the name of the function:

>>> @jm.njit_tmpl
... def nice(x, y):
...     return x + y

You can register directly from a function:

>>> @jm.register("nice")
... def _(x: float, y:float) -> nt.float64:
...     pass

or again taking the function name:

>>> @jm.register
... def nice(x: float, y:float) -> nt.float64:
...     pass

$ pip install numbakit-anjit

$ conda install -c conda-forge numbakit-anjit

in [Lazy mode], compilation will be deferred until the first function
execution. Numba will infer the argument types at call time, and
generate optimized code based on this information. Numba will also
be able to compile separate specializations depending on the input
types.

In which you can also tell Numba the function signature you are expecting.
[..] In this case, the corresponding specialization will be compiled by the
decorator, and no other specialization will be allowed. This is useful
if you want fine-grained control over types chosen by the compiler (for
example, to use single-precision floats).