Performance Python: Solving The 2D Diffusion Equation With numpy

The sliding scale between performance and coding simplicity has led to a spectrum of languages over the years. In the world of scientific computing, the two major languages are of course C and FORTRAN. Both of these are compiled, low level and high performance languages. Generally, one rarely considers an executed language for scientific grunt-work since interpreted languages tend to be drastically slower. In recent years, computers have become fast enough for a compromise, but still, most interpreted languages fall short of the task.

Python is different. On its own, it’s still too slow for general scientific application. However, its ability to easily interface with bits of binary code make it attractively extensible for scientific and numerical applications.

Numpy, a package of scientific computing tools for Python, extends Python’s data structures with a multidimensional array object, various derived objects and an assortment of routines for fast operations on arrays. As long as your problem can be cast into numpy’s array objects, operations can be carried out with speeds to rival a pure C implementation.

SciPy and matplotlib further extend numpy + Python to give functionality similar to Matlab making Python an excellent free and open-source alternative. Python has the added advantage of being a popular and complete programming language.

The aim of this article is to demonstrate how numpy can be used to dramatically increase the efficiency of the simple task of integrating the two dimensional diffusion equation. There will also be an example of advanced visualization with matplotlib, creating an animated image map to visualize the solution. The code in this article is nowhere near as elegant as it could be, but it is intentionally plain for the beginners eye.

The two-dimensional diffusion equation

We consider the partial differential equation

or more compactly, on the domain of the unit square (x,y between zero and one). This is the diffusion equation which models diffusive processes such as heat and chemical concentration.

We approximate by the discrete function , where , and . Applying finite (forward) difference approximations to the derivatives, then we obtain the partial difference equation

or

(1)  

A word on stability

The stability of the scheme requires that the following condition is met:

If we set the time interval at the highest value allowed by the condition, we will still see some noise in the solution. Reducing the time-step size can improve this.

Initial conditions, boundary values and characteristics of the solution

Equation (1) is a blueprint for integrating the solution forward in time, calculating the solution at each time-step based on the solution at the previous time-step. Since a previous time-step is required for each subsequent step, the initial conditions (i.e., ) must be given. Additionally, the derivatives in the Laplacian require boundary conditions.

For the example in this article, we will choose the initial conditions to be a ring centred at (0.5, 0.5), defined by

We will also choose the boundary conditions to be zero, in agreement with the initial conditions.

The solution given these conditions is illustrated in Fig. 1

Fig. 1: An illustration of the time evolution of diffusion (A) when t=0 (showing the initial conditions), (B) when m=5, (C) when m=200 and (D) when m=500.

Coding the problem in Python

Plain Python

I’d hope the reader is somewhat familiar with Python; but if not, it should be easy enough to follow along. Feel free to post questions, although I’d imagine that Google is good enough to answer most.

To integrate Eq. (1) through time, we set up a loop,

for m in range(1, timesteps+1):
	print "Computing u for m =", m
	evolve_ts(u, ui)

where m is our time index and timesteps holds the number of time-steps to calculate (say, 500). We have to define some function evolve_ts(u, ui), which will calculate u from ui.

Using plain Python syntax, the function evolve_ts should look something like this:

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)

where nx and ny contain the number of elements of u in the x- and y-direction (i.e., nx=1/dx where dx=). Although Python indexes from 0, we want the loops to start at index 1 and end at index nx-1 or ny-1 (as the case may be). This is because we wish to leave the boundary values at zero, as defined by our boundary conditions.

Here is a listing of a fully functional Python script which implements the above code. Copy and paste it into a document and execute it (you may have to chmod u+x to execute it under Linux/UNIX)

#!/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.

"""
import scipy as sp
import time

# 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.

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)

# Now, start the time evolution calculation...
tstart = time.time()
for m in range(1, timesteps+1):
	evolve_ts(u, ui)
	print "Computing u for m =", m
tfinish = time.time()

print "Done."
print "Total time: ", tfinish-tstart, "s"
print "Average time per time-step using numpy: ", ( tfinish - tstart )/timesteps, "s."

Executing the above code on my office workstation for 500 time-steps with dx=dy=0.01 takes about 88 seconds, averaging 0.1758s per time-step. Considering a compiled C-code implementation would execute all 500 time-steps in less than two-tenths of a second, the Python performance is pretty dismal.

Enhanced by numpy

Luckily, we can make a small change to the above code to boost the perfomance close to that of pure C.

Assuming you’ve run the above Python code, replace the definition of evolve_ts with the following:

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 )

In this version of evolve_ts, the loops for x and y have been replaced by a numpy expression. The benefits here are two-fold: (1) the syntax is leaner, replacing five lines of code with only one (albeit a larger one) and (2) numpy executes the expression using C. Hence, the performance of the code is drastically increased. On my computer, the numpy expression takes less than a quarter-second to execute (compare plain Python at 88s) with an average of 4.70094E-4s per time-step: nearly 400-times faster than the plain Python code.

The trick with the numpy expressions is the way numpy iterates through the slices. To understand exactly what the above statement does, you can use the official (tentative) numpy tutorial, starting with the section on slicing.

Using matplotlib to visualize

Fig 2. An example of the output from the matplotlib visualization, with white edges trimmed off for web presentation.

Matplotlib can be very simple and easy to use. In this particular case, I wanted to do something quite advanced: create an animated, real-time density map (Fig. 2) of the solution as it is calculated.

The following code uses the numpy enhanced update_ts function, and pops up a windows which renders the solution in real time, as it’s calculated. The performance is low, due to the rendering. Some details of the code are briefly explained in the comments; but I’ll generally leave it to the user to figure out what things do. I didn’t intend this example to be instructive but merely a showpiece for what can be done with relatively little code.

#!/usr/bin/env python
"""
Author: Timothy A.V. Teatro <http://www.timteatro.net>
Date  : Oct 25, 2010
Lisence: Creative Commons BY-SA
(http://creativecommons.org/licenses/by-sa/2.0/)

Description:
	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 u(x,y,t). This version uses a numpy
	expression which is evaluated in C, so the
	computation time is greatly reduced over plain
	Python code.

	This version also uses matplotlib to create an
	animation of the time evolution of the density.
"""
import scipy as sp
import matplotlib
matplotlib.use('GTKAgg') # Change this as desired.
import gobject
from pylab import *
# 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.

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):
	"""
	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 )

def updatefig(*args):
	global u, ui, m
	im.set_array(ui)
	manager.canvas.draw()
	# Uncomment the next two lines to save images as png
	# filename='diffusion_ts'+str(m)+'.png'
	# fig.savefig(filename)
	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 )
	ui = sp.copy(u)
	m+=1
	print "Computing and rendering u for m =", m
	if m >= timesteps:
		return False
	return True

fig = plt.figure(1)
img = subplot(111)
im = img.imshow( ui, cmap=cm.hot, interpolation='nearest', origin='lower')
manager = get_current_fig_manager()

m=1
fig.colorbar( im ) # Show the colorbar along the side

# once idle, call updatefig until it returns false.
gobject.idle_add(updatefig)
show()

Be Sociable, Share!

• 

Related posts:

1. Why Every Physicist Should Know Python If you deal with data sets that need to be...
2. Compiling the latest matplotlib from source As of the date of this post, the latest Ubuntu...