I also add animation using vpython but can't find 3d or surface version, so I planned to go to matplotlib surface plot route, :)

*(update: here it is, :) )*

#!/usr/bin/env python """ A program which uses an explicit finite difference scheme to solve the diffusion equation with fixed boundary values and a given initial value for the density. Two steps of the solution are stored: the current solution, u, and the previous step, ui. At each time- step, u is calculated from ui. u is moved to ui at the end of each time-step to move forward in time. http://www.timteatro.net/2010/10/29/performance-python-solving-the-2d-diffusion-equation-with-numpy/ he uses matplotlib I use visual python """ import scipy as sp import time from visual import * from visual.graph import * graph1 = gdisplay(x=0, y=0, width=600, height=400, title='x vs. T', xtitle='x', ytitle='T', foreground=color.black, background=color.white) # Declare some variables: dx=0.01 # Interval size in x-direction. dy=0.01 # Interval size in y-direction. a=0.5 # Diffusion constant. timesteps=500 # Number of time-steps to evolve system. t=0. nx = int(1/dx) ny = int(1/dy) dx2=dx**2 # To save CPU cycles, we'll compute Delta x^2 dy2=dy**2 # and Delta y^2 only once and store them. # For stability, this is the largest interval possible # for the size of the time-step: dt = dx2*dy2/( 2*a*(dx2+dy2) ) # Start u and ui off as zero matrices: ui = sp.zeros([nx,ny]) u = sp.zeros([nx,ny]) # Now, set the initial conditions (ui). for i in range(nx): for j in range(ny): if ( ( (i*dx-0.5)**2+(j*dy-0.5)**2 <= 0.1) & ((i*dx-0.5)**2+(j*dy-0.5)**2>=.05) ): ui[i,j] = 1 ''' def evolve_ts(u, ui): global nx, ny """ This function uses two plain Python loops to evaluate the derivatives in the Laplacian, and calculates u[i,j] based on ui[i,j]. """ for i in range(1,nx-1): for j in range(1,ny-1): uxx = ( ui[i+1,j] - 2*ui[i,j] + ui[i-1, j] )/dx2 uyy = ( ui[i,j+1] - 2*ui[i,j] + ui[i, j-1] )/dy2 u[i,j] = ui[i,j]+dt*a*(uxx+uyy) ''' def evolve_ts(u, ui): """ This function uses a numpy expression to evaluate the derivatives in the Laplacian, and calculates u[i,j] based on ui[i,j]. """ u[1:-1, 1:-1] = ui[1:-1, 1:-1] + a*dt*( (ui[2:, 1:-1] - 2*ui[1:-1, 1:-1] + ui[:-2, 1:-1])/dx2 + (ui[1:-1, 2:] - 2*ui[1:-1, 1:-1] + ui[1:-1, :-2])/dy2 ) # Now, start the time evolution calculation... #tstart = time.time() f1 = gcurve(color=color.blue) while True: rate(60) #for m in range(1, timesteps+1): if t<timesteps: t+=dt evolve_ts(u, ui) ui[:] = u[:] # I add this line to update ui value (not present in original code) #print "Computing u for m =", m f1.gcurve.pos = [] for i in arange(nx): f1.plot(pos=(i,u[nx/2,i])) #tfinish = time.time() #print "Done." #print "Total time: ", tfinish-tstart, "s" #print "Average time per time-step using numpy: ", ( tfinish - tstart )/timesteps, "s."