如何为Pandas DataFrame每行训练线性回归并生成斜率列

已通过以下代码创建Pandas DataFrame：

import numpy as np
import pandas as pd
    
ds = {'col1' : [11,22,33,24,15,6,7,68,79,10,161,12,113,147,115]}
df = pd.DataFrame(data=ds)

predFeature = []

for i in range(len(df)):
    predFeature.append(0)
    predFeature[i] = predFeature[i-1]+1

df['predFeature'] = predFeature

                
arrayTarget = []
arrayPred = []
target = np.array(df['col1'])
predFeature = np.array(df['predFeature'])

for i in range(len(df)):
    arrayTarget.append(target[i-4:i])
    arrayPred.append(predFeature[i-4:i])
        
df['arrayTarget'] = arrayTarget
df['arrayPred'] = arrayPred

生成的DataFrame如下：

col1  predFeature          arrayTarget         arrayPred
0     11            1                   []                []
1     22            2                   []                []
2     33            3                   []                []
3     24            4                   []                []
4     15            5     [11, 22, 33, 24]      [1, 2, 3, 4]
5      6            6     [22, 33, 24, 15]      [2, 3, 4, 5]
6      7            7      [33, 24, 15, 6]      [3, 4, 5, 6]
7     68            8       [24, 15, 6, 7]      [4, 5, 6, 7]
8     79            9       [15, 6, 7, 68]      [5, 6, 7, 8]
9     10           10       [6, 7, 68, 79]      [6, 7, 8, 9]
10   161           11      [7, 68, 79, 10]     [7, 8, 9, 10]
11    12           12    [68, 79, 10, 161]    [8, 9, 10, 11]
12   113           13    [79, 10, 161, 12]   [9, 10, 11, 12]
13   147           14   [10, 161, 12, 113]  [10, 11, 12, 13]
14   115           15  [161, 12, 113, 147]  [11, 12, 13, 14]

需求 #

• 新增一列slope，每行值为对应窗口数据的线性回归系数
• 前4行slope为NaN，因为窗口数组为空
• 从第5行开始，每行的slope由arrayPred（自变量）和arrayTarget（因变量）计算得出：
  • 第5行：自变量[1,2,3,4]，因变量[11,22,33,24]，斜率为0.10204081632653061
  • 第6行：自变量[2,3,4,5]，因变量[22,33,24,15]，斜率为-0.09090909090909091

解决方案 #

方法一：使用NumPy手动计算（高效轻量） #

利用np.polyfit直接计算一元线性回归的斜率，无需额外库：

def calculate_slope(x, y):
    # 窗口长度不足时返回NaN
    if len(x) < 2:
        return np.nan
    # 计算一阶多项式拟合的斜率
    slope, _ = np.polyfit(x, y, 1)
    return slope

# 应用到DataFrame每行
df['slope'] = df.apply(lambda row: calculate_slope(row['arrayPred'], row['arrayTarget']), axis=1)

方法二：使用Scikit-learn线性回归模型（适合复杂场景） #

若后续需要扩展到多元回归，可使用LinearRegression：

from sklearn.linear_model import LinearRegression

def calculate_slope_sklearn(x, y):
    if len(x) < 2:
        return np.nan
    # 转换为模型要求的二维数组格式
    x_arr = np.array(x).reshape(-1, 1)
    y_arr = np.array(y).reshape(-1, 1)
    # 拟合模型并返回斜率
    model = LinearRegression().fit(x_arr, y_arr)
    return model.coef_[0][0]

df['slope'] = df.apply(lambda row: calculate_slope_sklearn(row['arrayPred'], row['arrayTarget']), axis=1)

两种方法计算结果一致，NumPy方法在数据量较大时速度更优。

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