如何让NumPy einsum忽略数组中的NaN值？

问题核心在于np.einsum本身不具备跳过NaN的能力，只要运算路径中存在NaN，最终结果就会变成NaN。要实现忽略NaN的求和计算，需要结合np.nansum替代einsum的求和逻辑，具体修改如下：

修改后的代码 #

# 计算修剪后的斜率：用nansum跳过NaN元素，完成维度求和
slope_trimmed = (np.nansum(y_trimmed_norm * x_trimmed_norm, axis=1) /
                np.nansum(x_trimmed_norm * x_trimmed_norm, axis=1))

# 修正原代码笔误：beta_trimmed改为slope_trimmed
intercept_trimmed = y_trimmed_mean[:, 0] - slope_trimmed * x_trimmed_mean[:, 0]

原理说明 #

• 原代码中np.einsum('ij,ij->i', a, b)等价于对每个i维度，计算j维度上a*b的元素和，即np.sum(a*b, axis=1)。
• 当数组包含NaN时，np.sum会直接返回NaN，而np.nansum会自动跳过所有NaN元素，仅对有效数值求和，避免结果全为NaN的问题。

特殊情况处理 #

如果某个i对应的所有j元素都是NaN，np.nansum会返回0，此时分母为0会引发除零错误。可根据业务需求添加判断：

# 避免除零错误，将分母为0的情况设为NaN（或其他默认值）
denominator = np.nansum(x_trimmed_norm * x_trimmed_norm, axis=1)
slope_trimmed = np.where(denominator != 0, 
                         np.nansum(y_trimmed_norm * x_trimmed_norm, axis=1) / denominator,
                         np.nan)

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