1D-heat-equation.py

import numpy as np
import matplotlib.pyplot as plt
import time

L = 1.0 # length of 1-D heat-conducting object
Nx = 100 # number of spatial grid points
T = 10.0 # maximum time
Nt = 1000 # number of time steps
a = 0.005 # material proberty alpha

x = np.linspace(0, L, Nx+1)    # mesh points in space
dx = x[1] - x[0]

t = np.linspace(0, T, Nt+1)    # mesh points in time
dt = t[1] - t[0]

u   = np.zeros(Nx+1)           # unknown u at new time step
u_1 = np.zeros(Nx+1)           # u at the previous time step

# plotting boilerplate
fig = plt.figure()
ax = fig.add_subplot(111)
ax.set_ylim([0,1])
li, = ax.plot(x, u)
ax.relim() 
ax.autoscale_view(True,True,True)
fig.canvas.draw()
plt.show(block=False)

# definition of initial conditions
def initial(x):
	return x**2

# Set initial condition u(x,0) = initial(x)
for i in range(0, Nx+1):
    u_1[i] = initial(x[i])

# loop through every time step
for n in range(0, Nt):

    # Compute u at inner grid points
    for i in range(1, Nx):
        u[i] = u_1[i] + a*dt/(dx*dx)*(u_1[i-1] - 2*u_1[i] + u_1[i+1])

    # Appöy boundary conditions
    u[0] = 1.
    u[Nx] = 0.

    # Update u_1 before next step
    u_1[:]= u

    # plot every 10 time steps
    if n % 10 == 0:
	    li.set_ydata(u)
	    fig.canvas.draw()
	    time.sleep(0.001)
	    plt.savefig('frames/' + str(n).zfill(3) + '.png')