Notebook Texto y código sujeto bajo Creative Commons Attribution license, CC-BY-SA. (c) Original por Lorena A. Barba y Gilbert Forsyth en 2013, traducido por F.J. Navarro-Brull para CAChemE.org

Este notebook complementa el primer módulo interectivo online de CFD con Python y clases dadas por la profesora Lorena A. Barba, llamadas 12 Pasos para Navier-Stokes. En particular, abordaremos la cuestión del trabajar con Python en alto rendimiento.

Optimizando bucles con Numba¶

Recordarás nuestra exploración Operaciones matriciales (arrays) con NumPy donde se mostró como existen grandes incrementos de velocidad en la aplicación de nuestras discretizaciones mediante operaciones de arrays NumPy optimizadas (en lugar de muchos bucles anidados).

Numba es una herramienta que ofrece otro enfoque para la optimización de nuestro código Python. Numba es una librería de Python que convierte las funciones de Python en funciones compiladas en estilo C usando LLVM. Dependiendo del código original y el tamaño del problema, Numba puede proporcionar una aceleración significativa sobre código optimizado con NumPy.

Vamos a revisar la Ecuación de Laplace en 2D:

In [1]:

from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import matplotlib.pyplot as plt
import numpy as np

## Declaración de variables
nx = 81
ny = 81
c = 1
dx = 2.0/(nx-1)
dy = 2.0/(ny-1)

## Condiciones iniciales
p = np.zeros((ny,nx)) ##create a XxY vector of 0's

## Guías de la figura (malla)
x = np.linspace(0,2,nx)
y = np.linspace(0,1,ny)

## Condiciones de frontera
p[:,0] = 0		##p = 0 @ x = 0
p[:,-1] = y		##p = y @ x = 2
p[0,:] = p[1,:]		##dp/dy = 0 @ y = 0
p[-1,:] = p[-2,:]	##dp/dy = 0 @ y = 1

Esta es la función para iterar sobre la Ecuación de Laplace que escribimos en el Paso 9:

In [2]:

def laplace2d(p, y, dx, dy, l1norm_target):
    l1norm = 1
    pn = np.empty_like(p)

    while l1norm > l1norm_target:
        pn[:] = p[:]
        p[1:-1,1:-1] = (dy**2*(pn[2:,1:-1]+pn[0:-2,1:-1])+dx**2*(pn[1:-1,2:]+pn[1:-1,0:-2]))/(2*(dx**2+dy**2)) 
        p[0,0] = (dy**2*(pn[1,0]+pn[-1,0])+dx**2*(pn[0,1]+pn[0,-1]))/(2*(dx**2+dy**2))
        p[-1,-1] = (dy**2*(pn[0,-1]+pn[-2,-1])+dx**2*(pn[-1,0]+pn[-1,-2]))/(2*(dx**2+dy**2)) 
    
        p[:,0] = 0		##p = 0 @ x = 0
        p[:,-1] = y		##p = y @ x = 2
        p[0,:] = p[1,:]		##dp/dy = 0 @ y = 0
        p[-1,:] = p[-2,:]	##dp/dy = 0 @ y = 1
        l1norm = (np.sum(np.abs(p[:])-np.abs(pn[:])))/np.sum(np.abs(pn[:]))
     
    return p

Usemos el comando %%timeit (cell-magic) para ver lo rápido que se ejecuta:

In [28]:

%%timeit
laplace2d(p, y, dx, dy, .00001)

1 loops, best of 3: 206 us per loop

¡Vale! Nuestra función laplace2d tarda alrededor de 206 micro-segundos en completarse. Eso es muy rápido y tenemos que agradecérselo a las operaciones con arrays. Vamos a echar un vistazo al tiempo que tarda el uso de una versión Python más 'pura'.

In [29]:

def laplace2d_vanilla(p, y, dx, dy, l1norm_target):
    l1norm = 1
    pn = np.empty_like(p)
    nx, ny = len(y), len(y)

    while l1norm > l1norm_target:
        pn[:] = p[:]
        
        for i in range(1, nx-1):
            for j in range(1, ny-1):
                p[i,j] = (dy**2*(pn[i+1,j]+pn[i-1,j])+dx**2*(pn[i,j+1]-pn[i,j-1]))/(2*(dx**2+dy**2))
                          
        p[0,0] = (dy**2*(pn[1,0]+pn[-1,0])+dx**2*(pn[0,1]+pn[0,-1]))/(2*(dx**2+dy**2))
        p[-1,-1] = (dy**2*(pn[0,-1]+pn[-2,-1])+dx**2*(pn[-1,0]+pn[-1,-2]))/(2*(dx**2+dy**2)) 
    
        p[:,0] = 0		##p = 0 @ x = 0
        p[:,-1] = y		##p = y @ x = 2
        p[0,:] = p[1,:]		##dp/dy = 0 @ y = 0
        p[-1,:] = p[-2,:]	##dp/dy = 0 @ y = 1
        l1norm = (np.sum(np.abs(p[:])-np.abs(pn[:])))/np.sum(np.abs(pn[:]))
     
    return p

In [30]:

%%timeit
laplace2d_vanilla(p, y, dx, dy, .00001)

10 loops, best of 3: 32 ms per loop

La versión simple Python tarda 32 mili-segundos en completarse. Vamos a calcular el aumento de velocidad que ganamos en el uso de las operaciones array:

In [35]:

32*1e-3/(206*1e-6)

Out[35]:

155.33980582524273

Así NumPy nos da un aumento de velocidad 155x sobre el código normal de Python. Dicho esto, a veces poniendo en práctica nuestras discretizaciones en operaciones con arrays (vectorizadas) puede ser un poco complicado.

Vamos a ver lo que puede hacer Numba. Empecemos con la importación de la función especial autojit de la biblioteca numba:

In [1]:

from numba import autojit

Para integrar Numba con nuestra función existente, todo lo que tenemos que hacer es anteponer el comando @ autojit antes de nuestra declaración de def:

In [38]:

@autojit
def laplace2d_numba(p, y, dx, dy, l1norm_target):
    l1norm = 1
    pn = np.empty_like(p)

    while l1norm > l1norm_target:
        pn[:] = p[:]
        p[1:-1,1:-1] = (dy**2*(pn[2:,1:-1]+pn[0:-2,1:-1])+dx**2*(pn[1:-1,2:]+pn[1:-1,0:-2]))/(2*(dx**2+dy**2)) 
        p[0,0] = (dy**2*(pn[1,0]+pn[-1,0])+dx**2*(pn[0,1]+pn[0,-1]))/(2*(dx**2+dy**2))
        p[-1,-1] = (dy**2*(pn[0,-1]+pn[-2,-1])+dx**2*(pn[-1,0]+pn[-1,-2]))/(2*(dx**2+dy**2)) 
    
        p[:,0] = 0		##p = 0 @ x = 0
        p[:,-1] = y		##p = y @ x = 2
        p[0,:] = p[1,:]		##dp/dy = 0 @ y = 0
        p[-1,:] = p[-2,:]	##dp/dy = 0 @ y = 1
        l1norm = (np.sum(np.abs(p[:])-np.abs(pn[:])))/np.sum(np.abs(pn[:]))
     
    return p

Las únicas líneas que han cambiado son la línea @ autojit y también el nombre de la función, que se ha cambiado para que podamos comparar el rendimiento. Ahora veamos lo que sucede:

In [39]:

%%timeit
laplace2d_numba(p, y, dx, dy, .00001)

1 loops, best of 3: 137 us per loop

¡Ok! Así que no es un aumento de velocidad 155x como vimos entre la versión pura de Python y NumPy, pero es una ganancia nada trivial de rendimiento, especialmente teniendo en cuenta lo fácil que era ponerla en práctica. Otra característica interesante de Numba es que se puede usar el comando @autojit en funciones de operación no-array, también. Vamos a tratar de agregarlo a nuestra versión de Python pura:

In [41]:

@autojit
def laplace2d_vanilla_numba(p, y, dx, dy, l1norm_target):
    l1norm = 1
    pn = np.empty_like(p)
    nx, ny = len(y), len(y)

    while l1norm > l1norm_target:
        pn[:] = p[:]
        
        for i in range(1, nx-1):
            for j in range(1, ny-1):
                p[i,j] = (dy**2*(pn[i+1,j]+pn[i-1,j])+dx**2*(pn[i,j+1]-pn[i,j-1]))/(2*(dx**2+dy**2))
                          
        p[0,0] = (dy**2*(pn[1,0]+pn[-1,0])+dx**2*(pn[0,1]+pn[0,-1]))/(2*(dx**2+dy**2))
        p[-1,-1] = (dy**2*(pn[0,-1]+pn[-2,-1])+dx**2*(pn[-1,0]+pn[-1,-2]))/(2*(dx**2+dy**2)) 
    
        p[:,0] = 0		##p = 0 @ x = 0
        p[:,-1] = y		##p = y @ x = 2
        p[0,:] = p[1,:]		##dp/dy = 0 @ y = 0
        p[-1,:] = p[-2,:]	##dp/dy = 0 @ y = 1
        l1norm = (np.sum(np.abs(p[:])-np.abs(pn[:])))/np.sum(np.abs(pn[:]))
     
    return p

In [42]:

%%timeit
laplace2d_vanilla_numba(p, y, dx, dy, .00001)

1 loops, best of 3: 561 us per loop

561 micro-ssegundos. Eso no llega al 155x que vimos con NumPy, pero está cerca. Y todo lo que hicimos fue agregar una línea de código.

Así, tenemos:

Puro (Vanilla) Python: 32 milisegundos

NumPy Python: 206 microsegundos

Vanilla + Numba: 561 microsegundos

NumPy + Numba: 137 microsegundos

Es evidente que la combinación Numba + NumPy es la más rápida, pero la capacidad de optimizar rápidamente código con bucles anidados también puede ser muy útil en ciertas aplicaciones.

In [37]:

from IPython.core.display import HTML
def css_styling():
    styles = open("../styles/custom.css", "r").read()
    return HTML(styles)
css_styling()

Out[37]:

(La celda de arriba establece el formato de este notebook) Tomado originalmente de CamDavidsonPilon, @Cmrn_DP.)