1 #!/usr/bin/env python
 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  # used to compute the change in solution in some of the methods.
 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  # The actual iteration
 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  # The actual iteration
 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  # compiler keyword only needed on windows with MSVC installed
 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  # compiler keyword only needed on windows with MSVC installed
 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()