二维扩散模拟中网格更新的变量交换逻辑疑问

用户提供的二维扩散模拟代码如下：

import time

grid_shape = (640, 640)

@profile
def evolve(grid, dt, out, D=1.0):
    xmax, ymax = grid_shape  # width and height of the matrix

    for i in range(xmax):
        for j in range(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 * dt * (grid_xx + grid_yy)

def run_experiment(num_iterations):
    # setting up the initial conditions
    xmax, ymax = grid_shape
    next_grid = [[0.0] * ymax for x in range(xmax)]
    grid = [[0.0] * ymax for x in range(xmax)]

    # initial conditions for simulating a drop of dye in the middle of the simulated region
    block_low = int(grid_shape[0] * 0.4)
    block_high = int(grid_shape[0] * 0.5)

    for i in range(block_low, block_high):
        for j in range(block_low, block_high):
            grid[i][j] = 0.005

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

    return time.time() - start


run_experiment(100)

用户疑问：我理解evolve函数会原地修改next_grid，随后将grid赋值为next_grid以反映当前迭代的更新，但疑惑为何要将next_grid赋值为网格的前一状态grid？

解答 #

这么做核心是两点：保证计算正确性和节省内存开销：

• 扩散计算必须基于上一轮的完整状态。如果直接在grid上原地修改，计算过程中会覆盖掉还没用到的单元格原始值，导致后续相邻单元格的计算用的是已经更新过的数据，结果必然出错。所以必须用两个独立网格，一个存当前状态（输入），一个存计算后的新状态（输出）。
• 每次迭代交换两个变量的引用，是为了复用已有的网格空间。对于640x640的大数组，每次迭代都新建网格会带来频繁的内存分配和释放，拖慢程序速度。交换引用后，旧的grid就变成了下一轮的输出容器，不用重新创建数组，直接覆盖里面的值即可。
• 简单说就是两个网格轮流当"输入盘"和"输出盘"，既确保了计算逻辑的正确性，又最大化利用了已分配的内存资源。

内容的提问来源于stack exchange，提问作者disguisedtoast