如何按组对Pandas DataFrame每n天计算斜率（slope）？

给定如下Pandas DataFrame：

import pandas as pd

data = [['A', '2022-09-01', 2], ['A', '2022-09-02', 1], ['A', '2022-09-04', 3], ['A', '2022-09-06', 2],
        ['A', '2022-09-07', 1], ['A', '2022-09-07', 2], ['A', '2022-09-08', 4], ['A', '2022-09-09', 2],
        ['B', '2022-09-01', 2], ['B', '2022-09-03', 4], ['B', '2022-09-04', 2], ['B', '2022-09-05', 2],
        ['B', '2022-09-07', 1], ['B', '2022-09-08', 3], ['B', '2022-09-10', 2]]
df = pd.DataFrame(data = data, columns = ['group', 'date', 'value'])

df['date'] = pd.to_datetime(df['date'])
df['diff_days'] = (df['date']-df['date'].groupby(df['group']).transform('first')).dt.days

生成的DataFrame如下：

group       date  value  diff_days
0      A 2022-09-01      2          0
1      A 2022-09-02      1          1
2      A 2022-09-04      3          3
3      A 2022-09-06      2          5
4      A 2022-09-07      1          6
5      A 2022-09-07      2          6
6      A 2022-09-08      4          7
7      A 2022-09-09      2          8
8      B 2022-09-01      2          0
9      B 2022-09-03      4          2
10     B 2022-09-04      2          3
11     B 2022-09-05      2          4
12     B 2022-09-07      1          6
13     B 2022-09-08      3          7
14     B 2022-09-10      2          9

需要新增一列slope，按group分组后，每n天（本案例n=3）划分区间，区间从每组首个日期开始计算。利用diff_days和value列计算每个区间的斜率，区间内所有行共享同一斜率。

期望输出如下：

data = [['A', '2022-09-01', 2, 0, 0.43], ['A', '2022-09-02', 1, 1, 0.43], ['A', '2022-09-04', 3, 3, 0.43], ['A', '2022-09-06', 2, 5, -0.5],
        ['A', '2022-09-07', 1, 6, -0.5], ['A', '2022-09-07', 2, 6, -0.5], ['A', '2022-09-08', 4, 7, -2], ['A', '2022-09-09', 2, 8, -2],
        ['B', '2022-09-01', 2, 0, 0.14], ['B', '2022-09-03', 4, 2, 0.14], ['B', '2022-09-04', 2, 3, 0.14], ['B', '2022-09-05', 2, 4, -0.5],
        ['B', '2022-09-07', 1, 6, -0.5], ['B', '2022-09-08', 3, 7, -0.5], ['B', '2022-09-10', 2, 9, -0.5]]
df_desired = pd.DataFrame(data = data, columns = ['group', 'date', 'value', 'diff_days', 'slope'])

对应的DataFrame：

group        date  value  diff_days  slope
0      A  2022-09-01      2          0   0.43
1      A  2022-09-02      1          1   0.43
2      A  2022-09-04      3          3   0.43
3      A  2022-09-06      2          5  -0.50
4      A  2022-09-07      1          6  -0.50
5      A  2022-09-07      2          6  -0.50
6      A  2022-09-08      4          7  -2.00
7      A  2022-09-09      2          8  -2.00
8      B  2022-09-01      2          0   0.14
9      B  2022-09-03      4          2   0.14
10     B  2022-09-04      2          3   0.14
11     B  2022-09-05      2          4  -0.50
12     B  2022-09-07      1          6  -0.50
13     B  2022-09-08      3          7  -0.50
14     B  2022-09-10      2          9  -0.50

斜率计算示例：

• A组第一个3天区间：slope([0,1,3],[2,1,3])=0.43
• A组下一个3天区间：slope([5,6,6],[2,1,2])=-0.5
• A组再下一个区间：slope([7,8],[4,2])=-2.0

注意：数据中并非包含所有日期，需严格按每n天划分区间。

步骤1：按组划分n天区间 #

根据diff_days和n值，为每个组生成等宽区间，区间从0开始，每n天为一个分段（如0-3天、3-6天等），即使数据中没有连续日期也能正确划分：

n = 3

def create_bins(group, n):
    max_diff = group['diff_days'].max()
    # 生成区间边界：0,3,6,...直到覆盖最大diff_days
    bins = list(range(0, max_diff + n + 1, n))
    return pd.cut(group['diff_days'], bins=bins, labels=False, include_lowest=True)

df['interval'] = df.groupby('group').apply(create_bins, n=n).reset_index(drop=True)

步骤2：计算每个区间的斜率 #

使用线性回归计算每个区间内diff_days和value的斜率，保留两位小数：

import numpy as np

def calculate_slope(x, y):
    if len(x) <= 1:
        return np.nan  # 避免单一点无法计算斜率
    slope, _ = np.polyfit(x, y, 1)
    return round(slope, 2)

# 按组和区间分组计算斜率
slope_df = df.groupby(['group', 'interval']).apply(
    lambda x: calculate_slope(x['diff_days'], x['value'])
).reset_index(name='slope')

# 将斜率合并回原DataFrame
df = df.merge(slope_df, on=['group', 'interval'])

完整代码 #

import pandas as pd
import numpy as np

# 生成原始数据
data = [['A', '2022-09-01', 2], ['A', '2022-09-02', 1], ['A', '2022-09-04', 3], ['A', '2022-09-06', 2],
        ['A', '2022-09-07', 1], ['A', '2022-09-07', 2], ['A', '2022-09-08', 4], ['A', '2022-09-09', 2],
        ['B', '2022-09-01', 2], ['B', '2022-09-03', 4], ['B', '2022-09-04', 2], ['B', '2022-09-05', 2],
        ['B', '2022-09-07', 1], ['B', '2022-09-08', 3], ['B', '2022-09-10', 2]]
df = pd.DataFrame(data = data, columns = ['group', 'date', 'value'])

df['date'] = pd.to_datetime(df['date'])
df['diff_days'] = (df['date']-df['date'].groupby(df['group']).transform('first')).dt.days

# 划分区间
n = 3
def create_bins(group, n):
    max_diff = group['diff_days'].max()
    bins = list(range(0, max_diff + n +1, n))
    return pd.cut(group['diff_days'], bins=bins, labels=False, include_lowest=True)

df['interval'] = df.groupby('group').apply(create_bins, n=n).reset_index(drop=True)

# 计算斜率
def calculate_slope(x, y):
    if len(x) <=1:
        return np.nan
    slope, _ = np.polyfit(x, y, 1)
    return round(slope, 2)

slope_df = df.groupby(['group', 'interval']).apply(lambda x: calculate_slope(x['diff_days'], x['value'])).reset_index(name='slope')
df = df.merge(slope_df, on=['group', 'interval'])

# 移除临时的interval列（可选）
df = df.drop('interval', axis=1)
print(df)

运行后输出结果与期望完全一致，区间内所有行共享对应斜率。

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