Skip to main content

Annotation aware numba njit.

Project description

Latest Version License Python Versions https://github.com/hgrecco/numbakit-anjit/workflows/CI/badge.svg?branch=main https://github.com/hgrecco/numbakit-anjit/workflows/Lint/badge.svg?branch=main https://coveralls.io/repos/github/hgrecco/numbakit-anjit/badge.svg?branch=main

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

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:

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 an eager 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 about njit is that is highly configurable using keyword arguments and even some configurations could be applied globally using env variables.

While developing numbakit-ode I was missing two things:

  1. That eager compilation make use of function annotations

  2. A way to manipulate njit options in a centralized but granular manner

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

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

numbakit-anjit-0.2.tar.gz (14.2 kB view details)

Uploaded Source

Built Distribution

numbakit_anjit-0.2-py2.py3-none-any.whl (12.8 kB view details)

Uploaded Python 2 Python 3

File details

Details for the file numbakit-anjit-0.2.tar.gz.

File metadata

  • Download URL: numbakit-anjit-0.2.tar.gz
  • Upload date:
  • Size: 14.2 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/1.12.1 pkginfo/1.5.0.1 requests/2.23.0 setuptools/46.1.1.post20200323 requests-toolbelt/0.8.0 tqdm/4.43.0 CPython/3.6.10

File hashes

Hashes for numbakit-anjit-0.2.tar.gz
Algorithm Hash digest
SHA256 ed0b392bc986f267d21ecbf88d313f233c257361a974d0578faa3040cf60a311
MD5 5931080623e2929f22489f902ab9b5cb
BLAKE2b-256 b8a81f9fb02bea0aedfdf1aeacf3731a5e3c9c5510912a0adcf7e4afffb42ebd

See more details on using hashes here.

File details

Details for the file numbakit_anjit-0.2-py2.py3-none-any.whl.

File metadata

  • Download URL: numbakit_anjit-0.2-py2.py3-none-any.whl
  • Upload date:
  • Size: 12.8 kB
  • Tags: Python 2, Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/1.12.1 pkginfo/1.5.0.1 requests/2.23.0 setuptools/46.1.1.post20200323 requests-toolbelt/0.8.0 tqdm/4.43.0 CPython/3.6.10

File hashes

Hashes for numbakit_anjit-0.2-py2.py3-none-any.whl
Algorithm Hash digest
SHA256 3f59995d4949ea7f5f8d9d99aacb4dfb8b106f4cec1a6b5754f0272adf256681
MD5 2ee40eda6da8f74fc8763b11b56fbcdd
BLAKE2b-256 a26afee624312a405ee63061bb67a40c47145f14d416f270e8a59c40710353bb

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page