简单布林带策略实现
ID:383995
import pandas as pdimport matplotlib.pyplot as pltimport warningsimport tushare as tsimport seaborn as snssns.set()warnings.simplefilter("ignore")# 数据日期参数设置start_time = '0701'end_time = '0401'initial_account = 1000000# 布林带参数N = 20# 导入股票数据def get_ts_api():return ts.pro_api("your token here")def get_index_data(start, end):pro = get_ts_api()index_data = pro.index_daily(ts_code='000300.SH', start_date=start, end_date=end)return index_data# 指标计算# md=N日内的收盘价之和÷N# md=平方根N日的（C－MA）的两次方之和除以N# mb=N日的ma# up=MB+2×md# dn=MB－2×mddef calculate_boll(start, end):index_data = get_index_data(start, end)index_data.sort_index(ascending=False, inplace=True)df = index_data[['trade_date', 'open', 'close', 'high', 'low']]df['trade_date'] = pd.to_datetime(df['trade_date'], format="%Y%m%d")df.set_index('trade_date', drop=True, inplace=True)# 移动平均线和方差的计算ma = df['close'].rolling(N).mean()md = df['close'].rolling(N).std()# Boll带相关指标的计算mb = ma.copy()up = mb+2*mddn = mb-2*mddf['Boll_up'] = updf['Boll_down'] = dndf['Boll_mid'] = mbreturn df# 策略开发# 当Boll价格向上突破中轨时，增加持仓；当Boll价格突破上轨时，全仓杀入# 当Boll价格向下突破中轨时，减少持仓；当Boll价格突破下轨时，全仓卖出# 参数解释# up_status：价格是否向上突破上轨，取值为2,0# down_status: 价格是否向下突破下轨，取值分别为-2,0# mid_status：价格是否向上或向下穿过中轨，取值分别为1,-1,0# status: 将以上参数相加，得到全仓杀入、买入100、卖出100、全仓杀出信号分别为2,1，-1，-2，不交易信号为0def get_status(start, end):df = calculate_boll(start, end)up_status = pd.DataFrame(index=df.index, columns=['up_status'])up_status[(df['close'].shift(1) < df['Boll_up'].shift(1)) & (df['close'] > df['Boll_up'])] = 2up_status.fillna(0, inplace=True)down_status = pd.DataFrame(index=df.index, columns=['down_status'])down_status[(df['close'].shift(1) > df['Boll_down'].shift(1)) & (df['close'] < df['Boll_down'])] = -2down_status.fillna(0, inplace=True)mid_status = pd.DataFrame(index=df.index, columns=['mid_status'])mid_status[(df['close'].shift(1) < df['Boll_mid'].shift(1)) & (df['close'] > df['Boll_mid'])] = 1mid_status[(df['close'].shift(1) > df['Boll_mid'].shift(1)) & (df['close'] < df['Boll_mid'])] = -1mid_status.fillna(0, inplace=True)df['status'] = up_status['up_status']+down_status['down_status']+mid_status['mid_status']df.dropna(inplace=True)return dfclass CapitalInfo:account = initial_account # 账户金额quantity = 0 # 持有股票数# 只买卖100股def buy(c, price):if c.account > price*100:c.account -= price*100 # 每次买100股c.quantity += 100return cdef sell(c, price):if c.quantity > 100:c.account += price*100 # 每次卖200股c.quantity -= 100return c# 全仓杀入杀出def buy_all(c, price):q = 100*int(c.account/(price*100))if q > 0:c.account -= q*pricec.quantity += qreturn cdef sell_all(c, price):if c.quantity > 0:c.account += c.quantity*pricec.quantity = 0return cdef get_strategy_data(start, end, init):c = CapitalInfo()strategy_data = get_status(start, end)# 账户初始化c = buy(c, strategy_data['close'][0])strategy = pd.DataFrame(index=strategy_data.index, columns=['strategy'])strategy.iloc[0, 0] = c.account+c.quantity * strategy_data['close'][0]# 遍历每个交易日进行交易# 全仓杀入杀出for i in strategy_data.index:if strategy_data.loc[i, 'status'] == 2:c = buy_all(c, strategy_data.loc[i, 'close'])if strategy_data.loc[i, 'status'] == 1:c = buy(c, strategy_data.loc[i, 'close'])if strategy_data.loc[i, 'status'] == -1:c = sell(c, strategy_data.loc[i, 'close'])if strategy_data.loc[i, 'status'] == -2:c = sell_all(c, strategy_data.loc[i, 'close'])strategy.loc[i, 'strategy'] = c.account+c.quantity * strategy_data.loc[i, 'close']strategy_data = pd.concat([strategy_data, strategy], axis=1)strategy_data['strategy_revenue'] = (strategy_data['strategy']-init)/initstrategy_data['benchmark'] = (strategy_data['close']-strategy_data['close'][0])/strategy_data['close'][0]return strategy_datadef main(start, end, init):strategy_data = get_strategy_data(start, end, init)plt.figure(figsize=(20, 15))plt.plot(strategy_data['strategy_revenue'], 'r-')plt.plot(strategy_data['benchmark'], 'b-')plt.legend(['strategy', 'HS300'])plt.xlabel('Date')plt.ylabel('Revenue')plt.title("Boll_strategy")plt.show()
