[Python] 공적분을 활용한 Pair Trading

import numpy as np
import pandas as pd
import statsmodels
from statsmodels.tsa.stattools import coint
import matplotlib.pyplot as plt
np.random.seed(42)

data = pd.read_csv('./datas/data_all.csv')
tickers = ['Shinhan', 'KB']

# 인덱스 설정
data.set_index('Date', inplace=True)

# 데이터 정렬
data.sort_values('Date', ascending=True, inplace=True) # ascending=True 오름차순, False 내림차순

pd.concat([data[tickers[0]], data[tickers[1]]], axis=1).plot(figsize=(15,7))
plt.ylabel('Stock Price');
plt.xlabel('Time');
title = tickers[0] + 'and ' + tickers[1]
plt.title(title)

cumm_rtn = (1 + data.pct_change()) # 일일수익률
cumm_rtn = cumm_rtn.dropna()
cumm_rtn = cumm_rtn.tail(100)  

X = cumm_rtn.Shinhan
Y = cumm_rtn.KB
pd.concat([X, Y], axis=1).plot(figsize=(15,7))
plt.ylabel('Cumulative Return');
plt.xlabel('Time');
title = 'Cummulative Plot of ' + tickers[0] + 'and ' + tickers[1]
plt.title(title)

(Y/X).plot(figsize=(15,7))
plt.axhline((Y/X).mean(), color='red', linestyle='--')
plt.xlabel('Time')
plt.legend(['Price Ratio', 'Mean'])
plt.show()

score, pvalue, _ = coint(X,Y)
print('Correlation: ' + str(X.corr(Y)))
print('Cointegration test p-value: ' + str(pvalue))

import FinanceDataReader as fdr
fdr.__version__

# 한국거래소 상장종목 전체
# 용도 : 코드와 종목명 가져오기
df_krx = fdr.StockListing('KRX')
df_krx.head()

n = 10 # 생성할 종목수 지정
df_krx_list = df_krx['Symbol'].head(n*2) # 임시로 2배까지 루프

# 코스피 종목 추출
# 용도 : DataReader 사용시 종목별 기간 조회 건수가 상이하다.
#        코스피 데이터 날짜를 기준으로 가져오기 위해 사용
strt_dt = '2014-01-01' # 시작일 지정
end_dt = '2018-12-31'  # 종료일 지정
kospi_df = fdr.DataReader('KS11', strt_dt,end_dt)
kospi_df.reset_index(inplace = True)
kospi_df.head()

stock_df = pd.DataFrame()
stock_df['Date'] = kospi_df['Date']
print("동기간 KOSPI 생성일수 : ", len(kospi_df['Date']))
normal_cnt = 0
err_cnt = 0


for code in df_krx_list:
    stock = df_krx[df_krx.Symbol == code]
    code_nm = list(stock.Name)[0]


    try:
        temp = fdr.DataReader(code, strt_dt, end_dt)

        # 데이터일수가 시장보다 작으면 skip(최근 상장 데이터로 판단)
        if len(temp) < len(kospi_df['Date']):
            err_cnt += 1
            print("skip : (",err_cnt,")", code, code_nm, strt_dt, end_dt, ", 건수 : ", len(temp))
            continue

        temp.reset_index(inplace = True)
        temp_df = pd.merge(temp[['Date','Close']], kospi_df[['Date']], on='Date', how='right')
        stock_df[code_nm] = temp_df.Close
        normal_cnt += 1
        print("정상 : (",normal_cnt,")", code, code_nm, strt_dt, end_dt, ", 건수 : ", len(temp), "->", len(stock_df))
    except:
        err_cnt += 1
        print("skip : (",err_cnt,")", code, code_nm, strt_dt, end_dt, ", 건수 : ", len(temp))

    if normal_cnt == n:
        print('총', n,'개 생성 설정 / ', normal_cnt, '개 생성 완료')
        break # n개 종목 생성시 종료

stock_df.head()

# 데이터 정렬
stock_df.sort_values('Date', ascending=True, inplace=True) # ascending=True 오름차순, False 내림차순
# 결측치 채우기
stock_df.fillna(method='ffill', inplace=True)
# 인덱스 설정
stock_df.set_index('Date', inplace=True)
data = stock_df

# p-value가 지정된 값보다 작은 pair 쌍을 반환한다.
def find_cointegrated_pairs(data):
    n = data.shape[1]
    score_matrix = np.zeros((n, n))
    pvalue_matrix = np.ones((n, n))
    keys = data.keys()
    all_pairs = []
    pairs = []

    # result
    stock1 = []
    stock2 = []
    pvalue_list = []
    check_95 = []
    check_98 = []

    for i in range(n):
        for j in range(i+1, n):
            S1 = data[keys[i]]
            S2 = data[keys[j]]
            result = coint(S1, S2)
            score = result[0]
            pvalue = result[1]
            score_matrix[i, j] = score
            pvalue_matrix[i, j] = pvalue


            if pvalue < 0.05:
                pairs.append((keys[i], keys[j]))
                check_95.append('Y')
            else:
                check_95.append('N')

            if pvalue < 0.02:
                check_98.append('Y')
            else:
                check_98.append('N')

            # result
            stock1.append(keys[i])
            stock2.append(keys[j])
            pvalue_list.append(pvalue)


    pair_pvalue = pd.DataFrame()
    pair_pvalue['s1'] = stock1
    pair_pvalue['s2'] = stock2
    pair_pvalue['pvalue'] = pvalue_list
    pair_pvalue['check_95'] = check_95
    pair_pvalue['check_98'] = check_98

    pair_pvalue.sort_values('pvalue', ascending=True, inplace=True) # ascending=True 오름차순

    return score_matrix, pvalue_matrix, pair_pvalue, pairs

# Heatmap
instrumentIds = list(data.columns.values)

scores, pvalues, pair_pvalue, pairs = find_cointegrated_pairs(data)
import seaborn
m = [0,0.2,0.4,0.6,0.8,1]
seaborn.heatmap(pvalues, xticklabels=instrumentIds,
                yticklabels=instrumentIds, cmap='RdYlGn_r',
                mask = (pvalues >= 0.95))
plt.show()

# 유의한 pair 출력
print(pairs)

pair_pvalue.head(5)

def zscore(series):
    return (series - series.mean()) / np.std(series)

zscore(ratios).plot()
plt.axhline(zscore(ratios).mean())
plt.axhline(1.0, color='red')
plt.axhline(-1.0, color='green')
plt.show()

# 가장 유의성이 높은 2개 종목을 추출한다.
s1_nm = 'CJ'
s2_nm = 'CJ씨푸드'
S1 = data[s1_nm]
S2 = data[s2_nm]

ratios = S1 / S2
cut = int(len(ratios)*0.7)
train = ratios[:cut]
test = ratios[cut:]

S1_train = S1.iloc[:cut]
S2_train = S2.iloc[:cut]
S1_test = S1.iloc[cut:]
S2_test = S2.iloc[cut:]

ratios_mavg5 = train.rolling(window=5,
                               center=False).mean()
ratios_mavg60 = train.rolling(window=60,
                               center=False).mean()
std_60 = train.rolling(window=60,
                        center=False).std()
zscore_60_5 = (ratios_mavg5 - ratios_mavg60)/std_60
plt.figure(figsize=(15,7))
plt.plot(train.index, train.values)
plt.plot(ratios_mavg5.index, ratios_mavg5.values)
plt.plot(ratios_mavg60.index, ratios_mavg60.values)
plt.legend(['Ratio','5d Ratio MA', '60d Ratio MA'])
plt.ylabel('Ratio')
plt.show()

plt.figure(figsize=(15,7))
zscore_60_5.plot()
plt.axhline(0, color='black')
plt.axhline(1.0, color='red', linestyle='--')
plt.axhline(-1.0, color='green', linestyle='--')
plt.legend(['Rolling Ratio z-Score', 'Mean', '+1', '-1'])
plt.show()

# Plot the ratios and buy and sell signals from z score
plt.figure(figsize=(15,7))
train[60:].plot()
buy = train.copy()
sell = train.copy()
buy[zscore_60_5>-1] = 0
sell[zscore_60_5<1] = 0
buy[60:].plot(color='g', linestyle='None', marker='^')
sell[60:].plot(color='r', linestyle='None', marker='^')
x1,x2,y1,y2 = plt.axis()
plt.axis((x1,x2,ratios.min(),ratios.max()))
plt.legend(['Ratio', 'Buy Signal', 'Sell Signal'])
plt.show()

# Plot the prices and buy and sell signals from z score
plt.figure(figsize=(18,9))

S1_log = S1_train[60:].map(lambda x : np.log(x))
S2_log = S2_train[60:].map(lambda x : np.log(x))

S1_log[60:].plot(color='b')
S2_log[60:].plot(color='c')
buyR = 0*S1_log.copy()
sellR = 0*S1_log.copy()
# When buying the ratio, buy S1 and sell S2
buyR[buy!=0] = S1_log[buy!=0]
sellR[buy!=0] = S2_log[buy!=0]
# When selling the ratio, sell S1 and buy S2
buyR[sell!=0] = S2_log[sell!=0]
sellR[sell!=0] = S1_log[sell!=0]
buyR[60:].plot(color='g', linestyle='None', marker='^')
sellR[60:].plot(color='r', linestyle='None', marker='^')
x1,x2,y1,y2 = plt.axis()
plt.axis((x1,x2,min(S1_log.min(),S2_log.min())-1,max(S1_log.max(),S2_log.max())+1))
plt.legend([s1_nm, s2_nm, 'Buy Signal', 'Sell Signal'])
plt.show()

# Trade using a simple strategy
def trade(S1, S2, window1, window2):

    # If window length is 0, algorithm doesn't make sense, so exit
    if (window1 == 0) or (window2 == 0):
        return 0

    # Compute rolling mean and rolling standard deviation
    ratios = S1/S2
    ma1 = ratios.rolling(window=window1,
                               center=False).mean()
    ma2 = ratios.rolling(window=window2,
                               center=False).mean()
    std = ratios.rolling(window=window2,
                        center=False).std()
    zscore = (ma1 - ma2)/std

    # Simulate trading
    # Start with no money and no positions
    money = 0
    countS1 = 0
    countS2 = 0
    for i in range(len(ratios)):
        # Sell short if the z-score is > 1
        if zscore[i] > 1:
            money += S1[i] - S2[i] * ratios[i]
            countS1 -= 1
            countS2 += ratios[i]
            print('Selling Ratio %s %s %s %s'%(money, ratios[i], countS1,countS2))
        # Buy long if the z-score is < 1
        elif zscore[i] < -1:
            money -= S1[i] - S2[i] * ratios[i]
            countS1 += 1
            countS2 -= ratios[i]
            print('Buying Ratio %s %s %s %s'%(money,ratios[i], countS1,countS2))
        # Clear positions if the z-score between -.5 and .5
        # -0.5~0.5 사이인 경우 수익일 경우 이익 실현
        elif abs(zscore[i]) < 0.5 and (money + S1[i] * countS1 + S2[i] * countS2) > 0 :
            money += S1[i] * countS1 + S2[i] * countS2
            countS1 = 0
            countS2 = 0
            print('Exit pos %s %s %s %s'%(money,ratios[i], countS1,countS2))


    return money

money = trade(S1_train, S2_train, 78, 5)
money

money = trade(S1_test, S2_test, 78, 5)
money

# train
length_scores = [trade(S1_train,
                S2_train, l, 5)
                for l in range(120)]
best_length = np.argmax(length_scores)
print ('Best window length:', best_length)

# Test
length_scores2 = [trade(S1_test,
                  S2_test,l,5)
                  for l in range(120)]
print (best_length, 'day window:', length_scores2[best_length])
best_length2 = np.argmax(length_scores2)
print (best_length2, 'day window:', length_scores2[best_length2])

plt.figure(figsize=(15,7))
plt.plot(length_scores)
plt.plot(length_scores2)
plt.xlabel('Window length')
plt.ylabel('Score')
plt.legend(['Training', 'Test'])
plt.show()