如何将sum(ai*np.exp(bi*r2) for ai,bi in zip(a,b))改为纯numpy实现

核心思路是通过调整一维数组a、b的维度适配r2的形状，利用numpy广播机制完成批量运算后沿求和维度压缩，最终输出形状与r2完全一致。

完整代码实现 #

import numpy as np

# 自动适配任意维度的r2
a_ext = a.reshape(-1, *([1] * r2.ndim))
b_ext = b.reshape(-1, *([1] * r2.ndim))
result = np.sum(a_ext * np.exp(b_ext * r2), axis=0)

原理说明 #

• 假设a和b是长度为N的一维数组，r2.ndim表示r2的维度数，reshape(-1, *([1] * r2.ndim))会把a、b的形状调整为(N, 1, 1, ...)，后续1的数量等于r2的维度数
• 调整维度后的b_ext和r2做乘法时，numpy广播机制会自动将b_ext的尾轴对齐到r2的形状，得到形状为(N, *r2.shape)的运算结果
• 乘以a_ext后沿第0维（也就是a、b的长度维度）求和，求和后第0维消失，最终得到的result形状和r2完全一致

测试样例 #

你可以用以下代码验证效果：

# 测试参数
a = np.array([1, 2])
b = np.array([0.5, 0.1])
r2 = np.array([[1, 2], [3, 4]]) # 2维数组，shape为(2,2)

a_ext = a.reshape(-1, *([1] * r2.ndim))
b_ext = b.reshape(-1, *([1] * r2.ndim))
res = np.sum(a_ext * np.exp(b_ext * r2), axis=0)

print(res.shape) # 输出 (2, 2)，和r2形状完全匹配

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