如何高效处理Numpy二维数组边界内的邻域元素计算？

要高效实现你需要的邻域计算，完全可以用Numpy的向量化操作替代嵌套循环，以下是两种简洁且高性能的方案：

方案一：直接构造补零邻域数组 #

通过切片赋值的方式，为每个方向的邻域生成与psi同形状的数组，超出边界的位置自动填充0：

import numpy as np

# 假设psi是输入的(nx, ny)数组
nx, ny = psi.shape

# 构造上邻域（i-1,j）：第一行补0，其余行取psi的上一行
up = np.zeros_like(psi)
up[1:] = psi[:-1]

# 构造下邻域（i+1,j）：最后一行补0，其余行取psi的下一行
down = np.zeros_like(psi)
down[:-1] = psi[1:]

# 构造左邻域（i,j-1）：第一列补0，其余列取psi的左一列
left = np.zeros_like(psi)
left[:, 1:] = psi[:, :-1]

# 构造右邻域（i,j+1）：最后一列补0，其余列取psi的右一列
right = np.zeros_like(psi)
right[:, :-1] = psi[:, 1:]

# 代入自定义表达式f计算phi
phi = f(psi, up, down, left, right)

这个方案的核心是利用Numpy的切片操作批量赋值，所有操作都是向量化的，没有Python层面的循环，计算效率远高于嵌套循环。边界位置的邻域元素会自动保持为0，完美符合需求。

方案二：先补零再切片 #

先给psi的四周补一圈0，再通过切片直接提取各个邻域，逻辑更直观：

import numpy as np

# 给psi四周补一圈0，得到(nx+2, ny+2)的数组
padded_psi = np.pad(psi, pad_width=1, mode='constant', constant_values=0)

# 提取各个邻域（与原psi同形状）
center = padded_psi[1:-1, 1:-1]  # 对应原psi
up = padded_psi[:-2, 1:-1]       # 上邻域
down = padded_psi[2:, 1:-1]      # 下邻域
left = padded_psi[1:-1, :-2]     # 左邻域
right = padded_psi[1:-1, 2:]     # 右邻域

# 计算phi
phi = f(center, up, down, left, right)

这种方法适合需要更清晰邻域对应关系的场景，补零操作由np.pad高效完成，后续切片同样是向量化操作，性能与方案一相当。

示例验证 #

假设你的表达式f是简单的求和：

def f(center, up, down, left, right):
    return center + up + down + left + right

对于左上角元素phi[0][0]，会自动计算为psi[0][0] + 0 + psi[1][0] + 0 + psi[0][1]，完全符合边界补0的要求。

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