|
|
from typing import List, Union
|
|
|
|
|
|
import numpy as np
|
|
|
import pandas as pd
|
|
|
import pmdarima as pm
|
|
|
|
|
|
|
|
|
def train_ARIMA_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:
|
|
|
"""
|
|
|
使用ARIMA模型对时间序列数据进行预测。
|
|
|
|
|
|
Args:
|
|
|
endog (Union[np.ndarray, pd.Series]): 要分析的时间序列数据。
|
|
|
exog (Union[np.ndarray, pd.DataFrame], optional): 用于改进ARIMA模型的外生变量。默认为None。
|
|
|
exog_pred (Union[np.ndarray, pd.DataFrame], optional): 预测期间的外生变量，必须与训练期间的外生变量列数一致。默认为None。
|
|
|
steps (int, optional, default=20): 预测期的长度。
|
|
|
information_criterion (str, optional, default='aic'): 选择模型的信息准则，'aic'或'bic'。
|
|
|
|
|
|
Returns:
|
|
|
np.ndarray: 预测结果。
|
|
|
"""
|
|
|
model = pm.auto_arima(endog,
|
|
|
X=exog,
|
|
|
seasonal=False,
|
|
|
information_criterion=information_criterion)
|
|
|
|
|
|
pred = model.predict(n_periods=steps, X=exog_pred)
|
|
|
return pred
|
|
|
|
|
|
|
|
|
def ARIMA_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_ARIMA_model(endog=input_data[exog], steps=steps)
|
|
|
df_exog[exog] = pred
|
|
|
|
|
|
# 使用ARIMA模型对目标变量进行训练和预测，注意这里将df_exog作为预测阶段的外生变量传入
|
|
|
pred = train_ARIMA_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 |
