Nugroho's blog.: 2D Diffusion Equation using Python, Scipy, and VPython

Tuesday, November 24, 2015

2D Diffusion Equation using Python, Scipy, and VPython

I got it from here, but modify it here and there.

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 differencescheme to solve the diffusion equation with fixedboundary values and a given initial value for thedensity.Two steps of the solution are stored: the currentsolution, u, and the previous step, ui. At each time-step, u is calculated from ui. u is moved to ui at theend 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 matplotlibI use visual python"""import scipy as spimport timefrom 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^2dy2=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."`
.