将嵌套for循环转换为Numpy矩阵运算求助

先直接给出优化后的代码，再拆解关键逻辑：

import numpy as np

# 前提：xrfa、yrfa、xrft、yrft为一维Numpy数组，at、sigat为标量
# 生成所有i-h组合的dx矩阵（形状(Ma, Mt)）
dx = xrfa[:, np.newaxis] - xrft[np.newaxis, :]
# 生成所有j-k组合的dy矩阵（形状(Na, Nt)）
dy = yrfa[:, np.newaxis] - yrft[np.newaxis, :]

# 扩展维度，让dx²和dy²能广播匹配四维结构
dx_sq = dx[:, np.newaxis, :, np.newaxis] ** 2
dy_sq = dy[np.newaxis, :, np.newaxis, :] ** 2

# 计算最终四维数组Wat
Wat = at * np.exp(-(dx_sq + dy_sq) / (2 * sigat ** 2))

关键逻辑拆解： #

• 生成全量差矩阵：利用Numpy广播机制，把一维数组转成二维差矩阵。比如xrfa[:, np.newaxis]将长度为Ma的一维数组转为(Ma,1)的二维数组，xrft[np.newaxis, :]转为(1,Mt)，两者相减会自动广播成(Ma,Mt)的矩阵，每个元素对应原循环中xrfa[i]-xrft[h]的结果。
• 维度扩展适配四维结构：原Wat是(Ma,Na,Mt,Nt)的四维数组，通过给dx_sq添加第2、4维度（[:, np.newaxis, :, np.newaxis]），给dy_sq添加第1、3维度（[np.newaxis, :, np.newaxis, :]），让两者能广播成一致的四维形状，刚好对应原循环中每个(i,j,h,k)的组合。
• 向量化运算提速：所有运算由Numpy底层C实现，完全替代Python嵌套循环，速度会比原代码快几个数量级。

注意事项： #

• 确保输入的xrfa、yrfa等都是Numpy数组，而非Python列表，否则广播机制无法生效。
• 原代码中如果预先初始化了Wat数组，优化后无需提前创建，Numpy会自动生成对应形状的结果数组。

内容的提问来源于stack exchange，提问作者Grigoras Vlad Albastrel6