Quantrum ROC(200) & KST strategy test

quantrum.roc200.py

import talib  # библиотека для технического анализа
import numpy as np


# функция, выполняемая перед началом тестирования
def initialize(context):
    # сохраняем актив, с которым будем работать
    context.asset = symbol('SPY')
    context.strategy = 'roc_sma'
    set_benchmark(symbol('SPY'))
    
    # ставим событие для ребалансировки на открытии рынка
    schedule_function(rebalance, date_rules.every_day(), time_rules.market_open())
    

# ежедневная ребалансировка
def rebalance(context, data):
    # история для расчета средних за последние 400 торговых дней
    price_hist = data.history(context.asset, 'close', 400, '1d')
    
    # стратегия
    if context.strategy == 'price_sma200':
        # средние: длинная и короткая
        ma_long = talib.SMA(price_hist, timeperiod=200)
        ma_short = price_hist
        
        # сигнал на положении
        allow = ma_short[-1] >= ma_long[-1]
    elif context.strategy == 'roc200':
        price_hist = data.history(context.asset, 'close', 400, '1d')
        ma_long = talib.ROC(price_hist, timeperiod=200)
        
        # сигнал на положении
        allow = 0 < ma_long[-1]
        record(roc=ma_long[-1])
    elif context.strategy == 'roc200smooth':
        # средние: длинная и короткая
        roc200 = talib.ROC(price_hist, timeperiod=200)
        ma_long = (roc200 + np.roll(roc200, 1)*2 + np.roll(roc200, 2)*2 + np.roll(roc200, 3)) / 6
        
        # сигнал на положении
        allow = 0 < ma_long[-1]
        record(sroc=ma_long[-1])
    elif context.strategy == 'tsi':
        prices = data.history(context.asset, ['open', 'high', 'low', 'close'], 400, '1d')
        tsi = get_tsi(prices, 10, 50)  # 10, 50

        roc200 = talib.ROC(price_hist, timeperiod=200)
        for_sign = roc200.copy()
        tsi *= np.sign(for_sign)       
        
        ma_long = tsi
        
        # сигнал на положении
        allow = 0.5 < ma_long[-1]
        record(tsi=ma_long[-1])
    elif context.strategy == 'kst':
        prices = data.history(context.asset, ['close'], 400, '1d')
        kst, kst_signal = get_kst(prices, [10, 15, 20, 30], [10, 10, 10, 15], 9)
        
        # сигнал на положении
        allow = 0. <= kst[-1]
        record(kst=kst[-1])
    elif context.strategy == 'kst_signal':
        prices = data.history(context.asset, ['close'], 400, '1d')
        kst, kst_signal = get_kst(prices, [10, 15, 20, 30], [10, 10, 10, 15], 9)
        
        # сигнал на положении
        allow = 0. <= (kst - kst_signal)[-1]
        record(kst_hist=(kst - kst_signal)[-1])
    elif context.strategy == 'roc_sma':
        prices = data.history(context.asset, ['close'], 400, '1d')
        roc_sma = talib.SMA(talib.ROC(prices.close, timeperiod=5), timeperiod=200)
        
        # сигнал на положении
        allow = 0. <= roc_sma[-1]
        record(roc_sma=roc_sma[-1])
    elif context.strategy == 'roc5x200':
        ma_short = talib.ROC(price_hist, timeperiod=5)
        ma_long = talib.ROC(price_hist, timeperiod=200)
        
        # сигнал на положении
        allow = 0 < ma_short[-1] - ma_long[-1]
        record(rocx=ma_short[-1] - ma_long[-1])
    elif context.strategy == 'sma50_sma200':
        # средние: длинная и короткая
        ma_long = talib.SMA(price_hist, timeperiod=200)
        ma_short = talib.SMA(price_hist, timeperiod=50)
        
        # сигнал на пересечении
        #allow = ma_short[-1] >= ma_long[-1] and ma_short[-1] < ma_long[-1]
        # сигнал на положении
        allow = 0 < ma_short[-1] - ma_long[-1]
    else:
        print('Неизвестная стратегия')
        return
    
    record(allow=allow, price=price_hist[-1])
    
    # проверяем возможность торговли активом
    if data.can_trade(context.asset):
        if allow:
            # покупаем актив на 100% портфеля, если разрешено
            order_target_percent(context.asset, 1.)
        else:
            order_target_percent(context.asset, 0.)


def get_tsi(df, short, long):
    ratio = (df.close - df.close.shift(short)).abs() / talib.ATR(df.high.values, df.low.values, df.close.values, timeperiod=short)
    tsi = talib.SMA(talib.SMA(ratio.values, timeperiod=short), timeperiod=long) - 1  # substruct 1
    tsi[tsi < 0] = 0
    #print("Min: {0} Max: {1}".format(tsi[~np.isnan(tsi)].min(), tsi[~np.isnan(tsi)].max()))    
    return tsi

def get_kst(df, roc, sma, signal):
    kst = None
    for i, v in enumerate(roc):
        data = talib.SMA(talib.ROC(df.close, timeperiod=v), timeperiod=sma[i])
        if kst is None:
            kst = data * (i + 1)
        else:
            kst += data * (i + 1)
    return kst, talib.SMA(kst, timeperiod=signal)