使用ARIMA预测股票未来走势

from fileinput import close
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from statsmodels.tsa.arima_model import ARIMA
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
from statsmodels.stats.diagnostic import acorr_ljungbox

import statsmodels.api as sm
import matplotlib.font_manager as fm
import mplfinance as mpf

def arima_prediction(df,forecast_steps):

# 处理缺失值（如果有）
df = df.fillna(df.bfill())

# sm.graphics.tsa.plot_acf(ts,lags=30)
# sm.graphics.tsa.plot_pacf(ts,lags=30)

# 拟合ARIMA模型
p, d, q = 2, 0, 2  # 设置ARIMA模型的阶数（可以根据实际数据调整）
model = ARIMA(df, order=(p, d, q))
model_fit = model.fit()

# 预测未来1个月的K线数据

forecast= model_fit.forecast(steps=forecast_steps)
return forecast

def main():
code = '000001.SZ'
df = pd.read_csv("D:\csv_day_hq\"+code+".csv")
df['pushDate'] = pd.to_datetime(df['pushDate'])

df_p = df[df['pushDate']<='2023-08-25'].tail(250)
df_p.set_index('pushDate', inplace=True)

# print(df_p)
forecast_steps = 30
forecast_close = arima_prediction(df_p['close'],30)
forecast_open = arima_prediction(df_p['open'],30)
forecast_high = arima_prediction(df_p['high'],30)
forecast_low = arima_prediction(df_p['low'],30)

# 生成日期序列（未来1个月的日期）
last_date = df_p.index[-1]
print(last_date)
date_range = pd.date_range(start=last_date + pd.DateOffset(days=1), periods=forecast_steps, freq='D')
print(date_range)

# 创建未来1个月K线走势的DataFrame
forecast_df = pd.DataFrame({"pushDate":date_range,'open': forecast_open[0],'high': forecast_high[0],'low': forecast_low[0],'close': forecast_close[0]})
forecast_df.set_index('pushDate', inplace=True)
print(forecast_df)
lishi_df = df[['pushDate','open','high','low','close']]
lishi_df = lishi_df[lishi_df.shape[0]-100:]
lishi_df.set_index('pushDate', inplace=True)
forecast_df = pd.concat([lishi_df,forecast_df],ignore_index=False)
print(forecast_df)
# 自定义涨跌颜色
custom_colors = mpf.make_marketcolors(
    up='red',       # 上涨的颜色
    down='green',   # 下跌的颜色
    edge='black',   # K线边缘颜色
    wick='black',   # K线上下影线颜色
    volume='blue'   # 成交量条颜色
)

# 设置自定义样式
style = mpf.make_mpf_style(base_mpl_style="seaborn", marketcolors=custom_colors)
mpf.plot(forecast_df,type='candle', style=style, title='Kline', ylabel='price')

main()