如何高效计算大型3D NumPy数组中每个单元的邻域平均值？

你的三重Python循环效率极低的核心原因是Python循环本身的开销，尤其是处理大数组和多次迭代时。以下是两种高效的替代方案：

方案一：纯NumPy向量化实现 #

利用NumPy的切片操作直接计算所有邻域的和，完全避免Python循环，依赖NumPy的底层C优化，速度提升显著。

import numpy as np

def fun_vectorized(mat):
    # 计算六个方向邻域的和
    neighbors_sum = (
        mat[:-2, 1:-1, 1:-1] + mat[2:, 1:-1, 1:-1] +  # r轴正负方向
        mat[1:-1, :-2, 1:-1] + mat[1:-1, 2:, 1:-1] +  # c轴正负方向
        mat[1:-1, 1:-1, :-2] + mat[1:-1, 1:-1, 2:]    # d轴正负方向
    )
    # 将平均值赋值给中间区域（与原代码边界处理逻辑一致）
    mat[1:-1, 1:-1, 1:-1] = neighbors_sum / 6
    return mat

# 使用示例
# R=C=D=200
# mat = np.random.rand(R, C, D)  # 示例数组
for iter in range(1000):
    mat = fun_vectorized(mat)

优势 #

• 无额外依赖，仅需NumPy
• 速度比三重循环快100~1000倍（取决于数组大小）
• 内存效率高，切片操作仅引用原数组数据，无额外复制

方案二：SciPy 3D卷积实现 #

如果需要更灵活的邻域操作（比如后续修改核），可以使用SciPy的ndimage.convolve，通过定义3D卷积核实现平均值计算。

import numpy as np
from scipy.ndimage import convolve

# 定义3D卷积核：仅六个邻域位置为1，中心为0，预先除以6得到平均值核
kernel = np.zeros((3, 3, 3), dtype=np.float64)
kernel[0, 1, 1] = 1  # r-1方向
kernel[2, 1, 1] = 1  # r+1方向
kernel[1, 0, 1] = 1  # c-1方向
kernel[1, 2, 1] = 1  # c+1方向
kernel[1, 1, 0] = 1  # d-1方向
kernel[1, 1, 2] = 1  # d+1方向
kernel /= 6

def fun_convolve(mat):
    # 卷积计算，mode='constant'保证边界处理与原代码一致（边界值不更新）
    conv_result = convolve(mat, kernel, mode='constant')
    mat[1:-1, 1:-1, 1:-1] = conv_result[1:-1, 1:-1, 1:-1]
    return mat

# 使用示例
for iter in range(1000):
    mat = fun_convolve(mat)

优势 #

• 灵活性高，修改核即可实现不同的邻域计算逻辑
• 内置边界处理模式可选（如reflect、wrap等），满足不同需求

额外优化建议 #

1. 数据类型优化：将数组设为float32而非默认的float64，可减少一半内存占用，提升缓存命中率，进一步加快计算速度（500×500×500的float32数组仅占500MB）。
2. 避免原地修改（可选）：如果需要保留原始数据，可在函数内创建副本进行计算，但会增加内存开销。

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