NumPy计算积分时Energy与cos函数数组形状不匹配如何解决

编写代码计算0到pi区间积分时触发数组形状不匹配报错，原代码如下：

import numpy as np
from math import pi,cos
vtheta=np.linspace(0.0,pi,1000)
def my_function(x):
    Energy = np.arange(2.1,300.1,0.1)
    return ((1.0)/(Energy-1+np.cos(x)))
print (my_function(vtheta).sum())

运行后提示Energy数组与cos(x)返回数组形状不一致，无法得到计算结果。

• 直接触发报错的原因是numpy广播规则不匹配：代码中np.arange(2.1,300.1,0.1)生成的Energy是长度为2980的一维数组，作为采样点传入的vtheta是长度为1000的一维数组，两个长度不等的一维数组无法执行逐元素运算，因此抛出形状错误。
• 原代码存在两处隐性逻辑错误：
  • 导入的math.cos为标量计算函数，无法直接处理numpy数组输入，应替换为numpy内置的np.cos
  • 数值积分需要在采样点求和后乘以采样步长，原代码直接求和无法得到正确的积分结果

根据计算需求分两种场景调整：

场景1：需要遍历所有Energy值，分别计算每个能量对应的积分结果 #

通过给Energy增加长度为1的维度，让两个数组满足numpy广播要求，运算后沿角度采样轴求和，再乘以积分步长即可：

import numpy as np
from math import pi

vtheta = np.linspace(0.0, pi, 1000)
dx = vtheta[1] - vtheta[0]  # 角度采样步长

def my_function(x):
    Energy = np.arange(2.1, 300.1, 0.1)
    # Energy调整为(2980, 1)的列向量，可与长度1000的角度数组广播为二维数组
    return 1.0 / (Energy[:, None] - 1 + np.cos(x))

# 沿角度轴（第1维度）求和后乘步长，得到每个Energy对应的积分值
integral_result = my_function(vtheta).sum(axis=1) * dx
print(integral_result)

场景2：仅需计算固定Energy参数下的单积分结果 #

将Energy定义为固定标量值即可，不需要做维度调整：

import numpy as np
from math import pi

vtheta = np.linspace(0.0, pi, 1000)
dx = vtheta[1] - vtheta[0]
fixed_E = 2.1  # 替换为目标能量值

def my_function(x, E):
    return 1.0 / (E - 1 + np.cos(x))

integral_result = my_function(vtheta, fixed_E).sum() * dx
print(integral_result)

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