Python stattools.coint方法代碼示例

# 需要導入模塊: from statsmodels.tsa import stattools [as 別名]
# 或者: from statsmodels.tsa.stattools import coint [as 別名]
def coint(df, intercept = True, sig_level = 0.01):
    """
    Find pairs (of 2 time series) that passes the cointegration test.

    Parameters
    ----------
    df: pandas dataframe
        each column is the time series of a certain stock
    
    intercept: boolean
        if True, OLS and ADF test are done manually
        if False, the coint() function from statsmodels.tsa.stattools, which 
        does not include intercept term while doing OLS regression, is used.
    
    sig_level: if p_value of cointegration test is below this level, then we
        can reject the NULL hypothesis, which says that the two series are not
        cointegrated
    
    Return
    ------
    A list of tuples of the form (name of stock 1, name of stock 2, p_value of
    cointegration test).
    
    """
    cointegrated_pairs = []
    
    stock_names = df.columns.values.tolist()
    N = len(stock_names)
    
    stock_pairs = list(itertools.combinations(stock_names, 2))
    
    for pair in stock_pairs:
        stock_1, stock_2 = pair

        p_value = 0

        if not intercept:
            p_value = smts.coint(df[stock_1].values.astype(float), df[stock_2].values.astype(float), trend='c')[1]
        else:
            Y = df[stock_1].values.astype(float)
            X = df[stock_2].values.astype(float)
            X = sm.add_constant(X)

            model = sm.OLS(Y, X)
            results = model.fit()
            intercept, slope = results.params

            p_value = smts.adfuller(results.resid)[1]

        if p_value < sig_level and slope > 0:
            cointegrated_pairs.append(tuple([stock_1, stock_2, p_value]))

    return cointegrated_pairs 

# 需要導入模塊: from statsmodels.tsa import stattools [as 別名]
# 或者: from statsmodels.tsa.stattools import coint [as 別名]
def get_p_value(x, y):
        _, p_val, _ = coint(x, y)
        return p_val 

# 需要導入模塊: from statsmodels.tsa import stattools [as 別名]
# 或者: from statsmodels.tsa.stattools import coint [as 別名]
def find_cointegrated_pairs(data):
    n = data.shape[1]
    pvalue_matrix = np.ones((n, n))
    keys = data.keys()
    pairs = []
    for i in range(n):
        for j in range(i+1, n):
            result = coint(data[keys[i]], data[keys[j]])
            pvalue_matrix[i, j] = result[1]
            if result[1] < 0.02:
                pairs.append((keys[i], keys[j]))
    return pvalue_matrix, pairs