如何按指标分组提取时间序列预测结果的指定周期数据

需求说明 #

针对多指标时间序列预测结果数据集，需按以下规则提取数据：

• 对每个indicator_name对应的子数据集，提取actual_value非空（拟合周期）的最后3条数据（同步保留对应fitted_value）
• 提取predicted_value非空（预测周期）的前2条数据

样本数据 #

indicator_name        date  actual_value  fitted_value  predicted_value
0             CPI  2023-07-31          -7.7     -9.280851              NaN
1             CPI  2023-08-31          72.5      5.181248              NaN
2             CPI  2023-09-30         126.3     75.801026              NaN
3             CPI  2023-10-31         147.9    123.940712              NaN
4             CPI  2023-11-30           NaN           NaN       113.341466
5             CPI  2023-12-31           NaN           NaN       103.223588
6             CPI  2024-01-31           NaN           NaN       117.865400
7             PPI  2022-09-30          25.8     12.700000              NaN
8             PPI  2022-12-31           6.0     12.800000              NaN
9             PPI  2023-03-31           6.1      8.800000              NaN
10            PPI  2023-06-30           7.9      9.700000              NaN
11            PPI  2023-09-30           8.1     11.000000              NaN
12            PPI  2023-12-31           NaN           NaN        11.500000
13            PPI  2024-03-31           NaN           NaN        13.700000
14            PPI  2024-06-30           NaN           NaN        12.800000
15            PPI  2024-09-30           NaN           NaN        15.300000

预期结果 #

indicator_name        date  actual_value  fitted_value  predicted_value
0            CPI  2023-08-31          72.5      5.181248              NaN
1            CPI  2023-09-30         126.3     75.801026              NaN
2            CPI  2023-10-31         147.9    123.940712              NaN
3            CPI  2023-11-30           NaN           NaN       113.341466
4            CPI  2023-12-31           NaN           NaN       103.223588
5            PPI  2023-03-31           6.1      8.800000              NaN
6            PPI  2023-06-30           7.9      9.700000              NaN
7            PPI  2023-09-30           8.1     11.000000              NaN
8            PPI  2023-12-31           NaN           NaN        11.500000
9            PPI  2024-03-31           NaN           NaN        13.700000

实现代码（中文注释） #

import pandas as pd

def slice_actual_data(df, group_cols, n: int=3):
    """
    提取每个分组中actual_value非空的最后n条数据
    :param df: 原始数据集
    :param group_cols: 分组列名列表，此处为['indicator_name']
    :param n: 需要提取的最后n条数据，默认3
    :return: 切片后的实际值数据集
    """
    # 筛选actual_value非空的行，按日期排序
    actual_df = df[df['actual_value'].notna()].sort_values('date')
    # 按分组取最后n条数据
    return actual_df.groupby(group_cols).tail(n)

def slice_predicted_data(df, group_cols, n: int=2):
    """
    提取每个分组中predicted_value非空的前n条数据
    :param df: 原始数据集
    :param group_cols: 分组列名列表，此处为['indicator_name']
    :param n: 需要提取的前n条数据，默认2
    :return: 切片后的预测值数据集
    """
    # 筛选predicted_value非空的行，按日期排序
    predicted_df = df[df['predicted_value'].notna()].sort_values('date')
    # 按分组取前n条数据
    return predicted_df.groupby(group_cols).head(n)

def get_final_result(df):
    # 分别提取实际值和预测值切片
    actual_part = slice_actual_data(df, ['indicator_name'])
    predicted_part = slice_predicted_data(df, ['indicator_name'])
    # 合并结果并按指标和日期排序，重置索引
    return pd.concat([actual_part, predicted_part]).sort_values(['indicator_name', 'date']).reset_index(drop=True)

代码说明 #

原代码未区分实际值行和预测值行，直接对全表分组取头尾会导致数据混选。改进后的逻辑：

1. 先拆分两类数据：筛选actual_value非空的拟合周期行、predicted_value非空的预测周期行
2. 分别按indicator_name分组，提取对应条数的目标数据
3. 合并两类结果并排序，得到符合预期的最终数据集

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