如何优化DataFrame列移位计算斜率生成新列的低效代码？

我需要优化一段Python代码，实现以下功能：

• 将DataFrame的指定列移位x次，让每行包含x个历史值
• 结合当前行的值组成数组，计算线性回归斜率
• 将斜率存入对应的新列

原代码可运行但效率极低，处理1万行数据耗时约30秒，原代码如下：

from datetime import datetime

import math
import pandas as pd
import numpy as np

def trending(row, label, amt, prepend = False, postpend = False, reverse = True):
    arr = []

    if prepend:
        arr.append(row[label])

    for i in range(amt):
        arr.append(row[f'{label}_{i}'])

    if postpend:
        arr.append(row[label])

    if reverse:
        arr.reverse()

    if any(x == None or math.isnan(x) for x in arr):
        return 0

    return slope(arr)

def slope(arr):
    if len(arr) <= 1:
        return 0
    coeffs = np.polyfit(range(len(arr)), arr, 1)
    slope = coeffs[-2]
    return float(slope)

def trends(df):
    start = datetime.now()

    look_back = 5

    for i in range(look_back):
        j = i + 1
        df2 = df[['close', 'pvi', 'nvi', 'smi', 'roc', 'macd', 'histogram', 'percent_b', 'height']]
        df2 = df2.add_suffix(f'_{i}')
        df2 = df2.shift(j)
        df = pd.concat([df, df2], axis=1)

    def close_trend(row):
        return trending(row, 'close', look_back, True, False)

    def pvi_trend(row):
        return trending(row, 'pvi', look_back, True, False)

    def nvi_trend(row):
        return trending(row, 'nvi', look_back, True, False)

    def smi_trend(row):
        return trending(row, 'smi', look_back, True, False)

    def macd_trend(row):
        return trending(row, 'macd', look_back, True, False)

    def roc_trend(row):
        return trending(row, 'roc', look_back, True, False)

    def histogram_trend(row):
        return trending(row, 'histogram', look_back, True, False)

    def percent_b_trend(row):
        return trending(row, 'percent_b', look_back, True, False)

    def height_trend(row):
        return trending(row, 'height', look_back, True, False)
    
    df['close_trend'] = df.apply(close_trend, axis=1)
    df['pvi_trend'] = df.apply(pvi_trend, axis=1)
    df['nvi_trend'] = df.apply(nvi_trend, axis=1)
    df['smi_trend'] = df.apply(smi_trend, axis=1)
    df['macd_trend'] = df.apply(macd_trend, axis=1)
    df['roc_trend'] = df.apply(roc_trend, axis=1)
    df['histogram_trend'] = df.apply(histogram_trend, axis=1)
    df['percent_b_trend'] = df.apply(percent_b_trend, axis=1)
    df['height_trend'] = df.apply(height_trend, axis=1)

    print(f'Trends function took {datetime.now() - start}')

    return df

data = []

for i in range(10000):
    data.append([545.9, 0.3333398862,   0.01673619117,  0.2111060119,   55.95725508,    1.100447539,    0.8652411735,   0.8219623901,   1.808441041,    46.79554862])

columns = ['close', 'macd', 'histogram', 'roc', 'rsi', 'pvi', 'nvi', 'percent_b', 'height', 'smi']

df = pd.DataFrame(columns=columns, data=data)

df = trends(df)

print(df)

优化后的代码 #

from datetime import datetime
import pandas as pd
import numpy as np

def trends_optimized(df):
    start = datetime.now()
    look_back = 5
    target_cols = ['close', 'pvi', 'nvi', 'smi', 'roc', 'macd', 'histogram', 'percent_b', 'height']
    
    # 预计算x轴值，避免重复创建
    x = np.arange(look_back + 1)
    
    # 向量化斜率计算函数
    def calc_slope(window):
        if np.isnan(window).any():
            return 0.0
        coeffs = np.polyfit(x, window, 1)
        return coeffs[0]
    
    for col in target_cols:
        # 滚动窗口包含当前行+前look_back行
        rolling_window = df[col].rolling(window=look_back+1, min_periods=look_back+1)
        # raw=True直接传递numpy数组，提升效率
        df[f'{col}_trend'] = rolling_window.apply(calc_slope, raw=True)
        # 填充窗口不足的行（前look_back行）为0
        df[f'{col}_trend'] = df[f'{col}_trend'].fillna(0.0)
    
    print(f'Trends function took {datetime.now() - start}')
    return df

# 测试数据
data = []
for i in range(10000):
    data.append([545.9, 0.3333398862, 0.01673619117, 0.2111060119, 55.95725508, 
                 1.100447539, 0.8652411735, 0.8219623901, 1.808441041, 46.79554862])
columns = ['close', 'macd', 'histogram', 'roc', 'rsi', 'pvi', 'nvi', 'percent_b', 'height', 'smi']
df = pd.DataFrame(columns=columns, data=data)

df = trends_optimized(df)
print(df)

核心优化点 #

• 替换逐行操作：用rolling.apply(raw=True)替代df.apply(axis=1)，利用Pandas向量化特性大幅提升效率
• 简化移位逻辑：直接使用滚动窗口rolling，避免手动创建移位列和多次concat，减少内存开销
• 消除重复代码：统一遍历所有目标列，移除重复的单列处理函数
• 预计算公共数据：提前生成x轴数组，避免每次计算斜率时重复创建
• 高效空值处理：用np.isnan批量判断空值，替代逐元素检查

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