Python lent sur-boucles et des centaines de recherches d'attributs. Utilisez Numba?

Question

Je suis en train de travailler sur une mise en œuvre simple en SPH (hydrodynamique des particules lissées, pas pertinent ici) en python. Le code fonctionne, mais l'exécution est un peu lente. Je dois souvent comparer des particules individuelles avec un certain nombre de voisins. Dans une implémentation antérieure, je gardais toutes les positions des particules et toutes les distances-à-chaque-particule existante dans de grands tableaux numpy -> jusqu'à un certain point c'était assez rapide. Mais visuellement pas agréable et n ** 2. Maintenant je le veux propre et simple avec les classes + kdTree pour accélérer la recherche du voisin.

Tout cela se passe dans ma classe de simulation globale. De plus, il y a une classe appelée "particule" qui contient toutes les informations individuelles. Je crée des centaines d'instances avant de les boucler.

def calculate_density(self): 
    #Using scipys advanced nearest neighbour seach magic 
    tree = scipy.spatial.KDTree(self.particle_positions) 

    #here we go... loop through all existing particles. set attributes.. 
    for particle in self.my_particles: 

     #get the indexes for the nearest neighbours 
     particle.index_neighbours = tree.query_ball_point(particle.position,self.h,p=2) 


     #now loop through the list of neighbours and perform some additional math 
     particle.density = 0 
     for neighbour in particle.index_neighbours: 


      r = np.linalg.norm(particle.position - self.my_particles[neighbour].position) 
      particle.density += particle.mass * (315/(64*math.pi*self.h**9)) *(self.h**2-r**2)**3 

J'ai chronométré 0,2717630863189697 pour seulement 216 particules.

Maintenant, je me demande: que faire pour l'accélérer? La plupart des outils en ligne comme "Numba" montrent comment ils accélèrent les fonctions individuelles mathématiques. Je ne sais pas lequel choisir. Sur un sidenode, je ne peux même pas faire fonctionner Numba dans ce cas. Je reçois un message d'erreur looong. Et j'espérais que c'était aussi simple que de claquer "@jit" devant.

Je connais ses boucles avec les appels d'attribut qui écrasent ma performance de toute façon - pas la recherche mathématique ou la recherche de voisin. Malheureusement, je suis un novice de la programmation et j'ai aimé l'approche propre que j'ai eu à travailler ici :(des pensées?

Answer 1

Ces types de calculs intensifs en boucle sont lents en Python.Ainsi, la première chose que vous voulez faire est Pour voir si vous pouvez vectoriser ces opérations et vous débarrasser des boucles, les calculs seront effectués dans les bibliothèques C ou Fortran et vous obtiendrez beaucoup de rapidité, si vous pouvez le faire, c'est la voie à suivre. Cependant, certaines opérations sont simplement intrinsèquement gourmandes en boucle, dans ce cas, l'utilisation de Cython vous aidera beaucoup - vous pouvez généralement espérer une accélération de 60X + lorsque vous cytonisez votre boucle. aussi eu des expériences similaires avec numba - quand ma fonction devient compliquée, il a échoué à faire fa ster, donc d'habitude je viens d'utiliser Cython.

Codage en Cython n'est pas trop mauvais - beaucoup plus facile que le code en C parce que vous pouvez accéder à des tableaux numpy facilement via memoryviews. Un autre avantage est qu'il est assez facile de paralléliser la boucle avec openMP, ce qui peut vous donner des accélérations supplémentaires de 4X + (bien sûr, selon le nombre de cœurs de votre machine), donc votre code peut être des centaines de fois plus rapide. Un problème est que pour obtenir la vitesse optimale, vous devez supprimer tous les appels python dans votre boucle, ce qui signifie que vous ne pouvez pas appeler les fonctions numpy/scipy. Donc, vous devez convertir tree.query_ball_point et np.linalg.norm partie en Cython pour une vitesse optimale.