simple hmm model

simple_hmm.py

def simple_evaluate(data, states, hi=1.1, lo=0.9):
    skip = len(data) - len(states)
    chg = data['close'].pct_change()[skip:]
    n = max(states) + 1
    buys, sells = np.zeros(len(states)), np.zeros(len(states))
    for i in range(n):
        state = (states == i)
        dex = np.append(0, state[:-1])
        V = (chg.multiply(dex, axis=0)+1).cumprod()
        if V[-1] > hi:
            buys = buys + state
        elif V[-1] < lo:
            sells = sells + state
    
    buys, sells = np.append(0, buys[:-1]), np.append(0, sells[:-1])
    tot_return = (chg.multiply(buys, axis=0)+1).cumprod() - (chg.multiply(sells, axis=0)+1).cumprod()
    return tot_return

def fit_predict(df, lag=7, comps=5):
    sz = int(len(df) * 0.66)
    train, test = df.iloc[0:sz], df.iloc[sz:len(df)]

    chg = train.close.pct_change()[lag:]
    chg_lag = train.close.pct_change(lag)[lag:]

    X = np.column_stack([chg,chg_lag])
    hmm = GaussianHMM(n_components = comps, covariance_type='diag',n_iter = 5000).fit(X)
    states_train = hmm.predict(X)
    
    chg = test.close.pct_change()[lag:]
    chg_lag = test.close.pct_change(lag)[lag:]

    X = np.column_stack([chg,chg_lag])
    states_test = hmm.predict(X)
    return simple_evaluate(train, states_train), simple_evaluate(test, states_test)