简单布林带策略实现

简单布林带策略实现
ID:383995

import pandas as pd
import matplotlib.pyplot as plt
import warnings
import tushare as ts
import seaborn as sns
sns.set()
warnings.simplefilter("ignore")


# 数据日期参数设置
start_time = '20180701'
end_time = '20190401'
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×md
def 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*md
    dn = mb-2*md
    df['Boll_up'] = up
    df['Boll_down'] = dn
    df['Boll_mid'] = mb
    return 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，不交易信号为0
def 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'])] = 2
    up_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'])] = -2
    down_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'])] = 1
    mid_status[(df['close'].shift(1) > df['Boll_mid'].shift(1)) & (df['close'] < df['Boll_mid'])] = -1
    mid_status.fillna(0, inplace=True)

    df['status'] = up_status['up_status']+down_status['down_status']+mid_status['mid_status']
    df.dropna(inplace=True)
    return df


class CapitalInfo:
    account = initial_account  # 账户金额
    quantity = 0  # 持有股票数


# 只买卖100股
def buy(c, price):
    if c.account > price*100:
        c.account -= price*100  # 每次买100股
        c.quantity += 100
    return c


def sell(c, price):
    if c.quantity > 100:
        c.account += price*100  # 每次卖200股
        c.quantity -= 100
    return c


# 全仓杀入杀出
def buy_all(c, price):
    q = 100*int(c.account/(price*100))
    if q > 0:
        c.account -= q*price
        c.quantity += q
    return c


def sell_all(c, price):
    if c.quantity > 0:
        c.account += c.quantity*price
        c.quantity = 0
    return c


def 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)/init
    strategy_data['benchmark'] = (strategy_data['close']-strategy_data['close'][0])/strategy_data['close'][0]
    return strategy_data


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