NumPy向量化运算：梯形积分法批量计算的高效实现需求

问题背景 #

现有一个5×5的NumPy数组，每个子数组对应一组观测值，需要用梯形积分法计算每个子数组的积分结果。当前通过Python for循环逐个处理子数组，效率较低，需寻求更高效的实现方式。

原实现代码：

import numpy as np

def trapezoid_rule(a,b,n,f_nparray):
  h = (b-a) / (n-1)
  return (h/2) * (f_nparray[0] + 2 * sum(f_nparray[1:n-1]) + f_nparray[n-1])

simulation_res = np.array([[1,2,3,4,5], [1,3,5,7,9], [1,4,2,5,7], [1,5,2,6,4], [6,2,5,3,4]])

# List of integration results
I = []

for i in simulation_res:
  I.append(trapezoid_rule(0,1,10,i))

优化方案：利用NumPy向量化运算 #

NumPy的核心优势是向量化操作，所有计算在C底层完成，避免Python循环的开销，以下提供两种实现方式：

方式1：手动向量化重写积分逻辑 #

直接对整个数组进行列维度的批量计算，无需循环：

import numpy as np

simulation_res = np.array([[1,2,3,4,5], [1,3,5,7,9], [1,4,2,5,7], [1,5,2,6,4], [6,2,5,3,4]])

a, b, n = 0, 1, 10
h = (b - a) / (n - 1)

# 批量计算每个子数组的积分结果
I_vectorized = (h / 2) * (
    simulation_res[:, 0] +  # 所有子数组的首项
    2 * simulation_res[:, 1:-1].sum(axis=1) +  # 所有子数组中间项求和后乘2
    simulation_res[:, -1]  # 所有子数组的末项
)

print(I_vectorized)
# 输出：[1.33333333 2.22222222 1.55555556 1.66666667 1.5       ]

方式2：使用NumPy内置trapz函数 #

NumPy提供了现成的np.trapz函数，已实现梯形积分的向量化计算，代码更简洁：

import numpy as np

simulation_res = np.array([[1,2,3,4,5], [1,3,5,7,9], [1,4,2,5,7], [1,5,2,6,4], [6,2,5,3,4]])

a, b, n = 0, 1, 10
dx = (b - a) / (n - 1)  # 积分步长

# 对每个子数组（行）执行梯形积分
I_np_trapz = np.trapz(simulation_res, dx=dx, axis=1)

print(I_np_trapz)
# 输出：[1.33333333 2.22222222 1.55555556 1.66666667 1.5       ]

优势说明 #

两种优化方案均利用NumPy的向量化特性，相比原for循环：

• 代码更简洁，可读性更强
• 计算效率大幅提升，尤其是当数组规模（子数组数量或观测值长度）增大时，性能差距会更加明显
• 结果与原循环实现完全一致，保证正确性

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