Acelera el código mediante vectorización en Python: elimina los bucles para aumentar el rendimiento hasta 1800 veces

La primera opción que suele venir a la cabeza cuando se necesita realizar una misma operación sobre diferentes valores es mediante el uso de un bucle. Lo que en Python se puede hacer mediante el uso de for o while. Esta es una forma natural de operar, primero se ejecuta la instrucción sobre el primer elemento, después sobre el segundo y así hasta que se termina. Si se conoce el número de pasos necesario para terminar lo habitual es usar for, en caso contrario while. Pero en lenguajes modernos, como es el caso de Python, si se trabaja con un gran volumen de datos, cuando es posible, suele ser mejor vectorizar las operaciones para aumentar el rendimiento. Las matrices de NumPy y los DataFrame de Pandas ofrecen esta opción, por lo que es necesario conocer cómo funcionan y las ventajas que ofrece. En esta publicación se va a explicar cómo se puede utilizar la vectorización en Python y las increíbles mejoras de rendimiento que ofrece.

¿Qué es la vectorización?

La vectorización es una técnica que se puede aplicar sobre las matrices de NumPy y los DataFrame de Pandas mediante la cual se aplica una misma operación sobre un conjunto de datos. Pudiendo aplicar sobre todo la matriz o DataFrame o solo sobre una parte. Esta técnica se aprende cuando se empieza a trabajar con las librerías, aunque no se le suele dar la importancia que merece. Por ejemplo, si se desea multiplicar una matriz de NumPy por dos no hay que escribir un bucle for, solamente hay que escribir algo como vec * 2, donde vec es un ndarray.

El código vectorizado no solo es más compacto que usar un bucle for, y por lo tanto, más fácil de leer, sino que también es generalmente más rápido. Algo que se nota especialmente al trabajar con grandes conjuntos de datos. Veamos algunos ejemplos de las mejoras de rendimiento que se pueden obtener mediante el uso de la vectorización en Python.

Sumar los valores de un vector

Si se tiene un vector con miles o millones de valores y se desea obtener la suma de todos ellos se puede hacer con un bucle for. Una posible forma de hacer esto es la siguiente.

import time

# Inicia el contador de tiempo
start = time.time()

# Suma en un bucle los valores entre 0 y limit - 1
limit = 50000000
result = 0
for item in range(0, limit):
    result += item
    
print(f'La suma entre 0 y {limit - 1} es {result}')

# Calcula el tiempo transcurrido
end = time.time()
print(end - start)

La suma entre 0 y 49999999 es 1249999975000000
2.1216330528259277

Lo que ha tardado unos 2,12 segundos. En NumPy existe la función sum que puede realizar esta operación mucho más rápido. Simplemente se tiene que escribir np.sum(np.arange(limit)). Con el siguiente código se puede comprobar que esta función es mucho más rápida que el bucle del ejemplo anterior.

import numpy as np

# Inicia el contador de tiempo
start = time.time()

# Suma en un bucle los valores entre 0 y limit - 1
limit = 50000000
result = np.sum(np.arange(limit))

print(f'La suma entre 0 y {limit - 1} es {result}')

# Calcula el tiempo transcurrido
end = time.time()
print(end - start)

La suma entre 0 y 49999999 es 1249999975000000
0.1663508415222168

Ahora solo tarda 0,16 segundos, consiguiendo mejorar el rendimiento en un factor 12. Aunque el factor exacto puede variar, lo normal es obtener una mejora de un orden de magnitud. En la práctica, esto se traduce en reducir de minutos a segundos o, como es el caso, de segundo a décimas de segundos el tiempo necesario para finalizar la operación.

Realizar operaciones sobre las filas de un DataFrame

La vectorización también ofrece importantes mejoras de rendimiento al realizar operaciones sobre las filas de los DataFrame. Por ejemplo, se puede crear un dataFrame aleatorio con el siguiente código.

import pandas as pd

rows = 100000
cols = 4

np.random.seed(0)

df = pd.DataFrame(np.random.rand(rows, cols), columns=['A', 'B', 'C', 'D'])

df.head()

          A         B         C         D
0  0.548814  0.715189  0.602763  0.544883
1  0.423655  0.645894  0.437587  0.891773
2  0.963663  0.383442  0.791725  0.528895
3  0.568045  0.925597  0.071036  0.087129
4  0.020218  0.832620  0.778157  0.870012

Ahora se puede sumar todos los elementos de cada una de las filas mediante un bucle for. Algo que puede implementar como se muestra a continuación.

# Inicia el contador de tiempo
start = time.time()

df['result'] = 0
for idx, row in df.iterrows():
    df.at[idx, 'result'] = row["A"] + row["B"] + row["C"] + row["D"]

# Calcula el tiempo transcurrido
end = time.time()
print(end - start)

df.head()

1.5942940711975098
          A         B         C         D    result
0  0.548814  0.715189  0.602763  0.544883  2.411649
1  0.423655  0.645894  0.437587  0.891773  2.398909
2  0.963663  0.383442  0.791725  0.528895  2.667724
3  0.568045  0.925597  0.071036  0.087129  1.651807
4  0.020218  0.832620  0.778157  0.870012  2.501007

En este caso el tiempo que ha necesitado el programa para terminar es de 1,59 segundos. La forma de vectorizar este bucle también es bastante sencilla, solamente se tiene que escribir una línea.

# Inicia el contador de tiempo
start = time.time()

df['result'] = df["A"] + df["B"] + df["C"] + df["D"]

# Calcula el tiempo transcurrido
end = time.time()
print(end - start)

df.head()

0.0008599758148193359
          A         B         C         D    result
0  0.548814  0.715189  0.602763  0.544883  2.411649
1  0.423655  0.645894  0.437587  0.891773  2.398909
2  0.963663  0.383442  0.791725  0.528895  2.667724
3  0.568045  0.925597  0.071036  0.087129  1.651807
4  0.020218  0.832620  0.778157  0.870012  2.501007

Lo que ha reducido el tiempo de ejecución a solamente 0,00086 segundos. Obteniendo un factor de mejora aún mayor que en el caso anterior: 1853. ¡Una mejora de tres órdenes de magnitud! Además, con un código más compacto y fácil de leer y entender.

Operaciones condicionales en un DataFrame

La vectorización también se puede aplicar cuando se necesita realizar una operación diferente dependiendo de los valores de los elementos de una fila. Un ejemplo sencillo puede ser usar un valor si el elemento de la columna A es menor que 0,5, en caso contrario otro valor si B es menor que 0,5 y, finalmente en caso contrario, otro. Este ejemplo tan sencillo es lo que se muestra a continuación.

# Inicia el contador de tiempo
start = time.time()

for idx, row in df.iterrows():
    if row['A'] > 0.5:
        df.at[idx, 'result'] = 1
    elif (row['B'] > 0.5):
        df.at[idx, 'result'] = 2
    else:
        df.at[idx, 'result'] = 3
        
# Calcula el tiempo transcurrido
end = time.time()
print(end - start)

df.head()

1.4399909973144531
          A         B         C         D  result
0  0.548814  0.715189  0.602763  0.544883     1.0
1  0.423655  0.645894  0.437587  0.891773     2.0
2  0.963663  0.383442  0.791725  0.528895     1.0
3  0.568045  0.925597  0.071036  0.087129     1.0
4  0.020218  0.832620  0.778157  0.870012     2.0

Aunque el problema no es muy útil si que se puede ver que hacer esto requiere 1,44 segundos. Algo que se puede hacer más rápido mediante vectorización.

# Inicia el contador de tiempo
start = time.time()

df.loc['result'] = 3
df.loc[row['B'] > 0.5, 'result'] = 2
df.loc[row['A'] > 0.5, 'result'] = 1

# Calcula el tiempo transcurrido
end = time.time()
print(end - start)

df.head()

0.02192378044128418
          A         B         C         D  result
0  0.548814  0.715189  0.602763  0.544883     1.0
1  0.423655  0.645894  0.437587  0.891773     2.0
2  0.963663  0.383442  0.791725  0.528895     1.0
3  0.568045  0.925597  0.071036  0.087129     1.0
4  0.020218  0.832620  0.778157  0.870012     2.0

Ahora solamente es necesario 0,021 segundos, mejorando el rendimiento en un factor 65.

Este ejemplo puede necesitar una pequeña explicación. Lo que se ha hecho es invertir las condiciones, en primer lugar, se ha asignado a todos los registros de la columna 'result' el valor asignado a los casos que no cumple la primera ni la segunda condición. Luego, cuando se cumple la segunda condición se asigna el valor correspondiente, sobreescribiendo el anterior. Finalmente se hace lo propio con la primera condición. Como se sobrescriben los valores es necesario aplicarlos en el orden contrario al usado en el bucle for.

Conclusiones

En esta publicación se ha visto cómo es posible aumentar la velocidad del código mediante vectorización en Python. Llegando a mejorar en más de un factor 1800 en algunos casos. Cuando se trabaja con pequeños volúmenes de datos, se puede usar bucles for por simplicidad, pero si no es así, es conveniente optimizar el código para introducir la vectorización.

La vectorización es clave para poder optimizar el código en Python, por eso en publicaciones anteriores ya se han visto otros ejemplos con los que se pude mejorar el rendimiento. La publicación de hoy expande esta con más casos de uso.

Image by Hans from Pixabay