Différents facteurs font qu'un code numérique écrit en Python peut souffrir de certaines lenteurs
On s'intéresse ici à la résolution des équations de Saint-Venant par une méthode de Lattice Boltzmann appelée $D_2Q_9$.
Elle est constituée de 3 étapes simples
def transport(f):
f[:, 1:, 1] = f[:, :-1, 1]
f[1:, :, 2] = f[:-1, :, 2]
...
def collision(f, m):
f2m(f, m)
relaxation(m)
m2f(m, f)
def relaxation(m):
m[:, :, 3] += c*(-2*m[:, :, 0] + 3.0*m[:, :, 1]**2 + 3.0*m[:, :, 2]**2 - m[:, :, 3])
m[:, :, 4] += c*(m[:, :, 0] + 1.5*m[:, :, 1]**2 + 1.5*m[:, :, 2]**2 - m[:, :, 4])
...
def one_time_step(f, m):
periodic_bc(f)
transport(f)
f2m(f, m)
relaxation(m)
m2f(m, f)
Définir des types statiques permettant à Cython de comprendre que nous ne sommes plus dans une partie Python mais dans une partie pouvant être écrite facilement en C et donc optimisée.
def transport(f):
f[:, 1:, 1] = f[:, :-1, 1]
...
def transport(double[:, :, ::1] f):
f[:, 1:, 1] = f[:, :-1, 1]
...
def transport(double[:, :, ::1] f):
cdef:
int i, j
int nx = f.shape[0]
int ny = f.shape[1]
for i in xrange(nx-1, 0, -1):
for j in xrange(ny):
f[i, j, 1] = f[i-1, j, 1]
...
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
def transport(double[:, :, ::1] f):
cdef:
int i, j
int nx = f.shape[0]
int ny = f.shape[1]
for i in xrange(nx-1, 0, -1):
for j in xrange(ny):
f[i, j, 1] = f[i-1, j, 1]
...
from parakeet import jit
@jit
def transport(f):
...
from numba import jit
@jit("void(f8[:,:,:])")
def transport(f):
...
#pythran export transport(float[][][])
def transport(f):
...
On prend notre schéma $D_2Q_9$ sur une grille de taille $1024\times 1024$ et on réalise à chaque fois 100 essais sur les fonctions testées. On calcule le temps moyen.
Versions utilisées
Attention: le stockage de f et de m est (ns, nx, ny) pour NumPy et (nx, ny, ns) pour Numba, parakeet, pythran et Cython.
On reprend la phase de transport.
def one_time_step(f1, f2):
nx, ny, ns = f1.shape
floc = np.zeros(ns)
mloc = np.zeros(ns)
periodic_bc(f1)
for i in range(1, nx-1):
for j in range(1, ny-1):
getf(f1, floc, i, j)
f2m_loc(floc, mloc)
relaxation_loc(mloc)
m2f_loc(mloc, floc)
setf(f2, floc, i, j)