Heat Equation part 2 a slight modification. Full article can be found at http://www.firsttimeprogrammer.blogspot.com/2015/07/heat-equation-part-2-slight-modification.html

heatEquation2.py

import numpy as np
from numpy import pi
import matplotlib.pyplot as plt
import matplotlib.animation as animation

fig = plt.figure()
fig.set_dpi(100)
ax1 = fig.add_subplot(1,1,1)

#Diffusion constant
k = 2

#Scaling factor
scale = 5

#Length of the rod (0,L) on the x axis
L = 3*pi

#Initial contitions
x0 = np.linspace(0,L,10000)
t0 = 0
temp0 = 5    #Temperature of the thermostat

#Increment
dt = 0.02

#Heat equation
def u(x,t):
    return temp0 + scale*np.exp(-k*t)*np.sin(x)

#Gradient of the Heat Equation
def grad_u(x,t):
                               #df/dx              #df/dt
    return scale*np.array([np.exp(-k*t)*np.cos(x),-k*np.exp(-k*t)*np.sin(x)])

a = []

for i in range(500):
    value = u(x0,t0) + grad_u(x0,t0)[1]*dt
    t0 = t0 + dt
    a.append(value)

k = 0
def animate(i):
    global k
    x = a[k]
    k += 1
    ax1.clear()
    plt.plot(x0,x,color='red',label='Temperature while cooling after heating')
    plt.plot(x0,temp0*np.ones(x0.shape),color='green',label='Thermostat temperature of the rod')
    plt.grid(True)
    plt.legend()
    plt.ylim([-4,2.5*scale])
    plt.xlim([-0.2,L+0.2])
    
anim = animation.FuncAnimation(fig,animate,frames=360,interval=20)
plt.show()