如何复现时间序列STS变换滚动窗口训练测试选择动画？

核心逻辑 #

要复现skforecast同款动画，需精准模拟固定窗口训练集滚动、每次重新拟合模型、单步/多步预测的完整流程，用matplotlib动画模块逐帧渲染关键环节。

Google Colab适配代码 #

1. 依赖准备 #

!pip install pandas numpy matplotlib scikit-learn
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from sklearn.linear_model import LinearRegression

2. 生成模拟时间序列 #

np.random.seed(42)
time_idx = pd.date_range(start='2020-01-01', periods=100, freq='D')
y = np.cumsum(np.random.randn(100)) + 100
df = pd.DataFrame({'value': y}, index=time_idx)

3. 回测参数配置 #

fixed_train_size = 20  # 固定训练集大小
n_in = 9  # STS变换回溯步长
n_out = 1  # 单步预测，需多步可修改对应数值
total_test_steps = len(df) - fixed_train_size

4. 动画核心实现 #

# 初始化画布
fig, ax = plt.subplots(figsize=(10, 6))
ax.set_xlabel('时间')
ax.set_ylabel('序列值')
ax.set_title('固定训练集滚动原点回测')

# 初始化绘图元素
train_line, = ax.plot([], [], 'b-', label='训练集')
actual_line, = ax.plot([], [], 'r-', label='真实测试值')
pred_line, = ax.plot([], [], 'g--', label='预测值')
ax.legend()

# 固定坐标轴范围，避免动画跳变
ax.set_xlim(df.index[0], df.index[-1])
ax.set_ylim(df['value'].min()-5, df['value'].max()+5)

# 存储预测结果
preds = []

def update(frame):
    # 划分当前训练/测试集
    train_end = fixed_train_size + frame
    train_data = df.iloc[:train_end]
    test_data = df.iloc[train_end:train_end+n_out]
    
    # 实现STS变换（series_to_supervised逻辑）
    def series_to_supervised(data, n_in, n_out):
        cols = []
        # 输入特征：回溯n_in步
        for i in range(n_in, 0, -1):
            cols.append(data.shift(i))
        # 目标值：未来n_out步
        for i in range(0, n_out):
            cols.append(data.shift(-i))
        agg = pd.concat(cols, axis=1)
        agg.dropna(inplace=True)
        return agg.values
    
    # 处理训练数据
    train_sts = series_to_supervised(train_data[['value']], n_in, n_out)
    X_train, y_train = train_sts[:, :n_in], train_sts[:, n_in:]
    
    # 拟合模型
    model = LinearRegression()
    model.fit(X_train, y_train)
    
    # 生成预测输入（训练集最后n_in个值）
    last_window = train_data['value'].values[-n_in:].reshape(1, -1)
    current_pred = model.predict(last_window)
    
    # 更新预测列表
    if frame < total_test_steps:
        preds.extend(current_pred.flatten())
    
    # 更新绘图数据
    train_line.set_data(train_data.index, train_data['value'])
    actual_line.set_data(df.index[:train_end+len(test_data)], df['value'][:train_end+len(test_data)])
    pred_idx = df.index[fixed_train_size : fixed_train_size+len(preds)]
    pred_line.set_data(pred_idx, preds)
    
    return train_line, actual_line, pred_line

# 创建并显示动画
ani = FuncAnimation(fig, update, frames=total_test_steps, interval=500, blit=True)
from IPython.display import HTML
HTML(ani.to_jshtml())

动画效果优化点 #

• 固定坐标轴范围：提前锁定x/y轴极限，避免帧渲染时坐标轴频繁缩放
• 启用Blit模式：blit=True仅更新变化区域，大幅提升流畅度
• 调整帧间隔：interval=500控制每帧停留时长，可根据需求修改为300-1000毫秒
• 多步预测适配：修改n_out参数即可切换多步预测可视化，STS变换会自动生成对应目标序列

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