matplotlib animation

#!/usr/bin/env python
# -*- coding: utf-8 -*-

#importa bibliotecas
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation as anim

#dados do problema
L = 1 #comprimento da barra [m]
a = 9.7e-5 #difusão térmica [m²/s]
T0 = 800 #temp inicial [°C]
Tr = 0 #temp de resfriamento [°C]

tmax = 3600 #instante final [s]

N = 200 #tamanho da malha X
dx = L / N #Delta x
dt = dx**2 / (2 * a) #calculando o máximo Delta t apra a estabilidade

x = np.linspace(0, L, N+1) #criando a malha na barra
T = np.zeros(N+1) #array com as temperaturas num dados instante

print(f'Malha: {N} pontos')

T[:] = T0 #impondo a condição inicial
T[[0, -1]] = Tr #condição de contorno

#variáveis aux
f = a * dt/dx**2 #alpha Delta / delta x² ( = 0,5 se uso o máximo de Delta t)
fm = 1 - 2 * f

#parte gráfica e animação
fig = plt.figure() #gera figura
plt.ylim(0, 850) #limites para o eixo y
plt.xlabel('x [m]') #nome do eixo x
plt.ylabel('T [Cº]') #nome do eixo y

plt.plot(x, [T0]*(N+1),'--') #plota linha
perfil, = plt.plot(x, T)
tempo = plt.text(L/2, 820, f't=0 s') #plota  o tempo


def update(i):
    T[1:-1] = f * (T[:-2] + T[2:]) + fm * T[1:-1] #função da troca de calor
    perfil.set_ydata(T) #atualiza o gráfico
    tempo.set_text(f't = {i*dt:.2f} s') #atualiza o tempo
    return perfil, tempo

an = anim(fig, update, interval = 1, blit = True) #função de animação
plt.show() #fim