pyplot: realmente lento creando heatmaps

Question

Tengo un bucle que ejecuta el cuerpo alrededor de 200 veces. En cada iteración de ciclo, realiza un cálculo sofisticado, y luego como depuración, deseo producir un mapa de calor de una matriz NxM. Pero, generar este mapa de calor es insoportablemente lento y ralentiza significativamente un algoritmo ya lento.

Mi código es a lo largo de las líneas:

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 matriz, de numpy, no es enorme --- 300 x 600 de dobles. Incluso si no guardo la figura y en su lugar actualizo un gráfico en la pantalla, es incluso más lento.

Seguramente debo estar abusando de pyplot. (Matlab puede hacer esto, no hay problema.) ¿Cómo puedo acelerar esto?

Answer 1

trate de poner plt.clf() en el bucle para borrar la cifra actual:

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 no lo hace, el bucle disminuye a medida que la máquina se esfuerza por asignar más y más memoria de la figura.

Answer 2

Creo que esto es un poco más rápido:

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 llamadas gca que exige gcf que comprueba para ver si hay una figura; si no, crea uno. Al instanciar manualmente la figura primero, no necesita hacer todo eso.