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()