Here's one approach using Mayavi - read the comments for some explanation:
import numpy as np
import time
# import mayavi's mlab API for scripting
from mayavi import mlab
###variable declarations
nx = 31
ny = 31
nt = 17
nu=.05
dx = 2.0/(nx-1)
dy = 2.0/(ny-1)
sigma = .25
dt = sigma*dx*dy/nu
x = np.linspace(0,2,nx)
y = np.linspace(0,2,ny)
u = np.ones((ny,nx)) ##create a 1xn vector of 1's
un = np.ones((ny,nx)) ##
###Assign initial conditions
u[.5/dy:1/dy+1,.5/dx:1/dx+1]=2 ##set hat function I.C. : u(.5<=x<=1 && .5<=y<=1 ) is 2
X,Y = np.meshgrid(x,y)
###Run through nt timesteps
u[.5/dy:1/dy+1,.5/dx:1/dx+1]=2
# create a surface from grid-shaped data
surf = mlab.mesh(X,Y,u[:])
t = time.time()
max_framerate = 10
for n in range(nt+1):
un[:] = u[:]
u[1:-1,1:-1]=un[1:-1,1:-1]+nu*dt/dx**2*(un[2:,1:-1]-2*un[1:-1,1:-1]+un[0:-2,1:-1])+nu*dt/dy**2* (un[1:-1,2:]-2*un[1:-1,1:-1]+un[1:-1,0:-2])
u[0,:]=1
u[-1,:]=1
u[:,0]=1
u[:,-1]=1
# the mlab_source attribute of surf represents the data we're plotting.
# it has x, y and z attributes as you'd expect. here we only need to
# update the z attribute
surf.mlab_source.z = u
# there's no need to call any equivalent to matplotlib's draw() or show()
# functions - another draw event gets triggered automatically whenever
# surf's data source gets modified
# put a pause in here to control the maximum framerate
while time.time() - t < (1./max_framerate):
pass
t = time.time()
You can find out more about animating data in Mayavi here.