import tushare as ts
pro = ts.pro_api(' ++++++++++++++++ ')  # 到tushare官网注册账号获取token
ts_code = '000001'  
ts_code = ts_code+'.SZ'
ts_code
'000001.SZ'
df = pro.daily(ts_code=ts_code, start_date='20210101', end_date='20210531')
df.head()

import numpy as np
import pandas as pd
import statsmodels.tsa.api as smt
from statsmodels.tsa.ar_model import AR
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False #当坐标轴有负号的时候可以显示负号
import warnings
warnings.filterwarnings("ignore")
def draw_ac_pac(series, nlags=30):
    fig = plt.figure(figsize=(10,5))
    # 设置子图
    ts_ax = fig.add_subplot(311)
    acf_ax = fig.add_subplot(312)
    pacf_ax = fig.add_subplot(313)
    # 绘制图像
    ts_ax.set_title('time series')
    acf_ax.set_title('autocorrelation coefficient')
    pacf_ax.set_title('partial autocorrelation coefficient')
    ts_ax.plot(series)
    smt.graphics.plot_acf(series, lags=nlags, ax=acf_ax)
    smt.graphics.plot_pacf(series, lags=nlags, ax=pacf_ax)
    # 自适应布局
    plt.tight_layout()
    plt.show()
# 将日期设为索引
df.index = pd.to_datetime(df['trade_date'], format='%Y%m%d')
# 可视化自相关和偏自相关系数
draw_ac_pac(df['close'], nlags=30)

# 划分训练数据和测试数据
train = df['close'][:-10]
test = df['close'][-10:]
model_fit = AR(train).fit()
params = model_fit.params
p = model_fit.k_ar  # 即时间序列模型中常见的p，即AR(p), ARMA(p,q), ARIMA(p,d,q)中的p。
history = train[-p:]   
history = np.hstack(history).tolist()
test = np.hstack(test).tolist()
predictions = []
for t in range(len(test)):
    lag = history[-p:]
    yhat = params[0]
    for i in range(p):
        yhat += params[i+1] * lag[p-1-i]
    predictions.append(yhat)
    obs = test[t]
    history.append(obs)
print('MSE指标：',np.mean((np.array(test) - np.array(predictions))**2))  # 得到mean_squared_error, MSE
fig = plt.figure(figsize=(15,5))
plt.plot(test)
plt.plot(predictions, color='r')
plt.show()
MSE指标： 0.988321573905643

import ffn
df_plot = pd.DataFrame({'predictions':predictions})
result=ffn.calc_total_return(df_plot['predictions'])
ann_result=ffn.annualize(result,10,one_year=250)  # 根据10天数据计算交易日250天的年化收益率
ann_result
-0.9898424754209083