|
|
"""
|
|
|
模块功能：该模块用于季节性自回归移动平均 (SARIMA) 预测模型的实施。
|
|
|
|
|
|
函数：
|
|
|
train_SARIMA_model(endog, exo, exog_pred, steps): 使用 SARIMA 模型对时间序列数据进行预测。
|
|
|
run(input_data, forecast_target, exog_columns, steps): 执行模型训练和数据预测等步骤，并返回预测结果。
|
|
|
"""
|
|
|
import pandas as pd
|
|
|
import pmdarima as pm
|
|
|
import numpy as np
|
|
|
|
|
|
from typing import List, Union
|
|
|
|
|
|
|
|
|
def train_SARIMA_model(endog: Union[np.ndarray, pd.Series],
|
|
|
exog: Union[np.ndarray, pd.DataFrame] = None,
|
|
|
exog_pred: Union[np.ndarray, pd.DataFrame] = None,
|
|
|
steps: int = 20,
|
|
|
information_criterion: str = 'aic') -> np.ndarray:
|
|
|
"""
|
|
|
使用SARIMA模型对时间序列数据进行预测。
|
|
|
|
|
|
Args:
|
|
|
endog (Union[np.ndarray, pd.Series]): 用于训练和预测的数据，可以是numpy数组或者pandas序列。
|
|
|
exog (Union[np.ndarray, pd.DataFrame], optional): 外生变量，用于提高ARIMA模型的预测精度。
|
|
|
exog_pred (Union[np.ndarray, pd.DataFrame], optional): 预测阶段使用的外生变量，它的列数必须和训练阶段的外生变量一致。
|
|
|
steps (int, optional, default=20): 预测期长度。
|
|
|
information_criterion (str, optional, default='aic'): 用于模型选择的信息准则，可以是 'aic' 或 'bic'。
|
|
|
|
|
|
Returns:
|
|
|
np.ndarray: 预测结果。
|
|
|
"""
|
|
|
model = pm.auto_arima(endog,
|
|
|
X=exog,
|
|
|
seasonal=True,
|
|
|
information_criterion=information_criterion)
|
|
|
|
|
|
pred = model.predict(n_periods=steps, X=exog_pred)
|
|
|
return pred
|
|
|
|
|
|
|
|
|
def run(input_data: pd.DataFrame,
|
|
|
forecast_target: str,
|
|
|
exog_columns: List[str],
|
|
|
steps: int = 20) -> pd.DataFrame:
|
|
|
"""
|
|
|
主运行函数，用以读取数据、训练模型、预测数据。
|
|
|
|
|
|
Args:
|
|
|
input_data (pd.DataFrame): 输入的时间序列数据。
|
|
|
forecast_target (str): 预测目标的列名。
|
|
|
exog_columns (List[str]): 外生变量的列名列表。
|
|
|
steps (int, optional, default=20): 预测步长
|
|
|
|
|
|
Returns:
|
|
|
pd.DataFrame: 预测结果的DataFrame对象。
|
|
|
"""
|
|
|
# 创建一个未来日期的索引，用于保存预测数据
|
|
|
future_index = pd.date_range(start=input_data.index.max() +
|
|
|
pd.Timedelta(days=1),
|
|
|
periods=steps)
|
|
|
|
|
|
# 创建一个用于保存预测外生变量的空数据帧
|
|
|
df_exog = pd.DataFrame(index=future_index)
|
|
|
|
|
|
# 循环每个外生变量，使用ARIMA模型进行训练和预测，然后将预测值保存到df_exog中
|
|
|
for exog in exog_columns:
|
|
|
pred = train_SARIMA_model(endog=input_data[exog], steps=steps)
|
|
|
df_exog[exog] = pred
|
|
|
|
|
|
# 使用ARIMA模型对目标变量进行训练和预测，注意这里将df_exog作为预测阶段的外生变量传入
|
|
|
pred = train_SARIMA_model(endog=input_data[forecast_target],
|
|
|
exog=input_data[exog_columns],
|
|
|
exog_pred=df_exog[exog_columns],
|
|
|
steps=steps,
|
|
|
information_criterion='bic')
|
|
|
|
|
|
# 根据预测值创建一个新的数据帧，用于保存预测的目标变量
|
|
|
forecast_df = pd.DataFrame(pred,
|
|
|
index=future_index,
|
|
|
columns=[forecast_target])
|
|
|
|
|
|
return forecast_df
|
