import pandas as pd
import numpy as npimport matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']        # 字体设置
import matplotlib
matplotlib.rcParams['axes.unicode_minus']=False    # 负号显示问题from arch.unitroot import ADF                     # pip install arch
import statsmodels.api as sm

#导入数据
import tushare as ts
#后面自己研究何种股票，何种方式
data1 = ts.get_k_data('002966', '2013-06-01', '2021-12-31')[['date','close']]
data2 = ts.get_k_data('600919', '2013-06-01', '2021-12-31')[['date','close']]
#建立新的数据
sh = pd.merge(data1,data2,on='date')
sh.set_index('date',inplace = True)
stocks_pair = ['601988','600000']
sh.columns = stocks_pair
sh.head()

#计算相关性，这里使用日线级别，计算相关性
sh[['601988','600000']].corr()

#计算价差
sh['jc'] = sh['600000']-sh['601988']
sh.head()

#计算新的价差数据组
sh = sh['jc'].reset_index()
sh.set_index('date',inplace=True)
sh.index = pd.to_datetime(sh.index)
sh.head()

#计算每周收益，百分比
sh['ret'] = sh['jc'].pct_change()
#计算每周收益，固定值
sh['ret2']  = sh['jc']-sh['jc'].shift(1) 

sh['jc'].plot()

#设计信号
sh['sig'] = sh['jc'].apply(lambda x:-1 if(x>8) else 0)#价差超过-7，买601988，卖60000
sh['sig2'] = sh['jc'].apply(lambda x:1 if(x<5) else 0)#价差低于-9，卖601988，60000

#未出现信号，则继续持有
sh[['sig','sig2']].replace(0,np.nan,inplace=True)
sh[['sig','sig2']].bfill(inplace=True)

#计算每次交易收益
sh['rr'] = (sh['ret2']*sh['sig2']+sh['ret2']*sh['sig2'])
#收益画图。图为累计可以获得的价差累加值收益
sh['lr'] = sh['rr'].cumsum()
sh['lr'].plot()

#粗略计算收益率,9年获利0.08
0.66/(2.15+5.815)

#转换为周线级别
period_type = 'W'
sh2=pd.DataFrame()
#sh2[['收盘价','开盘价','最高价','最低价']] = 0
sh2['收盘价'] = sh['jc'].resample(period_type).last()
sh2['开盘价'] = sh['jc'].resample(period_type).first()
sh2['最高价'] = sh['jc'].resample(period_type).max()
sh2['最低价'] = sh['jc'].resample(period_type).min()
sh2.head()#计算每周收益，百分比
sh2['ret'] = sh2['收盘价'].pct_change()
#计算每周收益，固定值
sh2['ret2']  = sh2['收盘价']-sh2['收盘价'].shift(1) #设计信号
sh2['sig'] = sh2['收盘价'].apply(lambda x:-1 if(x>8) else 0)
sh2['sig2'] = sh2['收盘价'].apply(lambda x:1 if(x<7) else 0)#未出现信号，则继续持有
sh2[['sig','sig2']].replace(0,np.nan,inplace=True)
sh2[['sig','sig2']].bfill(inplace=True)
#计算每次交易收益
sh2['rr'] = (sh2['ret2']*sh2['sig2']+sh2['ret2']*sh2['sig2'])
#收益画图。图为累计可以获得的价差累加值收益
sh2['lr'] = sh2['rr'].cumsum()
sh2['lr'].plot()