How to accelerate math calculations with Numba and Cython

Python is an interpretation that results in a pretty slow execution of a simple code. To accelerate it, it is possible to use some precompilation techniques.

In this tutorial, I will show you how to use Numba and Cython libraries to accelerate linear mathematical calculations up to 20 times.

Fibanacci numbers in python

I have already shown, how to calculate Fibonacci numbers. For this example, I will use very simple code in the file fib.py:


def fibonacci(n):
    fib1, fib2 = 0, 1
    for i in range (2, n + 1):
        fib1, fib2 = fib2, fib1 + fib2
    return fib2;


import timeit
from fib import fibonacci as p_fib

t_p = timeit.timeit(lambda: p_fib(90), number=5_000_000)

print(f"Python fib: {t_p:.2f}")

Numba optimisation

Numba coding

For linear, non-vectorised operations, Numba can do just-in-time optimisation while saving time on permanent interpretation of the source code.


import numba

@numba.njit
def fibonacci(n):
    fib1, fib2 = 0, 1
    for i in range (2, n + 1):
        fib1, fib2 = fib2, fib1 + fib2
    return fib2;

So, the only changes are to import library numba and then use @numba.njit in front of the function to be optimised.


import timeit
from fib import fibonacci as p_fib
from nfib import fibonacci as n_fib

t_p = timeit.timeit(lambda: p_fib(90), number=5_000_000)
t_n = timeit.timeit(lambda: n_fib(90), number=5_000_000)

print(f"Python fib: {t_p:.2f}")
print(f"Numba fib: {t_n:.2f}")
print(f"Acceleration: Numba:{t_p/t_n:.1f}x")

Numba acceleration

Cython optimisation

Optimisation for proper C Cython is more complicated and it requires some knowledge of C language.

Cython coding


# cython language_level=3, boundscheck=False, wraparound=False

cpdef fibonacci(int n):
    cdef long fib1 = 0, fib2 = 1
    cdef int i
    for i in range (2, n + 1):
        fib1, fib2 = fib2, fib1 + fib2
    return fib2;

It is important to have a proper C definition of all variables used. Then, it is necessary to compile it with the command:


cythonize -b -i -a cfib.pyx

The options mean that we will build in place our code with detailed annotations. Analysis of the annotation file will allow you to understand, how effectively your code is translated.


import timeit
from fib import fibonacci as p_fib
from nfib import fibonacci as n_fib
from cfib import fibonacci as c_fib

t_p = timeit.timeit(lambda: p_fib(90), number=5_000_000)
t_n = timeit.timeit(lambda: n_fib(90), number=5_000_000)
t_c = timeit.timeit(lambda: c_fib(90), number=5_000_000)

print(f"Python fib: {t_p:.2f}")
print(f"Numba fib: {t_n:.2f}")
print(f"Numba fib: {t_c:.2f}")
print(f"Acceleration: Numba:{t_p/t_n:.1f}x;  Cython:{t_p/t_c:.1f}x")

Cython acceleration

Python vs Numba vs Cython

Variables

It is important to remember, that in Python, variables can be very long. In Numba and Cython, variables are limited by the C definition!

In fact, Numba can work in Python and C mode: both jit and njit decorators can be used for JIT compilation, but njit is a shortcut specifically for cases where you want to ensure that the function is compiled in the nopython mode.

Choosing Between Numba and Cython:

For quick and easy speedups with minimal modifications, especially for numerical code, Numba may be preferred.

If you need more control over optimizations, integration with existing C code, or fine-tuned performance, Cython might be a better choice.

Ultimately, the choice between Numba and Cython depends on your specific use case, development preferences, and the nature of the code you're working with.