pyplot: très lent création de heatmaps

Question

J'ai une boucle qui exécute le corps environ 200 fois. Dans chaque itération de boucle, il effectue un calcul sophistiqué, puis en débogage, je souhaite produire une carte thermique d'une matrice NxM. Mais, générer cette heatmap est insupportablement lent et ralentit considérablement un algorithme déjà lent.

Mon code est le long des lignes:

import numpy 
import matplotlib.pyplot as plt 
for i in range(200): 
    matrix = complex_calculation() 
    plt.set_cmap("gray") 
    plt.imshow(matrix) 
    plt.savefig("frame{0}.png".format(i)) 

La matrice, de numpy, n'est pas énorme --- 300 x 600 doubles. Même si je ne sauvegarde pas la figure et que je mets à jour un tracé à l'écran, c'est encore plus lent.

Sûrement je dois abuser de pyplot. (Matlab peut le faire, pas de problème.) Comment accélérer cela?

Answer 1

Essayez de mettre plt.clf() dans la boucle pour effacer le chiffre actuel:

for i in range(200): 
    matrix = complex_calculation() 
    plt.set_cmap("gray") 
    plt.imshow(matrix) 
    plt.savefig("frame{0}.png".format(i)) 
    plt.clf() 

Si vous ne le faites pas, la boucle ralentit que la machine se démène pour allouer de plus en plus de mémoire pour la figure.

Answer 2

Je pense que cela est un peu plus rapide:

import matplotlib.pyplot as plt 
from matplotlib import cm 
fig = plt.figure() 
ax = fig.add_axes([0.1,0.1,0.8,0.8]) 
for i in range(200): 
    matrix = complex_calculation() 
    ax.imshow(matrix, cmap=cm.gray) 
    fig.savefig("frame{0}.png".format(i)) 

plt.imshow appels gca qui appelle gcf qui vérifie s'il y a une figure; sinon, il en crée un. En instanciant manuellement la figure d'abord, vous n'avez pas besoin de faire tout cela.