Package: numbakit-anjit
numbakit-anjit: leveraging numba to speed up ODE integrationnumbakit-anjit (nbkanjit) to assistNumbaintensive project
by providinganjit, an annotation aware numba jit decorator and
manager object to handle Jit configuration.It runs in Python 3.7+ depending onNumba. It is licensed under
BSD.It is extremely easy and natural to use:>>>fromnumbaimporttypesasnt>>>importnbkanjit>>>@nbkanjit.anjit...deffunc(x:nt.float64,y:nt.float64)->nt.float64:...returnx+yYou can also use Python types:>>>@nbkanjit.anjit...deffunc(x:float,y:float)->float:...returnx+ywhich are mapped to numba types.You can use:>>>fromnbkanjitimportFunctionasF_>>>@nbkanjit.anjit...deffunc1(x:int,y:float)->float:...returnx+y>>>deffunc2(x:int,y:F_(funct1)._return)->float:...returnx+yYou can also use the annotation of any argument. For example,F_(func).xin this case is equivalent toint. Or even the
full functionF_(func)that will returnFunctionType(float64(int, float64))It also provides a manager to encapsulate (and reuse different parameters)>>>importnbkanjit>>>jm=nbkanjit.JitManager(cache=True)>>>@jm.anjit...deffunc(x:float,y:float)->nt.float64:...returnx+yeven to be applied in to the standard numba njit.>>>jm=nbkanjit.JitManager(cache=True)>>>@jm.njit...deffunc(x,y):...returnx+yAnd 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):>>>importnbkanjit>>>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")...defother_func(x,y):...returnx+yor using the name of the function:>>>@jm.njit_tmpl...defnice(x,y):...returnx+yYou can register directly from a function:>>>@jm.register("nice")...def_(x:float,y:float)->nt.float64:...passor again taking the function name:>>>@jm.register...defnice(x:float,y:float)->nt.float64:...passQuick InstallationTo install numbakit-anjit, simply (soon):$pipinstallnumbakit-anjitor utilizing conda, with the conda-forge channel (soon):$condainstall-cconda-forgenumbakit-anjitand then simply enjoy it!WhyNumbanjitis awesome. Simple to use, produces the appropriate machine code
once that the function is called. As theNumba docssays: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.But numba also has aneager mode: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).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.Another think we like aboutnjitis that is highly configurable using keyword
arguments and even some configurations could be applied globally using env
variables.While developingnumbakit-odeI was missing two things:That eager compilation make use of function annotationsA way to manipulatenjitoptions in a centralized but
granular mannerSo,numbakit-anjitwas born.numbakit-anjit is maintained by a community. SeeAUTHORSfor a complete list.To review an ordered list of notable changes for each version of a project,
seeCHANGES