Keras simple neural network with template structures (v) - RNN LSTM Regressor Recurrent Neural Networks

# -*- coding: utf-8 -*-
import numpy as np
np.random.seed(1337)
import matplotlib.pyplot as plt
from keras.models import Sequential
from keras.layers import LSTM,TimeDistributed,Dense
from keras.optimizers import Adam

BATCH_START = 0 
TIME_STEPS = 20 
BATCH_SIZE = 50
INPUT_SIZE = 1 
OUTPUT_SIZE = 1
CELL_SIZE = 20
LR = 0.006

def get_batch():
    global BATCH_START,TIME_STEPS
    # xs shape(50,20,)
    #xs=np.arange(0,0+20*50).reshape(50,20)
    xs = np.arange(BATCH_START,BATCH_START+TIME_STEPS*BATCH_SIZE).reshape((BATCH_SIZE,TIME_STEPS)) / (10*np.pi)
    seq = np.sin(xs)
    res = np.cos(xs)
    BATCH_START += TIME_STEPS
    #plt.plot(xs[0,:],res[0,:],'r',xs[0,:],seq[0,:],'b--')
    #plt.show()
    return [seq[:,:,np.newaxis], res[:,:,np.newaxis],xs]

#get_batch()
#exit()
    
    
model = Sequential()

model.add(LSTM(output_dim= Cell_size, 
               return_sequences = True, # each time the output of a point the Output 
               batch_input_shape = (BATCH_SIZE, TIME_STEPS, INPUT_SIZE), 
               Stateful = True, # Is there a link between batch and batch 
               # before the final step in a batch after batch, and a the first step is linked 
        )) 

model.add (TimeDistributed (Dense (OUTPUT_SIZE))) # Dense each output connection, be calculated for each time point 

ADAM = Adam (LR) 
model.compile (Optimizer = ADAM, 
              Loss = ' MSE ' ,) 

Print ( 'Training ------------')
for step in range(501):
    # data shape = (batch_num,steps,inputs/output)
    X_batch, Y_batch, xs = get_batch()
    cost = model.train_on_batch(X_batch, Y_batch)
    pred = model.predict(X_batch,BATCH_SIZE)
    plt.plot(xs[0,:], Y_batch[0].flatten(),'r',xs[0,:],pred.flatten()[:TIME_STEPS],'b--')
    plt.ylim((-1.2,1.2))
    plt.draw()
    plt.pause(0.5)
    if step % 10 == 0:
        print('train cost',cost)