如何用Numpy替代循环优化二维网格随机系统模拟速度？

问题背景 #

我正在尝试在二维网格上模拟一个随机系统：每当随机时钟触发时，一个像素的颜色会发生变化，可能还会改变某个随机相邻像素的颜色。现有代码能实现预期效果，但在1000x1000的大网格下速度不佳（100x100网格约耗时7秒，500x500网格约45秒），推测step()函数中的循环是性能瓶颈。

核心问题 #

1. 如何利用Numpy函数替代该循环，尤其是处理相邻数组元素的访问与修改？
2. 还有哪些简单方法可以提升模拟速度？
3. 曾考虑用np.where()规避循环和内部条件判断，但不知如何访问并修改随机相邻元素；尝试过Numba但未成功，请问值得继续尝试吗？

解决方案 #

一、用Numpy矢量化操作替代循环 #

1. 批量随机选择目标像素

不用逐个循环选像素，直接一次性生成所有需要更新的像素坐标：

# 假设每次step要更新N个像素，N可按时钟触发频率设置
N = int(grid.size * 0.01)  # 比如每次更新1%的像素
rows = np.random.randint(0, grid.shape[0], size=N)
cols = np.random.randint(0, grid.shape[1], size=N)

2. 处理随机相邻像素的矢量化方法

针对每个选中的像素，随机选择一个相邻方向（上下左右），然后计算相邻像素的坐标，同时处理边界情况：

# 定义四个方向的偏移量：上、下、左、右
dirs = np.array([[-1,0], [1,0], [0,-1], [0,1]])
# 随机为每个选中像素选一个方向
dir_idx = np.random.randint(0, 4, size=N)
selected_dirs = dirs[dir_idx]

# 计算相邻像素坐标
neigh_rows = rows + selected_dirs[:,0]
neigh_cols = cols + selected_dirs[:,1]

# 边界处理：把超出网格的相邻像素坐标修正为自身（或者忽略，根据需求调整）
neigh_rows = np.clip(neigh_rows, 0, grid.shape[0]-1)
neigh_cols = np.clip(neigh_cols, 0, grid.shape[1]-1)

3. 批量修改像素值

直接通过索引批量更新目标像素和相邻像素，完全替代循环内的逐个修改：

# 修改目标像素颜色（示例：随机切换0/1，可替换为你的颜色逻辑）
grid[rows, cols] = 1 - grid[rows, cols]
# 随机决定是否修改相邻像素（比如50%概率）
modify_neigh = np.random.choice([True, False], size=N)
grid[neigh_rows[modify_neigh], neigh_cols[modify_neigh]] = 1 - grid[neigh_rows[modify_neigh], neigh_cols[modify_neigh]]

这种方式把所有循环操作转化为Numpy的矢量化计算，避免了Python级别的循环开销，速度会有数量级的提升。

二、其他简单提速方法 #

• 减少不必要的内存操作：避免在step里频繁创建新数组，尽量复用已有的数组（比如提前定义dirs、预分配rows/cols数组）。
• 降低精度或简化颜色逻辑：如果颜色用整数类型（比如uint8）而非float，内存占用更少，计算更快。
• 并行化单步操作：用numpy.random的批量生成替代逐个随机数生成，或者把大网格拆分成小块并行处理（适合超大规模网格）。
• 减少可视化开销：如果模拟过程中实时绘图，把绘图频率降低（比如每10步画一次），或者用更快的绘图模式（比如matplotlib的blit模式）。

三、关于Numba的可行性 #

绝对值得继续尝试！Numba对循环密集型的Python代码优化效果非常显著，尤其是这种涉及数组索引和随机操作的场景。

如果之前尝试失败，大概率是这几个原因：

1. 没有给函数加上@numba.jit或@numba.njit装饰器（njit是纯编译模式，速度更快）。
2. 函数内使用了Numba不支持的Python特性（比如某些复杂的类方法、动态类型变量）。
3. 没有处理边界条件的分支逻辑（Numba对分支的优化需要明确的类型注解）。

给你一个简单的Numba适配示例：

import numba as nb

@nb.njit
def step_numba(grid, update_prob, neigh_modify_prob):
    rows, cols = grid.shape
    # 随机选择要更新的像素数量
    N = int(rows * cols * update_prob)
    for _ in range(N):
        r = np.random.randint(0, rows)
        c = np.random.randint(0, cols)
        # 修改当前像素
        grid[r, c] = 1 - grid[r, c]
        # 随机选择相邻方向
        dir = np.random.randint(0, 4)
        nr, nc = r, c
        if dir == 0 and r > 0:
            nr = r - 1
        elif dir == 1 and r < rows-1:
            nr = r + 1
        elif dir == 2 and c > 0:
            nc = c - 1
        elif dir == 3 and c < cols-1:
            nc = c + 1
        # 随机修改相邻像素
        if np.random.random() < neigh_modify_prob:
            grid[nr, nc] = 1 - grid[nr, nc]
    return grid

注意：第一次调用Numba函数会有编译延迟，但后续调用速度会比纯Python循环快几十倍甚至上百倍。

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