預期效果
根據輸入的數據爬取一段時期內每天的股價信息(以上證指數爲例),根據15日均價制定簡易的股票交易策略,並對結果作圖展示。
代碼實現
import json
import requests
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
def run():
''' 主程序, 用來調度各個重要流程 '''
kline = load_sse()
df = init_df(kline)
df = strategy(df)
df = backtest(df)
draw(df,days)
df.to_csv('result.csv', index = False)
def load_sse():
''' 獲取上交所的上證指數K線, 最近N個交易日數據 '''
response = requests.get(
# 'http://yunhq.sse.com.cn:32041/v1/sh1/dayk/000001?callback=jQuery111205234775875526079_1542185571865&select=date%2Copen%2Chigh%2Clow%2Cclose%2Cvolume&begin=-2000&end=-1&_=1542185571881',
'http://yunhq.sse.com.cn:32041/v1/sh1/dayk/000001?callback=jQuery111205234775875526079_1542185571865&select=date%2Copen%2Chigh%2Clow%2Cclose%2Cvolume&begin=-'+ begin +'&end=-'+ end +'&_=1542185571881',
headers={'Referer': 'http://www.sse.com.cn/market/price/trends/'}
)
# 針對結果進行格式處理
json_str = response.text[42:-1]
data = json.loads(json_str)
return data['kline']
def init_df(kline):
''' 根據K線數據,創建含有日期與收盤價的矩陣 '''
df = pd.DataFrame({})
df['date'] = [x[0] for x in kline]
#kline中包含日期、開盤價、最高價、最低價、收盤價等信息
df['close'] = [x[1] for x in kline]
return df
def strategy(df):
# 連續15天數據,計算平均值,作爲當天的平均價格指標
window_size = 15
df['avg'] = df['close'].rolling(window_size).apply(lambda x: sum(x) / len(x))
def avg_buy(x):
''' 做多策略 '''
min_percent = 0.995
max_percent = 1.005
# 追漲,當我們的價格超過了均線一定程度時
if (x[1] / x[0]) < min_percent:
return 'open buy'
# 殺跌,當我們的價格低於均線一定程度時
if (x[1] / x[0]) > max_percent:
return 'close buy'
# 其他情況不操作
return 'wait'
# df['action'] = avg_buy([df['close'], df['avg']])
df['action'] = df[['close', 'avg']].apply(avg_buy, axis=1)
return df
def backtest(df):
''' 迴歸測試 '''
global shares, cash
amount = 1000000
shares = 0
cash = amount
def run_strategy(row):
''' 把每天的數據執行策略 '''
global shares, cash
action = row['action']
close = row['close']
# 資產 = 現金 + 股票價值
liquidate = cash + shares * close
message = 'nothing'
# 策略要求開倉做多,而且當前空倉時,做多
if action == 'open buy' and shares == 0:
shares = int(cash / close)
cash -= shares * close
message = 'open buy ' + str(shares)
# 策略要求平倉,而且當前有倉時,平掉
if action == 'close buy' and shares > 0:
message = 'close buy ' + str(shares)
cash += shares * close
shares = 0
return [message, shares, cash, liquidate]
rows = df[['close', 'action']].apply(run_strategy, axis=1)
df['message'], df['shares'], df['cash'], df['liquidate'] = zip(*rows)
return df
def draw(df,days):
''' 畫圖 '''
# 創建畫板
fig = plt.figure(figsize=(10, 5))
# 準備橫座標
count = df.count()['close']
index = np.arange(count)
df['index'] = index
# 設置橫座標的刻度與顯示標籤
limit = days
plt.xticks(index[::limit], df['date'][::limit])
# 收盤價與資產的兩套座標系
ax_close = plt.gca()
ax_liquidate = ax_close.twinx()
# 畫收盤價曲線
ax_close.set(xlabel='Date', ylabel='close')
l_close, = ax_close.plot(index, df['close'], 'black', label='close')
l_avg, = ax_close.plot(index, df['avg'], 'pink', label='avg')
# 畫資產曲線
ax_liquidate.set(ylabel = 'liquidate')
l_liquidate, = ax_liquidate.plot(index, df['liquidate'], 'blue', label='liquidate')
def drawAction(row):
if row['message'] == 'nothing':
return
color = ''
marker = 'o'
size = 12
if row['action'] == 'open buy':
color='r'
if row['action'] == 'close buy':
color='g'
ax_close.scatter(row['index'], row['close'], s=size, color=color, zorder=2, marker=marker)
df[['index', 'action', 'message', 'close']].apply(drawAction, axis=1)
# 給兩條線都提供一個圖例說明
plt.legend(handles=[l_close, l_avg, l_liquidate])
plt.show()
if __name__ == '__main__':
begin=input('從前多少天:')
end=input('到最近幾天:')
days=input('橫座標日期間隔天數:')
days=int(days)
run()