Est-il possible de vectoriser cette triple boucle en Python / Numpy ?

Question

Je cherche à accélérer mon code qui prend actuellement un peu plus d'une heure pour s'exécuter en Python / Numpy. La majeure partie du temps de calcul se produit dans la fonction ci-dessous.

Je tente de vectoriser Z, mais je trouve cela plutôt difficile pour une triple boucle for. Pourrais-je éventuellement implémenter la fonction numpy.diff quelque part? Jetez un oeil:

def MyFESolver(KK,D,r,Z):
    global tdim
    global xdim
    global q1
    global q2
    for k in range(1,tdim):
        for i in range(1,xdim-1):
            for j in range (1,xdim-1):
                Z[k,i,j]=Z[k-1,i,j]+r*q1*Z[k-1,i,j]*(KK-Z[k-1,i,j])+D*q2*(Z[k-1,i-1,j]-4*Z[k-1,i,j]+Z[k-1,i+1,j]+Z[k-1,i,j-1]+Z[k-1,i,j+1])
    return Z

tdim = 75 xdim = 25

Answer 1

Je suis d'accord, c'est délicat car les conditions aux limites (BCs) sur les quatre côtés ruinant la structure simple de la matrice de rigidité. Vous pouvez vous débarrasser des boucles d'espace comme suit :

from pylab import *
from scipy.sparse.lil import lil_matrix
tdim = 3;     xdim = 4;  r = 1.0;  q1, q2 = .05, .05; KK= 1.0; D = .5  #valeurs aléatoires
Z = ones((tdim, xdim, xdim))
#Parcours dans le temps
for k in range(1,tdim):
    Z_prev = Z[k-1,:,:] #peut avoir besoin d'être aplanie
    Z_up = Z_prev[1:-1,2:]
    Z_down = Z_prev[1:-1,:-2]

    Z_left = Z_prev[:-2,1:-1]
    Z_right = Z_prev[2:,1:-1]

    centre_term  = (q1*r*(Z_prev[1:-1,1:-1] + KK) - 4*D*q2)* Z_prev[1:-1,1:-1] 

    Z[k,1:-1,1:-1]= Z_prev[1:-1,1:-1]+ centre_term + q2*(Z_up+Z_left+Z_right+Z_down)

Mais je ne pense pas que vous puissiez vous débarrasser de la boucle temporelle...

Je pense que l'expression :

Z_up = Z_prev[1:-1,2:]

fait une copie en numpy, alors que ce que vous voulez est une vue - si vous pouvez comprendre comment faire cela - cela devrait être encore plus rapide (de combien?)

Enfin, je suis d'accord avec le reste des réponses - d'expérience, ce genre de boucles est mieux réalisé en C puis enveloppé dans numpy. Mais ce qui précède devrait être plus rapide que l'original...

Answer 2

Cela ressemble à un cas idéal pour Cython. Je suggérerais d'écrire cette fonction en Cython, elle sera probablement des centaines de fois plus rapide.