NumPy二维高斯乘积求和优化 解决大数组内存占满问题

问题核心 #

当前计算需求为对每个偏移量$(x0_t, y0_t)$，计算固定二维分布$g1$与中心偏移的二维分布$g2_t$在1001×1001网格上的逐点乘积和，即$\sum_{i,j} g1_{ij} \cdot g2_{t,ij}$。现有全量广播实现会生成尺寸为1001×1001×N的中间数组（N为时间步总数，最高可达1e5），内存占用超过硬件上限；逐时间步循环的向量化程度低，运行速度无法满足需求；缩小网格尺寸会损失计算精度，且不适配非高斯类的通用分布场景。

优化方案 #

1. 分块批量计算（通用首选，无精度损失） #

这是适配性最强的方案，不受分布类型限制，可同时兼顾内存占用和运行速度：

• 核心逻辑：不要一次性对所有N个时间步做广播，而是根据机器可用内存，将时间步序列切分为若干大小合适的块，块内沿用原有广播+numexpr多核加速的逻辑计算，单块计算完成后立即释放中间数组，再计算下一块。
• 内存控制：单块中间数组的内存占用为 1001 * 1001 * block_size * 8字节（float64类型），比如设置block_size=1000时，单块内存约8GB；block_size=250时单块内存约2GB，可根据硬件配置灵活调整。
• 性能表现：块内为全向量化运算，无额外计算开销，运行速度和全量广播方案几乎一致，比逐点循环快2~3个数量级。

参考实现代码：

import numpy as np
import numexpr as ne

def gaussian2d(x,y,x0,y0,x_std,y_std):
   return np.exp(-(x-x0)**2/x_std**2-(y-y0)**2/y_std**2)

# 初始化固定网格与g1分布
x = np.linspace(-5,5,1001)
y = np.linspace(-5,5,1001)
xx, yy = np.meshgrid(x,y)
g1 = gaussian2d(xx, yy, 0, 0, 0.25, 0.25)
X_std = 0.75
Y_std = 0.75
dx = x[1] - x[0]
dy = y[1] - y[0]

# 时间步偏移量
n_steps = 25000
x0 = np.random.rand(n_steps)
y0 = np.random.rand(n_steps)

# 分块计算
block_size = 1000  # 按可用内存调整
final = np.empty(n_steps, dtype=np.float64)
for block_idx in range(0, n_steps, block_size):
    block_slice = slice(block_idx, block_idx + block_size)
    x0_b = x0[block_slice]
    y0_b = y0[block_slice]
    
    X = xx[:,:,np.newaxis] - x0_b
    Y = yy[:,:,np.newaxis] - y0_b
    g2_block = ne.evaluate('exp(-(X)**2/(2*X_std**2)-(Y)**2/(2*Y_std**2))')
    # 直接沿网格维度求和，避免数组转置开销
    final[block_slice] = np.sum(g2_block * g1[:,:,np.newaxis], axis=(0,1)) * dx * dy

2. 高斯场景解析解计算（性能最优，零大数组开销） #

如果两个分布均为二维高斯分布，且当前网格已经覆盖高斯99.9%以上的取值范围（当前设置的[-5,5]区间对应最大标准差0.75时，覆盖超过6σ范围，完全满足要求），可以直接使用两个高斯内积的解析公式计算，完全不需要生成网格相关的大数组：

• 核心原理：二维高斯可分离为x、y两个独立一维高斯的乘积，两个高斯分布的内积（网格求和近似连续积分）存在闭式解，不需要逐网格计算。
• 性能表现：内存开销仅为O(N)，N取1e6时内存占用不到10MB，计算速度比分块方案快10倍以上，结果误差小于网格离散本身的误差。

对应公式：
x方向一维高斯内积：$I_x = \frac{\sqrt{2\pi}\sigma_{1x}\sigma_{2x}}{\sqrt{\sigma_{1x}^{2+\sigma_{2x}}2}} \cdot \exp\left( -\frac{(x0 - 0)^{2}{2(\sigma_{1x}}2+\sigma_{2x}^2)} \right)$
y方向同理得到$I_y$，最终结果为$I_x * I_y$。

参考实现代码：

sigma1x, sigma1y = 0.25, 0.25
sigma2x, sigma2y = 0.75, 0.75

# 预计算常数项
const_x = np.sqrt(2*np.pi) * sigma1x * sigma2x / np.sqrt(sigma1x**2 + sigma2x**2)
const_y = np.sqrt(2*np.pi) * sigma1y * sigma2y / np.sqrt(sigma1y**2 + sigma2y**2)
var_sum_x = sigma1x**2 + sigma2x**2
var_sum_y = sigma1y**2 + sigma2y**2

final = const_x * const_y * np.exp( -x0**2/(2*var_sum_x) - y0**2/(2*var_sum_y) )

注意：该方法仅适用于高斯分布场景，如果后续替换为其他类型分布，需换回分块计算方案。

3. FFT互相关加速（超大规模N场景可选） #

如果后续时间步N增长到1e5以上，且g2为平移不变核（即g2取值仅和网格点与中心的相对坐标有关），该计算本质是g1与g2核的二维互相关运算，可以通过FFT快速计算全平面的互相关结果，再通过插值取所有(x0,y0)对应位置的值即可。该方案内存开销仅为两个1001×1001的FFT数组，和N大小完全无关，N越大速度优势越明显，但需要处理边界补零、插值精度问题，适合超大规模N的通用核场景。

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