从含ID与时间步的表格数据提取X_train、y_train并转换为指定形状

核心逻辑 #

按ID分组提取特征与目标值，通过分组操作和数组形状调整，得到符合要求的X_train（三维数组）和y_train（二维数组）。

代码实现 #

首先导入依赖库：

import pandas as pd
import numpy as np

1. 确保数据按ID升序排列（保证X_train[0]对应ID=1）

df_sorted = df.sort_values('ID').reset_index(drop=True)

2. 生成特征数组X_train
   按ID分组提取特征列x1、x2，转换为numpy数组后调整为(9, 24, 2)的形状：

X_train = df_sorted.groupby('ID')[['x1', 'x2']].apply(lambda group: group.values).to_numpy()
# 强制调整形状（确保分组结果符合24条/ID的预期）
X_train = X_train.reshape(9, 24, 2)

3. 生成目标数组y_train
   按ID提取唯一的y值（同一ID下y值一致），转换为(9, 1)的形状：

y_train = df_sorted.groupby('ID')['y'].first().to_numpy().reshape(9, 1)

结果验证 #

可以通过以下代码验证X_train[0,:,:]是否对应ID=1的特征数据：

id1_features = df[df['ID'] == 1][['x1', 'x2']].values
print(np.array_equal(X_train[0], id1_features))  # 输出True即为正确

内容的提问来源于stack exchange，提问作者Muhammad Ikhwan Perwira