如何加速DataFrame逐时间步的np.trapz积分循环计算？

问题描述 #

我有一个形状为[16,60000]的DataFrame，列对应不同时间步。需要在每个时间步沿两个不同轴使用np.trapz计算积分，尝试了以下三种方式但均未获得明显性能提升：

• 经典循环+列表追加
• 经典循环+数组存储
• 列表推导式

最小可复现代码 #

import numpy as np
import pandas as pd    
import time

time_start = time.time()
# Read data
df_data = pd.DataFrame(np.random.randn(16, 60000))
x_values = np.array([0.  , 0.03, 0.1 , 0.2 , 0.3 , 0.4 , 0.5 , 0.6 , 0.7 , 0.8 , 0.9 ,
                     0.85, 0.7 , 0.5 , 0.2 , 0.05])
x_values2 = np.array([0.   ,  0.043,  0.083,  0.114,  0.13 ,  0.134,  0.124,  0.102,
                      0.078,  0.056,  0.03 , -0.006, -0.02 , -0.055, -0.069, -0.042])

# Get sample characteristics
Ns = df_data.shape[1]
times = range(Ns)

lt_data = [df_data.iloc[:,i] for i in times]
a = np.array([-np.trapz(y=data, x=x_values) for data in lt_data])
b = np.array([np.trapz(y=data, x=x_values2) for data in lt_data])

time_end = time.time()
elapsed = time_end - time_start
print(f'Elapsed: {elapsed:.1f}s')

优化方案：向量化运算+减少不必要的转换 #

核心是彻底抛弃Python循环，利用numpy的向量化特性批量计算所有积分，同时避免pandas索引操作的额外开销。

优化后的代码

import numpy as np
import pandas as pd    
import time

time_start = time.time()
# 若原始数据是DataFrame，直接转成numpy数组
df_data = pd.DataFrame(np.random.randn(16, 60000))
data_arr = df_data.values

x_values = np.array([0.  , 0.03, 0.1 , 0.2 , 0.3 , 0.4 , 0.5 , 0.6 , 0.7 , 0.8 , 0.9 ,
                     0.85, 0.7 , 0.5 , 0.2 , 0.05])
x_values2 = np.array([0.   ,  0.043,  0.083,  0.114,  0.13 ,  0.134,  0.124,  0.102,
                      0.078,  0.056,  0.03 , -0.006, -0.02 , -0.055, -0.069, -0.042])

# 直接对整个数组计算积分，axis=0表示沿第0轴（每一列的16个元素）积分
a = -np.trapz(y=data_arr, x=x_values, axis=0)
b = np.trapz(y=data_arr, x=x_values2, axis=0)

time_end = time.time()
elapsed = time_end - time_start
print(f'Elapsed: {elapsed:.4f}s')

优化原理

1. 向量化运算：np.trapz的axis参数可指定积分轴，直接对(16,60000)数组批量计算，所有操作在numpy底层C实现中完成，彻底避免Python循环的性能损耗。
2. 减少pandas开销：直接操作numpy数组，省去df_data.iloc[:,i]这类索引操作的额外时间，仅需一次从DataFrame到numpy数组的转换。

性能对比

原代码运行时间约2-3秒（取决于机器配置），优化后的代码运行时间通常在0.01秒以内，性能提升可达数百倍。

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