[Python] 평균회귀를 활용한 Long/Short 전략

import FinanceDataReader as fdr
fdr.__version__

%matplotlib inline
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

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

# 코스피 종목 추출
# 용도 : 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)

def stock_reader(kospi_df, code_list, n=0):

    if n == 0:
        n = len(code_list)

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

    for code in code_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
            code_nm_list.append(code_nm)
            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.sort_values('Date', ascending=True, inplace=True) # ascending=True 오름차순, False 내림차순
        # 결측치 채우기
        stock_df.fillna(method='ffill', inplace=True)

    return stock_df, code_nm_list

code_list = ['055550']
stock_df, code_nm_list = stock_reader(kospi_df, code_list)

stock_df.head(3)

prices = pd.DataFrame(stock_df[code_nm_list[0]])

code_nm = code_nm_list[0]
prices['mu'] = [prices[code_nm][:i].mean() for i in range(len(prices))]
# Plot the price and the moving average
plt.figure(figsize=(15,7))
plt.plot(prices[code_nm])
plt.plot(prices['mu']);
plt.show()

zscores = [(prices[code_nm][i] - prices['mu'][i]) / np.std(prices[code_nm][:i]) for i in range(len(prices))]

# Start with no money and no positions
money = 0
count = 0
for i in range(len(prices)):
    # Sell short if the z-score is > 1
    if zscores[i] > 1:
        money += prices[code_nm][i]
        count -= 1
    # Buy long if the z-score is < 1
    elif zscores[i] < -1:
        money -= prices[code_nm][i]
        count += 1
    # Clear positions if the z-score between -.5 and .5
    elif abs(zscores[i]) < 0.5:
        money += count*prices[code_nm][i]
        count = 0
print(money)

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

stock_df, code_nm_list = stock_reader(kospi_df, df_krx_list, n)

stock_df.set_index('Date', inplace=True)
data = stock_df
data.head(3)

data = np.array(data)

# 첫 주의 수익률
wreturns = (data[4] - data[0])/data[0]  # 4일 = 1 weak
# Rank
order = wreturns.argsort()
ranks = order.argsort()

# 첫 주 이후, 한달 수익률
mreturns = (data[5+4*4] - data[5])/data[5]  # 첫주이후 5day, 한달 후 5일 + 4day*4weak
order2 = mreturns.argsort()
ranks2 = order2.argsort()

# Plot (가로 : 첫주 수익률, 세로 : 첫주 이후 한달 수익률)
plt.figure(figsize=(15,7))
plt.scatter(wreturns, mreturns)
plt.xlabel('Returns for the first week')
plt.ylabel('Returns for the following month')

# 비교할 종목간 첫주 수익률 기준으로 랭킹 top 20%, bottom 20%를 찾는다.
# ranks는 첫주 수익률 순위를 의미(낮을 수록 1, 낮은걸 long한다.)
# ranks는 첫주 이후 한달 수익률 순위를 의미

top_cnt = int(len(ranks) * 0.2)
bottom_cnt = int(len(ranks) * 0.8)

longs = np.array([int(x < top_cnt)for x in ranks])
shorts = np.array([int(x >= bottom_cnt) for x in ranks])
print('Going long in:', [code_nm_list[i] for i in range(len(code_nm_list)) if longs[i]])
print('Going short in:', [code_nm_list[i] for i in range(len(code_nm_list)) if shorts[i]])

# Resolve all positions and calculate how much we would have earned
print('Yield:', sum((data[-1] - data[4])*(longs - shorts)))