Optimizando Python con la librería Numba

Como comentamos en el post anterior, Python es un lenguaje de programación muy utilizado hoy en día, entre otras razones porque es cómodo para programar a alto nivel, tiene muchas librerías y habilita código fácilmente leíble y mantenible. A pesar de todas estas ventajas, cabe destacar que es notoriamente lento, por lo que es conveniente conocer algunas herramientas que permitan acelerar el código en Python.

Python versus Numba

Una de las formas en que se puede acelerar Python es usando la librería Numba. Para ilustrar esto, se tomará como ejemplo el cálculo del número π utilizando el método de Monte Carlo, de forma similar a nuestro anterior post (ejemplo extraído de la página de Numba).

Acelerar Phyton
Acelerar Phyton

El código es muy sencillo. Se calcula el número π muestrando un número aleatorio de puntos y utilizando el ratio entre el área de un cuadrado y un círculo. El área del círculo es estimada como el número de puntos dentro del círculo, y la del cuadrado como el total de puntos.

Montecarlo
Montecarlo

Utilizando 100.000 puntos, la función de Python tarda 30.1ms, mientras que usando Numba tarda 896μms, es decir, 33.55x veces más rápido.

Pero… ¿Qué es Numba?

Numba es una librería de Python que da la capacidad de compilar el código tan pronto como se ejecuta (es decir, es un compilador Just-in-Time), consiguiendo así velocidades propias de código en C sin salir de la simplicidad de Python.

Y todo esto haciendo modificaciones mínimas sobre el código en Python, como se puede ver en el ejemplo. Además de compilar el código, Numba tiene la capacidad de detectar el tipo de procesador sobre el que se está ejecutando y hacer uso de optimizaciones específicas del hardware, además de usar diferentes threads (cosa que por sí solo no puede hacer Python).

Por si eso fuera poco, Numba está diseñado para trabajar con Numpy, una de las más famosas librerías de Python para computación numérica.

Uno de los inconvenientes de Numba es que no es amigable con estructuras de datos dinámicas (es decir, que pueden incrementar o reducir su tamaño arbitrariamente) tales como adición a listas o diccionarios, además de que puede ser complicado de debuguear. Además, todo código escrito con Numba requiere de un sistema que soporte su compilador JIT.

¿Cómo usar Numba?

Para que un programa de Python que usa Numba pueda funcionar se necesitan cuatro cosas, además de tener Numba instalado, claro.

(1) Importar la librería Numba desde el script, tal que así: import numba.

(2) Que el código que se desea acelerar esté claramente definido en una función. Encima de dicha función se añadirá el decorador @numba.njit. Este decorador indica a Numba que es esa función la que debe intentar compilar. Si por algún motivo no la puede compilar devolverá un error. Si no se quiere que devuelva un error y se prefiere que haga un fallback a Python normal, entonces el decorador se debe cambiar por @numba.jit.

(3) Que el código no contenga estructuras de datos dinámicas, y preferiblemente que las estructuras de datos que haya sean propias de Numpy. Es decir, cosas tales como .append() a una lista deben ser evitadas.

(4) Que los bucles no sean del tipo for element in elements, sino que se itere por los índices, tal como for i in range(len(elements)).

Como se ha podido observar, usar Numba para partes del código que puedan ser cuellos de botella desde un punto de vista computacional puede proporcionar grandes incrementos en velocidad, y a cambio de muy poco esfuerzo.

Esperamos que te haya sido útil este artículo. Si tienes entre manos algún proyecto de ingeniería y crees que podemos ayudarte, aquí te dejo el enlace donde nos podrás contactar y explicarnos más a cerca de ello.