使用 arima算法 基于Python 帮我写一个回归类的代码实现 数据划分为训练集与测试集
时间: 2024-03-22 07:42:56 浏览: 15
好的,我可以为您提供一个基于Python的ARIMA回归类的代码实现,并将数据划分为训练集和测试集。下面是代码实现:
``` python
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from statsmodels.tsa.arima.model import ARIMA
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
class ARIMAModel:
def __init__(self, p, d, q):
self.p = p
self.d = d
self.q = q
def fit(self, data):
self.model = ARIMA(data, order=(self.p, self.d, self.q))
self.model_fit = self.model.fit()
def predict(self, start, end):
return self.model_fit.predict(start=start, end=end)
def split_data(data, test_size):
train_size = int(len(data) * (1 - test_size))
train, test = data[0:train_size], data[train_size:]
return train, test
# 生成示例数据
data = pd.Series([0.1, 0.5, 0.3, 0.7, 0.2, 0.8, 0.6, 0.9, 0.4, 1.0])
# 划分训练集和测试集
train, test = split_data(data, test_size=0.2)
# ARIMA模型拟合
model = ARIMAModel(p=1, d=1, q=0)
model.fit(train)
# 预测
predictions = model.predict(start=len(train), end=len(train)+len(test)-1)
# 计算均方误差
mse = mean_squared_error(test, predictions)
print('Mean Squared Error: %.3f' % mse)
# 可视化预测结果和真实结果
plt.plot(test, color='blue', label='Actual')
plt.plot(predictions, color='red', label='Predicted')
plt.legend()
plt.show()
```
在这个代码实现中,我们首先定义了一个ARIMAModel类,其中包含了p、d、q三个参数,分别代表ARIMA模型中的自回归、差分和移动平均的阶数。我们还定义了fit()和predict()两个方法,用于拟合ARIMA模型和进行预测。
接下来,我们使用split_data()函数将数据划分为训练集和测试集。然后,我们使用ARIMAModel类对训练集进行拟合,使用predict()方法对测试集进行预测,并计算均方误差。
最后,我们使用matplotlib库可视化预测结果和真实结果。您可以根据自己的需要修改示例数据和参数,并使用该代码实现您自己的ARIMA回归模型。