Accelerating Your Python Code with Numba

Numba is a just-in-time (JIT) compiler for Python that allows developers to accelerate their Python code without having to write C or C++ extensions. It is an open-source project that is widely used in the scientific computing community to speed up numerical calculations.

In this article, we’ll explore what Numba is, how it works, and some examples of how it can be used to speed up Python code.

What is Numba?

Numba is a Python package that compiles Python code to machine code at runtime using the LLVM compiler infrastructure. This allows Python code to be executed at near-native speeds, without the need for manual code optimization or the use of external libraries.

Numba is designed to work with NumPy arrays and other numerical data types. It provides a collection of decorators and functions that can be used to generate optimized machine code for a given function or method. These decorators and functions can be used to accelerate Python code by replacing interpreted Python code with compiled machine code.

How does Numba work?

Numba works by generating machine code for a Python function or method at runtime. When a decorated function is called, Numba compiles the Python bytecode to LLVM intermediate representation (IR). It then optimizes the IR using LLVM optimization passes and generates machine code for the target CPU architecture.

This process can take some time, so Numba includes a cache system that stores the generated machine code in memory. This cache allows Numba to avoid re-compiling the same function multiple times, which can significantly speed up the execution of subsequent calls.

Numba provides several decorators and functions that can be used to control the behavior of the JIT compiler. For example, the @jit decorator can be used to enable JIT compilation for a given function, while the @vectorize decorator can be used to generate SIMD vectorized code for a given function.

Examples of using Numba

Let’s take a look at an example of using Numba to speed up a loop. We’ll use the same @jit decorator to enable JIT compilation for the loop.

import numpy as np
from numba import jit

@jit
def add_numbers(n):
    s = 0
    for i in range(n):
        s += i
    return s

result = add_numbers(1000000)

In this example, we define a function called add_numbers that takes an integer n as an argument and returns the sum of the integers from 0 to n. We use the @jit decorator to enable JIT compilation for the loop.

We then call the add_numbers function with an argument of 1,000,000. The function is compiled to machine code at runtime and the sum of the integers is returned.

Lets take a loon at an another example from original documentation:

from numba import jit
import numpy as np
import time

x = np.arange(100).reshape(10, 10)

@jit(nopython=True)
def go_fast(a): # Function is compiled and runs in machine code
    trace = 0.0
    for i in range(a.shape[0]):
        trace += np.tanh(a[i, i])
    return a + trace

# DO NOT REPORT THIS... COMPILATION TIME IS INCLUDED IN THE EXECUTION TIME!
start = time.perf_counter()
go_fast(x)
end = time.perf_counter()
print("Elapsed (with compilation) = {}s".format((end - start)))

# NOW THE FUNCTION IS COMPILED, RE-TIME IT EXECUTING FROM CACHE
start = time.perf_counter()
go_fast(x)
end = time.perf_counter()
print("Elapsed (after compilation) = {}s".format((end - start)))

Output:

Elapsed (with compilation) = 0.33030009269714355s
Elapsed (after compilation) = 6.67572021484375e-06s

Conclusion

Numba is a powerful tool that can greatly improve the performance of your Python code. By using Numba, you can avoid the need to write and maintain complex C or C++ extensions, and still achieve near-native speeds for numerical calculations and other computationally intensive tasks.

Numba is a popular choice for scientific computing and data analysis tasks, but it can be used in a wide range of applications where performance is critical. With its easy-to-use decorators and functions, it is accessible to both beginners and experienced developers alike.

In conclusion, Numba is a valuable addition to any Python developer’s toolkit, and can help unlock the full potential of Python as a high-performance computing language.