Question

让我们将例6-2的伪代码写成正式代码，这样我们就能够更好地分析它的运行时性能。第一步是写出演化函数，它接受一个矩阵并返回其演化后的状态，见例6-3。

例6-3　纯Python 2阶扩散函数

grid_shape = (1024, 1024)

def evolve(grid, dt, D=1.0):
  xmax, ymax = grid_shape
  new_grid = [[0.0,] * ymax for x in xrange(xmax)]
  for i in xrange(xmax):
    for j in xrange(ymax):
      grid_xx = grid[(i+1)%xmax][j] + grid[(i-1)%xmax][j] - 2.0 * grid[i][j]
      grid_yy = grid[i][(j+1)%ymax] + grid[i][(j-1)%ymax] - 2.0 * grid[i][j]
      new_grid[i][j] = grid[i][j] + D * (grid_xx + grid_yy) * dt
return new_grid

　备忘　

其实相比预分配一个new_grid列表，还不如在for循环中用appends建立它。不仅速度比我们这样写要快很多，结果也是一样的。我们之所以使用现在这个方法是因为它能解释得更清楚。

全局变量grid_shape表示我们模拟的区域有多大。另外，在6.1节中解释过，我们使用的是周期性边界条件（这也是为什么我们需要对索引取模）。为了实际使用这段代码，我们必须初始化一个矩阵并对其调用evolve。例6-4中的代码是一个非常通用的初始化过程，将在本章多次被重用（它的性能将不会被分析，因为它只运行一次，不像evolve函数将被重复调用）。

例6-4　纯Python 2阶扩散初始化函数

def run_experiment(num_iterations):
  # Setting up initial conditions  ❶
  xmax, ymax = grid_shape
  grid = [[0.0,] * ymax for x in xrange(xmax)]

  block_low = int(grid_shape[0] * .4)
  block_high = int(grid_shape[0] * .5)
  for i in xrange(block_low, block_high):
    for j in xrange(block_low, block_high):
      grid[i][j] = 0.005

  # Evolve the initial conditions
  start = time.time()
  for i in range(num_iterations):
    grid = evolve(grid, 0.1)
  return time.time() - start

❶ 这里使用的初始条件跟图6-2的方形例子相同。

我们特地选择了足够小的dt和矩阵元素的值来让算法稳定。对这个算法的收敛特性更深入的讨论见Willian Press等人的著作《Numerical Recipes》第三版。

分配次数太多带来的问题

通过对纯Python演化函数使用line_profiler，我们就能开始理解运行变慢可能是什么原因导致的。见例6-5的分析输出，我们看到函数大多数时间花在了求导和更新矩阵上。^[1]这正是我们所需要的，因为这是一个纯粹的CPU密集型问题——任何不是花在CPU上的时间都是一个明显的优化对象。

例6-5　纯Python 2阶扩散函数分析

$ kernprof.py -lv diffusion_python.py
Wrote profile results to diffusion_python.py.lprof
Timer unit: 1e-06 s

File: diffusion_python.py
Function: evolve at line 8
Total time: 16.1398 s

Line #   Hits     Time  Per Hit  % Time  Line Contents
==============================================================
   8                     @profile
   9                     def evolve(grid, dt, D=1.0):
  10     10     39    3.9   0.0    xmax, ymax = grid_shape # ❶
  11  2626570  2159628    0.8  13.4    new_grid = ...
  12   5130     4167    0.8   0.0    for i in xrange(xmax): # ❷
  13  2626560  2126592    0.8  13.2      for j in xrange(ymax):
  14  2621440  4259164    1.6  26.4        grid_xx = ...
  15  2621440  4196964    1.6  26.0        grid_yy = ...
  16  2621440  3393273    1.3  21.0        new_grid[i][j] = ...
  17     10     10    1.0   0.0       return grid # ❸

❶ 这条语句用了这么长时间是因为grid_shape必须从本地名字空间获得（见4.2节）。

❷ 这行命中了5130次，意味着我们操作的矩阵的xmax为512。因为需要对xrange中的每一个值求值512次加上对循环终止条件求一次值，而这些一共发生了10次。

❸ 这行命中10次，告诉我们这个被分析函数运行了10次。

但是，输出也显示我们有20%的时间花费在分配new_grid列表上。这是一个浪费，因为new_grid的属性不会发生变化——无论我们传入evolve的值是什么，new_grid列表始终具有相同的形状和大小，包含同样的值。一个简单优化是只分配一次并重用它。这种优化类似于将重复代码移至循环外：

from math import sin

def loop_slow(num_iterations):
  """
  >>> %timeit loop_slow(int(1e4))
  100 loops, best of 3: 2.67 ms per loop
  """
  result = 0
  for i in xrange(num_iterations):
    result += i * sin(num_iterations) # ❶
  return result

def loop_fast(num_iterations):
  """
  >>> %timeit loop_fast(int(1e4))
  1000 loops, best of 3: 1.38 ms per loop
  """
  result = 0
  factor = sin(num_iterations)
  for i in xrange(num_iterations):
    result += i
  return result * factor

❶ sin(num_iterations)的值在循环内不会改变，所以不需要每次重新计算它。

我们可以对扩散代码进行类似的转换，如例6-6所示。在此，我们将例6-4的new_ grid实例化之后再传入我们的evolve函数。该函数做的事情跟之前一样：读取grid列表并写入新的new_grid列表。然后，我们可以简单交换new_grid和grid并继续。

例6-6　减少内存分配之后的纯Python 2阶扩散函数

def evolve(grid, dt, out, D=1.0):
  xmax, ymax = grid_shape
  for i in xrange(xmax):
    for j in xrange(ymax):
      grid_xx = grid[(i+1)%xmax][j] + grid[(i-1)%xmax][j] - 2.0 * grid[i][j]
      grid_yy = grid[i][(j+1)%ymax] + grid[i][(j-1)%ymax] - 2.0 * grid[i][j]
      out[i][j] = grid[i][j] + D * (grid_xx + grid_yy) * dt

def run_experiment(num_iterations):
  # 设置初始条件
  xmax,ymax = grid_shape
  next_grid = [[0.0,] * ymax for x in xrange(xmax)]
  grid = [[0.0,] * ymax for x in xrange(xmax)]

  block_low = int(grid_shape[0] * .4)
  block_high = int(grid_shape[0] * .5)
  for i in xrange(block_low, block_high):
    for j in xrange(block_low, block_high):
      grid[i][j] = 0.005

  start = time.time()
  for i in range(num_iterations):
    evolve(grid, 0.1, next_grid)
    grid, next_grid = next_grid, grid
  return time.time() – start

在例6-7的分析结果中^[2]，我们可以看到经过这个细小的修改后的版本给我们带来了21%的速度提升。这个结论跟之前对列表的append操作给我们的结论（见3.2.1节）类似：内存分配不便宜。每次当我们需要内存用于存储一个变量或列表，Python都必须花时间向操作系统申请更多内存空间，然后还要遍历新分配的空间来将它初始化为某个值。任何时候只要有可能，重用已分配的内存空间都能为我们带来速度提升。不过，要小心实现这些改动。速度提升固然可观，你还是需要通过分析来确保这是你想要的结果，而不是污染你代码库的垃圾。

例6-7　减少内存分配后的Python扩散函数的分析结果

$ kernprof.py -lv diffusion_python_memory.py
Wrote profile results to diffusion_python_memory.py.lprof
Timer unit: 1e-06 s

File: diffusion_python_memory.py
Function: evolve at line 8
Total time: 13.3209 s

Line #    Hits    Time  Per Hit  % Time  Line Contents
==============================================================
   8                     @profile
   9                     def evolve(grid, dt, out, D=1.0):
  10    10    15    1.5    0.0  xmax, ymax = grid_shape
  11    5130    3853    0.8    0.0  for i in xrange(xmax):
  12   2626560   1942976    0.7   14.6    for j in xrange(ymax):
  13   2621440   4059998    1.5   30.5      grid_xx = ...
  14   2621440   4038560    1.5   30.3      grid_yy = ...
  15   2621440   3275454    1.2   24.6      out[i][j] = ...