Easy Python Speedup Wins With Numba
By: Sebastian Witowski
Tags: Coding, Software Development, Performance, Programming, Python

If you have functions that do a lot of mathematical operations, use NumPy or rely heavily on loops, then there is a way to speed them up significantly with one line of code. Ok, two lines if you count the import.

Numba and the @jit decorator

Meet Numba and its @jit decorator. It changes how your code is compiled, often improving its performance. You don’t have to install any special tools (just the numba pip package), you don't have to tweak any parameters. All you have to do is:

Add the @jit decorator to a function

decorator to a function Check if it’s faster

Let’s see an example of code before and after applying Numba 's optimization.

Before

The only purpose of this code is to do some calculations and to “be slow.” Let’s see how slow (benchmarks are done with Python 3.8 — I describe the whole setup in the Introduction article on my blog):

Now, we add @jit to our code. The body of the function stays the same, and the only difference is the decorator. Don't forget to install Numba package with pip ( pip install numba ).

… and after. Can you spot 2 lines that have changed?

Let’s measure the execution time once more:

Using @jit decorator gave us a 120x speedup (217 / 1.76 = 123.295)! That’s a huge improvement for such a simple change!

Other features of Numba

@jit is the most common decorator from the Numba library, but there are others that you can use:

@njit — alias for @jit(nopython=True). In nopython mode, Numba tries to run your code without using the Python interpreter at all. It can lead to even bigger speed improvements, but it's also possible that the compilation will fail in this mode.

mode, Numba tries to run your code without using the Python interpreter at all. It can lead to even bigger speed improvements, but it's also possible that the compilation will fail in this mode. @vectorize and @guvectorize — produces ufunc and generalized ufunc used in NumPy.

and generalized used in NumPy. @jitclass — can be used to decorate the whole class.

@cfunc — declares a function to be used as a native callback (from C or C++ code).

There are also advanced features that let you, for example, run your code on GPU with @cuda.jit. This doesn’t work out of the box, but it might be worth the effort for some very computational-heavy operations.

Numba has plenty of configuration options that will further improve your code’s execution time if you know what you are doing. You can:

Disable GIL (Global Interpreter Lock) with nogil

Cache results with cache

Automatically parallelize functions with parallel .

Check out the documentation to see what you can do. And to see more real-life examples (like computing the Black-Scholes model or the Lennard-Jones potential), visit the Numba Examples page.

Conclusions

Numba is a great library that can significantly speed up your programs with minimal effort. Given that it takes less than a minute to install and decorate some slow functions, it's one of the first solutions that you can check when you want to quickly improve your code (without rewriting it).

It works best if your code: