1
2
3 """
4 This script compares different ways of implementing an iterative
5 procedure to solve Laplace's equation. These provide a general
6 guideline to using Python for high-performance computing and also
7 provide a simple means to compare the computational time taken by the
8 different approaches. The script compares functions implemented in
9 pure Python, Numeric, weave.blitz, weave.inline, fortran (via f2py)
10 and Pyrex. The function main(), additionally accelerates the pure
11 Python version using Psyco and provides some numbers on how well that
12 works. To compare all the options you need to have Numeric, weave,
13 f2py, Pyrex and Psyco installed. If Psyco is not installed the script
14 will print a warning but will perform all other tests.
15
16 The fortran and pyrex modules are compiled using the setup.py script
17 that is provided with this file. You can build them like so:
18
19 python setup.py build_ext --inplace
20
21
22 Author: Prabhu Ramachandran <prabhu_r at users dot sf dot net>
23 License: BSD
24 Last modified: Sep. 18, 2004
25 """
26
27 import numpy
28 from scipy import weave
29 from scipy.weave import converters
30 import sys, time
31
32 msg = """**************************************************
33 Please build the fortran and Pyrex modules like so:
34
35 python setup.py build_ext --inplace
36
37 You will require f2py and Pyrex.
38 **************************************************
39 """
40 build = 0
41 try:
42 import flaplace
43 except ImportError:
44 build = 1
45 try:
46 import pyx_lap
47 except ImportError:
48 build = 1
49 if build:
50 print msg
51
52
53 class Grid:
54
55 """A simple grid class that stores the details and solution of the
56 computational grid."""
57
58 def __init__(self, nx=10, ny=10, xmin=0.0, xmax=1.0,
59 ymin=0.0, ymax=1.0):
60 self.xmin, self.xmax, self.ymin, self.ymax = xmin, xmax, ymin, ymax
61 self.dx = float(xmax-xmin)/(nx-1)
62 self.dy = float(ymax-ymin)/(ny-1)
63 self.u = numpy.zeros((nx, ny), 'd')
64
65 self.old_u = self.u.copy()
66
67 def setBC(self, l, r, b, t):
68 """Sets the boundary condition given the left, right, bottom
69 and top values (or arrays)"""
70 self.u[0, :] = l
71 self.u[-1, :] = r
72 self.u[:, 0] = b
73 self.u[:,-1] = t
74 self.old_u = self.u.copy()
75
76 def setBCFunc(self, func):
77 """Sets the BC given a function of two variables."""
78 xmin, ymin = self.xmin, self.ymin
79 xmax, ymax = self.xmax, self.ymax
80 x = numpy.arange(xmin, xmax + self.dx*0.5, self.dx)
81 y = numpy.arange(ymin, ymax + self.dy*0.5, self.dy)
82 self.u[0 ,:] = func(xmin,y)
83 self.u[-1,:] = func(xmax,y)
84 self.u[:, 0] = func(x,ymin)
85 self.u[:,-1] = func(x,ymax)
86
87 def computeError(self):
88 """Computes absolute error using an L2 norm for the solution.
89 This requires that self.u and self.old_u must be appropriately
90 setup."""
91 v = (self.u - self.old_u).flat
92 return numpy.sqrt(numpy.dot(v,v))
93
94
95 class LaplaceSolver:
96
97 """A simple Laplacian solver that can use different schemes to
98 solve the problem."""
99
100 def __init__(self, grid, stepper='numeric'):
101 self.grid = grid
102 self.setTimeStepper(stepper)
103
104 def slowTimeStep(self, dt=0.0):
105 """Takes a time step using straight forward Python loops."""
106 g = self.grid
107 nx, ny = g.u.shape
108 dx2, dy2 = g.dx**2, g.dy**2
109 dnr_inv = 0.5/(dx2 + dy2)
110 u = g.u
111
112 err = 0.0
113 for i in range(1, nx-1):
114 for j in range(1, ny-1):
115 tmp = u[i,j]
116 u[i,j] = ((u[i-1, j] + u[i+1, j])*dy2 +
117 (u[i, j-1] + u[i, j+1])*dx2)*dnr_inv
118 diff = u[i,j] - tmp
119 err += diff*diff
120
121 return numpy.sqrt(err)
122
123 def numericTimeStep(self, dt=0.0):
124 """Takes a time step using a numeric expressions."""
125 g = self.grid
126 dx2, dy2 = g.dx**2, g.dy**2
127 dnr_inv = 0.5/(dx2 + dy2)
128 u = g.u
129 g.old_u = u.copy()
130
131
132 u[1:-1, 1:-1] = ((u[0:-2, 1:-1] + u[2:, 1:-1])*dy2 +
133 (u[1:-1,0:-2] + u[1:-1, 2:])*dx2)*dnr_inv
134
135 return g.computeError()
136
137 def blitzTimeStep(self, dt=0.0):
138 """Takes a time step using a numeric expression that has been
139 blitzed using weave."""
140 g = self.grid
141 dx2, dy2 = g.dx**2, g.dy**2
142 dnr_inv = 0.5/(dx2 + dy2)
143 u = g.u
144 g.old_u = u.copy()
145
146
147 expr = "u[1:-1, 1:-1] = ((u[0:-2, 1:-1] + u[2:, 1:-1])*dy2 + "\
148 "(u[1:-1,0:-2] + u[1:-1, 2:])*dx2)*dnr_inv"
149 weave.blitz(expr, check_size=0)
150
151 return g.computeError()
152
153 def inlineTimeStep(self, dt=0.0):
154 """Takes a time step using inlined C code -- this version uses
155 blitz arrays."""
156 g = self.grid
157 nx, ny = g.u.shape
158 dx2, dy2 = g.dx**2, g.dy**2
159 dnr_inv = 0.5/(dx2 + dy2)
160 u = g.u
161
162 code = """
163 #line 120 "laplace.py"
164 double tmp, err, diff;
165 err = 0.0;
166 for (int i=1; i<nx-1; ++i) {
167 for (int j=1; j<ny-1; ++j) {
168 tmp = u(i,j);
169 u(i,j) = ((u(i-1,j) + u(i+1,j))*dy2 +
170 (u(i,j-1) + u(i,j+1))*dx2)*dnr_inv;
171 diff = u(i,j) - tmp;
172 err += diff*diff;
173 }
174 }
175 return_val = sqrt(err);
176 """
177
178 err = weave.inline(code,
179 ['u', 'dx2', 'dy2', 'dnr_inv', 'nx','ny'],
180 type_converters = converters.blitz,
181 compiler = 'gcc')
182 return err
183
184 def fastInlineTimeStep(self, dt=0.0):
185 """Takes a time step using inlined C code -- this version is
186 faster, dirtier and manipulates the numeric array in C. This
187 code was contributed by Eric Jones. """
188 g = self.grid
189 nx, ny = g.u.shape
190 dx2, dy2 = g.dx**2, g.dy**2
191 dnr_inv = 0.5/(dx2 + dy2)
192 u = g.u
193
194 code = """
195 #line 151 "laplace.py"
196 double tmp, err, diff;
197 double *uc, *uu, *ud, *ul, *ur;
198 err = 0.0;
199 for (int i=1; i<nx-1; ++i) {
200 uc = u+i*ny+1;
201 ur = u+i*ny+2; ul = u+i*ny;
202 ud = u+(i+1)*ny+1; uu = u+(i-1)*ny+1;
203 for (int j=1; j<ny-1; ++j) {
204 tmp = *uc;
205 *uc = ((*ul + *ur)*dy2 +
206 (*uu + *ud)*dx2)*dnr_inv;
207 diff = *uc - tmp;
208 err += diff*diff;
209 uc++;ur++;ul++;ud++;uu++;
210 }
211 }
212 return_val = sqrt(err);
213 """
214
215 err = weave.inline(code,
216 ['u', 'dx2', 'dy2', 'dnr_inv', 'nx','ny'],
217 compiler='gcc')
218 return err
219
220 def fortranTimeStep(self, dt=0.0):
221 """Takes a time step using a simple fortran module that
222 implements the loop in fortran. Use f2py to compile
223 flaplace.f like so: f2py -c flaplace.c -m flaplace. You need
224 the latest f2py version for this to work. This Fortran
225 example was contributed by Pearu Peterson. """
226 g = self.grid
227 g.u, err = flaplace.timestep(g.u, g.dx, g.dy)
228 return err
229
230 def pyrexTimeStep(self, dt=0.0):
231 """Takes a time step using a function written in Pyrex. Use
232 the given setup.py to build the extension using the command
233 python setup.py build_ext --inplace. You will need Pyrex
234 installed to run this."""
235 g = self.grid
236 err = pyx_lap.pyrexTimeStep(g.u, g.dx, g.dy)
237 return err
238
239 def setTimeStepper(self, stepper='numeric'):
240 """Sets the time step scheme to be used while solving given a
241 string which should be one of ['slow', 'numeric', 'blitz',
242 'inline', 'fastinline', 'fortran']."""
243 if stepper == 'slow':
244 self.timeStep = self.slowTimeStep
245 elif stepper == 'numeric':
246 self.timeStep = self.numericTimeStep
247 elif stepper == 'blitz':
248 self.timeStep = self.blitzTimeStep
249 elif stepper == 'inline':
250 self.timeStep = self.inlineTimeStep
251 elif stepper.lower() == 'fastinline':
252 self.timeStep = self.fastInlineTimeStep
253 elif stepper == 'fortran':
254 self.timeStep = self.fortranTimeStep
255 elif stepper == 'pyrex':
256 self.timeStep = self.pyrexTimeStep
257 else:
258 self.timeStep = self.numericTimeStep
259
260 def solve(self, n_iter=0, eps=1.0e-16):
261 """Solves the equation given an error precision -- eps. If
262 n_iter=0 the solving is stopped only on the eps condition. If
263 n_iter is finite then solution stops in that many iterations
264 or when the error is less than eps whichever is earlier.
265 Returns the error if the loop breaks on the n_iter condition
266 and returns the iterations if the loop breaks on the error
267 condition."""
268 err = self.timeStep()
269 count = 1
270
271 while err > eps:
272 if n_iter and count >= n_iter:
273 return err
274 err = self.timeStep()
275 count = count + 1
276
277 return count
278
279
280 def BC(x, y):
281 """Used to set the boundary condition for the grid of points.
282 Change this as you feel fit."""
283 return (x**2 - y**2)
284
285
286 def test(nmin=5, nmax=30, dn=5, eps=1.0e-16, n_iter=0, stepper='numeric'):
287 iters = []
288 n_grd = numpy.arange(nmin, nmax, dn)
289 times = []
290 for i in n_grd:
291 g = Grid(nx=i, ny=i)
292 g.setBCFunc(BC)
293 s = LaplaceSolver(g, stepper)
294 t1 = time.clock()
295 iters.append(s.solve(n_iter=n_iter, eps=eps))
296 dt = time.clock() - t1
297 times.append(dt)
298 print "Solution for nx = ny = %d, took %f seconds"%(i, dt)
299 return (n_grd**2, iters, times)
300
301
302 def time_test(nx=500, ny=500, eps=1.0e-16, n_iter=100, stepper='numeric'):
303 g = Grid(nx, ny)
304 g.setBCFunc(BC)
305 s = LaplaceSolver(g, stepper)
306 t = time.clock()
307 s.solve(n_iter=n_iter, eps=eps)
308 return time.clock() - t
309
310
311 def main(n=500, n_iter=100):
312 print "Doing %d iterations on a %dx%d grid"%(n_iter, n, n)
313 for i in ['numeric', 'blitz', 'inline', 'fastinline', 'fortran',
314 'pyrex']:
315 print i,
316 sys.stdout.flush()
317 print "took", time_test(n, n, stepper=i, n_iter=n_iter), "seconds"
318
319 print "slow (1 iteration)",
320 sys.stdout.flush()
321 s = time_test(n, n, stepper='slow', n_iter=1)
322 print "took", s, "seconds"
323 print "%d iterations should take about %f seconds"%(n_iter, s*n_iter)
324
325 try:
326 import psyco
327 except ImportError:
328 print "You don't have Psyco installed!"
329 else:
330 psyco.bind(LaplaceSolver)
331 psyco.bind(Grid)
332 print "slow with Psyco (1 iteration)",
333 sys.stdout.flush()
334 s = time_test(n, n, stepper='slow', n_iter=1)
335 print "took", s, "seconds"
336 print "%d iterations should take about %f seconds"%\
337 (n_iter, s*n_iter)
338
339
340 if __name__ == "__main__":
341 main()