Applying Research

from sklearn.decomposition import PCA
from arch.unitroot.cointegration import engle_granger
from statsmodels.tsa.stattools import adfuller
import numpy as np
from matplotlib import pyplot as plt

def initialize(self):

    #1. Required: Five years of backtest history
    self.set_start_date(2014, 1, 1)

    #2. Required: Alpha Streams Models:
    self.set_brokerage_model(BrokerageName.ALPHA_STREAMS)

    #3. Required: Significant AUM Capacity
    self.set_cash(1000000)

    #4. Required: Benchmark to SPY
    self.set_benchmark("SPY")

    self.assets = ["SCHO", "SHY"]
    
    # Add Equity ------------------------------------------------ 
    for i in range(len(self.assets)):
        self.add_equity(self.assets[i], Resolution.MINUTE).symbol
        
    # Instantiate our model
    self.recalibrate()
    
    # Set a variable to indicate the trading bias of the portfolio
    self.state = 0
    
    # Set Scheduled Event Method For Kalman Filter updating.
    self.schedule.on(self.date_rules.week_start(), 
        self.time_rules.at(0, 0), 
        self.recalibrate)
    
    # Set Scheduled Event Method For Kalman Filter updating.
    self.schedule.on(self.date_rules.every_day(), 
        self.time_rules.before_market_close("SHY"), 
        self.every_day_before_market_close)
        
        
def recalibrate(self):
    qb = self
    history = qb.history(self.assets, 252*2, Resolution.DAILY)
    if history.empty: return
    
    # Select the close column and then call the unstack method
    data = history['close'].unstack(level=0)
    
    # Convert into log-price series to eliminate compounding effect
    log_price = np.log(data)
    
    ### Get Cointegration Vectors
    # Get the cointegration vector
    coint_result = engle_granger(log_price.iloc[:, 0], log_price.iloc[:, 1], trend="c", lags=0)
    coint_vector = coint_result.cointegrating_vector[:2]
    
    # Get the spread
    spread = log_price @ coint_vector
    
    ### Kalman Filter
    # Initialize a Kalman Filter. Using the first 20 data points to optimize its initial state. We assume the market has no regime change so that the transitional matrix and observation matrix is [1].
    self.kalman_filter = KalmanFilter(transition_matrices = [1],
                        observation_matrices = [1],
                        initial_state_mean = spread.iloc[:20].mean(),
                        observation_covariance = spread.iloc[:20].var(),
                        em_vars=['transition_covariance', 'initial_state_covariance'])
    self.kalman_filter = self.kalman_filter.em(spread.iloc[:20], n_iter=5)
    (filtered_state_means, filtered_state_covariances) = self.kalman_filter.filter(spread.iloc[:20])
    
    # Obtain the current Mean and Covariance Matrix expectations.
    self.current_mean = filtered_state_means[-1, :]
    self.current_cov = filtered_state_covariances[-1, :]
    
    # Initialize a mean series for spread normalization using the Kalman Filter's results.
    mean_series = np.array([None]*(spread.shape[0]-20))
    
    # Roll over the Kalman Filter to obtain the mean series.
    for i in range(20, spread.shape[0]):
        (self.current_mean, self.current_cov) = self.kalman_filter.filter_update(filtered_state_mean = self.current_mean,
                                                                filtered_state_covariance = self.current_cov,
                                                                observation = spread.iloc[i])
        mean_series[i-20] = float(self.current_mean)
    
    # Obtain the normalized spread series.
    normalized_spread = (spread.iloc[20:] - mean_series)
    
    ### Determine Trading Threshold
    # Initialize 50 set levels for testing.
    s0 = np.linspace(0, max(normalized_spread), 50)
    
    # Calculate the profit levels using the 50 set levels.
    f_bar = np.array([None]*50)
    for i in range(50):
        f_bar[i] = len(normalized_spread.values[normalized_spread.values > s0[i]])               / normalized_spread.shape[0]
        
    # Set trading frequency matrix.
    D = np.zeros((49, 50))
    for i in range(D.shape[0]):
        D[i, i] = 1
        D[i, i+1] = -1
        
    # Set level of lambda.
    l = 1.0
    
    # Obtain the normalized profit level.
    f_star = np.linalg.inv(np.eye(50) + l * D.T@D) @ f_bar.reshape(-1, 1)
    s_star = [f_star[i]*s0[i] for i in range(50)]
    self.threshold = s0[s_star.index(max(s_star))]
    
    # Set the trading weight. We would like the portfolio absolute total weight is 1 when trading.
    self.trading_weight = coint_vector / np.sum(abs(coint_vector))
    
        
def every_day_before_market_close(self):
    qb = self
    
    # Get the real-time log close price for all assets and store in a Series
    series = pd.Series()
    for symbol in qb.securities.Keys:
        series[symbol] = np.log(qb.securities[symbol].close)
        
    # Get the spread
    spread = np.sum(series * self.trading_weight)
    
    # Update the Kalman Filter with the Series
    (self.current_mean, self.current_cov) = self.kalman_filter.filter_update(filtered_state_mean = self.current_mean,
                                                                        filtered_state_covariance = self.current_cov,
                                                                        observation = spread)
        
    # Obtain the normalized spread.
    normalized_spread = spread - self.current_mean

    # ==============================
    
    # Mean-reversion
    if normalized_spread < -self.threshold:
        orders = []
        for i in range(len(self.assets)):
            orders.append(PortfolioTarget(self.assets[i], self.trading_weight[i]))
            self.set_holdings(orders)
            
        self.state = 1
            
    elif normalized_spread > self.threshold:
        orders = []
        for i in range(len(self.assets)):
            orders.append(PortfolioTarget(self.assets[i], -1 * self.trading_weight[i]))
            self.set_holdings(orders)
            
        self.state = -1
            
    # Out of position if spread recovered
    elif self.state == 1 and normalized_spread > -self.threshold or self.state == -1 and normalized_spread < self.threshold:
        self.liquidate()
        
        self.state = 0