如何优化含2D列表的Python草地扩散模拟代码？

1. 只跟踪待处理候选格子，避免全网格遍历 #

每次迭代中，只有与草地相邻的空地块才有可能转化为草地，无需遍历整个网格。可以用集合维护与草地相邻的空坐标：

• 初始时，将初始草地块的所有空邻居加入集合
• 每次迭代，仅遍历集合内的坐标，计算转化概率；对转化成功的格子：
  • 标记为草地
  • 将其未被标记的空邻居加入集合（自动去重）
• 该方式能大幅减少迭代处理的格子数量，模拟后期草地占比越高，效果越显著。

示例代码片段：

import numpy as np

class Biome:
    def __init__(self, size, initial_pos, max_steps):
        self.size = size
        self.grille = np.zeros((size, size), dtype=np.uint8)
        x, y = initial_pos
        self.grille[x, y] = 1
        # 初始化候选集合：初始草地的所有空邻居
        self.candidates = set()
        self._add_neighbors_to_candidates(x, y)
        self.max_steps = max_steps
        self.current_step = 0

    def _add_neighbors_to_candidates(self, x, y):
        for dx in (-1, 0, 1):
            for dy in (-1, 0, 1):
                if dx == 0 and dy == 0:
                    continue
                nx, ny = x + dx, y + dy
                if 0 <= nx < self.size and 0 <= ny < self.size and self.grille[nx, ny] == 0:
                    self.candidates.add((nx, ny))

    def step(self):
        if self.current_step >= self.max_steps or not self.candidates:
            return False
        
        to_convert = []
        new_candidates = set()
        for (x, y) in self.candidates:
            # 统计实际邻居数与草地邻居数
            total_neighbors = 0
            grass_neighbors = 0
            for dx in (-1, 0, 1):
                for dy in (-1, 0, 1):
                    if dx == 0 and dy == 0:
                        continue
                    nx, ny = x + dx, y + dy
                    if 0 <= nx < self.size and 0 <= ny < self.size:
                        total_neighbors += 1
                        grass_neighbors += self.grille[nx, ny]
            proba = grass_neighbors / total_neighbors if total_neighbors > 0 else 0
            if np.random.rand() < proba:
                to_convert.append((x, y))
        
        # 批量更新草地并收集新候选
        for (x, y) in to_convert:
            self.grille[x, y] = 1
            self._add_neighbors_to_candidates(x, y)
        
        # 更新候选集合：移除已转化的，加入新候选
        self.candidates -= set(to_convert)
        self.candidates.update(new_candidates)
        self.current_step += 1
        return True

2. 用Numba JIT编译核心计算逻辑 #

如果核心循环是纯Python实现，用Numba的@njit装饰器将其编译为机器码，可大幅提速统计邻居、计算概率的逻辑。

示例：

from numba import njit

@njit
def count_neighbors(grille, x, y, size):
    grass = 0
    total = 0
    for dx in (-1, 0, 1):
        for dy in (-1, 0, 1):
            if dx == 0 and dy == 0:
                continue
            nx = x + dx
            ny = y + dy
            if 0 <= nx < size and 0 <= ny < size:
                total += 1
                grass += grille[nx, ny]
    return grass, total

@njit
def process_candidates(grille, candidates, size):
    to_convert = []
    new_candidates = set()
    for (x, y) in candidates:
        grass, total = count_neighbors(grille, x, y, size)
        proba = grass / total if total > 0 else 0.0
        if np.random.rand() < proba:
            to_convert.append((x, y))
    
    # 标记转化格子并收集新候选
    for (x, y) in to_convert:
        grille[x, y] = 1
        for dx in (-1, 0, 1):
            for dy in (-1, 0, 1):
                if dx == 0 and dy == 0:
                    continue
                nx = x + dx
                ny = y + dy
                if 0 <= nx < size and 0 <= ny < size and grille[nx, ny] == 0:
                    new_candidates.add((nx, ny))
    
    remaining_candidates = set(candidates) - set(to_convert)
    return remaining_candidates.union(new_candidates), to_convert

之后在step方法中调用该JIT函数即可，Numba对numpy数组和基础数据类型支持良好。

3. 用Numpy向量化操作批量计算邻居数 #

若保留全网格遍历，可利用Numpy的滑动窗口工具批量计算每个格子的邻居草地数，替代逐元素循环：

• 给网格加一圈padding处理边界
• 生成3x3滑动窗口，批量统计每个窗口内的草地数量
• 用预先生成的掩码矩阵处理边界格子的实际邻居数

示例：

def compute_neighbor_counts(grille):
    size = grille.shape[0]
    # 加padding处理边界
    padded = np.pad(grille, pad_width=1, mode='constant', constant_values=0)
    # 生成3x3滑动窗口
    windows = np.lib.stride_tricks.sliding_window_view(padded, window_shape=(3,3))
    # 统计邻居草地数（减去自身）
    grass_counts = windows.sum(axis=(2,3)) - grille
    # 预先生成每个格子的总邻居数掩码
    neighbor_mask = np.ones((size, size), dtype=np.int8) * 8
    # 修正角落和边缘格子的邻居数
    neighbor_mask[0,0] = neighbor_mask[0,-1] = neighbor_mask[-1,0] = neighbor_mask[-1,-1] = 3
    neighbor_mask[0,1:-1] = neighbor_mask[-1,1:-1] = 5
    neighbor_mask[1:-1,0] = neighbor_mask[1:-1,-1] = 5
    return grass_counts, neighbor_mask

# 在Biome类中添加向量化迭代方法
def step_vectorized(self):
    if self.current_step >= self.max_steps or (self.grille == 1).all():
        return False
    grass_counts, neighbor_mask = compute_neighbor_counts(self.grille)
    # 仅处理空格子
    empty_mask = self.grille == 0
    proba = np.where(empty_mask, grass_counts / neighbor_mask, 0.0)
    # 批量判断需要转化的格子
    convert_mask = empty_mask & (np.random.rand(self.size, self.size) < proba)
    self.grille[convert_mask] = 1
    self.current_step += 1
    return True

4. 并行化批量模拟 #

若需大量独立模拟任务，可用多进程绕过Python GIL限制，并行处理不同模拟：

• 用multiprocessing.Pool或concurrent.futures.ProcessPoolExecutor
• 每个进程负责一个完整模拟，返回统计数据（如填满网格步数、草地占比等）

示例：

from multiprocessing import Pool

def run_single_simulation(params):
    size, initial_pos, max_steps = params
    biome = Biome(size, initial_pos, max_steps)
    while biome.step():
        pass
    return biome.current_step, (biome.grille == 1).mean()

# 准备100个模拟任务
tasks = [(50, (25,25), 1000) for _ in range(100)]
# 用4个进程并行处理
with Pool(processes=4) as pool:
    results = pool.map(run_single_simulation, tasks)

# 提取统计结果
total_steps = [r[0] for r in results]
mean_coverage = [r[1] for r in results]

5. 优化内存与数据类型 #

• 用np.uint8存储网格：每个元素仅占1字节，相比默认int类型内存占用更低，缓存命中率更高
• 避免不必要的数组复制：优先使用原地修改（如grille[mask] = 1），减少内存分配与拷贝开销

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