generating data for lstm example

Data

def _load_data(data, n_prev = INIC):
    docX = []
    docY = []
    for i in range(len(data)-n_prev):
        docX.append(data.iloc[i:i+n_prev].as_matrix())
        docY.append(data.iloc[i+n_prev].as_matrix())
    alsX = np.array(docX)
    alsY = np.array(docY)
    
    return(alsX, alsY)

def train_test_split(df, test_size=0.1):
    ntrn = round(len(df) * (1 - test_size))
    
    X_train, y_train = _load_data(df.iloc[0:ntrn])
    X_test, y_test = _load_data(df.iloc[ntrn:])
    
    return((X_train, y_train), (X_test, y_test))

import pandas as pd
flow = (list(range(1,10,1)) + list(range(10,1,-1)))*10000
data = pd.DataFrame({"a":flow, "b":flow})
data.b = data.b.shift(9)
data = data.iloc[10:]

(X_train, y_train), (X_test, y_test)  = train_test_split(data)