def main():
    # https://keras.io/layers/recurrent/
    num_samples, time_steps, input_dim, output_dim = 50_000, 10, 1, 1

    # Definition of the model.
    model_keras = NBeatsKeras(backcast_length=time_steps, forecast_length=output_dim,
                              stack_types=(NBeatsKeras.GENERIC_BLOCK, NBeatsKeras.GENERIC_BLOCK),
                              nb_blocks_per_stack=2, thetas_dim=(4, 4), share_weights_in_stack=True,
                              hidden_layer_units=64)

    model_pytorch = NBeatsPytorch(backcast_length=time_steps, forecast_length=output_dim,
                                  stack_types=(NBeatsPytorch.GENERIC_BLOCK, NBeatsPytorch.GENERIC_BLOCK),
                                  nb_blocks_per_stack=2, thetas_dim=(4, 4), share_weights_in_stack=True,
                                  hidden_layer_units=64)

    # Definition of the objective function and the optimizer.
    model_keras.compile(loss='mae', optimizer='adam')
    model_pytorch.compile(loss='mae', optimizer='adam')

    # Definition of the data. The problem to solve is to find f such as | f(x) - y | -> 0.
    # where f = np.mean.
    x = np.random.uniform(size=(num_samples, time_steps, input_dim))
    y = np.mean(x, axis=1, keepdims=True)

    # Split data into training and testing datasets.
    c = num_samples // 10
    x_train, y_train, x_test, y_test = x[c:], y[c:], x[:c], y[:c]
    test_size = len(x_test)

    # Train the model.
    print('Keras training...')
    model_keras.fit(x_train, y_train, validation_data=(x_test, y_test), epochs=20, batch_size=128)
    print('Pytorch training...')
    model_pytorch.fit(x_train, y_train, validation_data=(x_test, y_test), epochs=20, batch_size=128)

    # Save the model for later.
    model_keras.save('n_beats_model.h5')
    model_pytorch.save('n_beats_pytorch.th')

    # Predict on the testing set (forecast).
    predictions_keras_forecast = model_keras.predict(x_test)
    predictions_pytorch_forecast = model_pytorch.predict(x_test)
    np.testing.assert_equal(predictions_keras_forecast.shape, (test_size, model_keras.forecast_length, output_dim))
    np.testing.assert_equal(predictions_pytorch_forecast.shape, (test_size, model_pytorch.forecast_length, output_dim))

    # Predict on the testing set (backcast).
    predictions_keras_backcast = model_keras.predict(x_test, return_backcast=True)
    predictions_pytorch_backcast = model_pytorch.predict(x_test, return_backcast=True)
    np.testing.assert_equal(predictions_keras_backcast.shape, (test_size, model_keras.backcast_length, output_dim))
    np.testing.assert_equal(predictions_pytorch_backcast.shape, (test_size, model_pytorch.backcast_length, output_dim))

    # Load the model.
    model_keras_2 = NBeatsKeras.load('n_beats_model.h5')
    model_pytorch_2 = NBeatsPytorch.load('n_beats_pytorch.th')

    np.testing.assert_almost_equal(predictions_keras_forecast, model_keras_2.predict(x_test))
    np.testing.assert_almost_equal(predictions_pytorch_forecast, model_pytorch_2.predict(x_test))
Exemple #2
0
def main():
    args = get_script_arguments()
    device = torch.device(
        'cuda') if not args.disable_cuda and torch.cuda.is_available(
        ) else torch.device('cpu')
    forecast_length = 10
    backcast_length = 5 * forecast_length
    batch_size = 4  # greater than 4 for viz

    if args.task == 'm4':
        data_gen = batcher(get_m4_data(backcast_length, forecast_length),
                           batch_size=batch_size,
                           infinite=True)
    elif args.task == 'dummy':
        data_gen = dummy_data_generator(backcast_length,
                                        forecast_length,
                                        signal_type='seasonality',
                                        random=True,
                                        batch_size=batch_size)
    else:
        raise Exception('Unknown task.')

    print('--- Model ---')
    net = NBeatsNet(
        device=device,
        stack_types=[NBeatsNet.TREND_BLOCK, NBeatsNet.SEASONALITY_BLOCK],
        forecast_length=forecast_length,
        thetas_dims=[2, 8],
        nb_blocks_per_stack=3,
        backcast_length=backcast_length,
        hidden_layer_units=1024,
        share_weights_in_stack=False,
        nb_harmonics=None)

    # net = NBeatsNet(device=device,
    #                 stack_types=[NBeatsNet.GENERIC_BLOCK, NBeatsNet.GENERIC_BLOCK],
    #                 forecast_length=forecast_length,
    #                 thetas_dims=[7, 8],
    #                 nb_blocks_per_stack=3,
    #                 backcast_length=backcast_length,
    #                 hidden_layer_units=128,
    #                 share_weights_in_stack=False
    #                 )

    optimiser = optim.Adam(net.parameters())

    def plot_model(x, target, grad_step):
        if not args.disable_plot:
            print('plot()')
            plot(net, x, target, backcast_length, forecast_length, grad_step)

    max_grad_steps = 10000
    if args.test:
        max_grad_steps = 5

    simple_fit(net, optimiser, data_gen, plot_model, device, max_grad_steps)
Exemple #3
0
    def train(self, **kwargs):
        """Train NBEATS Model"""

        if os.path.isfile(CHECKPOINT_NAME):
            os.remove(CHECKPOINT_NAME)

        steps = kwargs.get("steps", 50)
        batch_size = kwargs.get("batch_size", 10)
        patience = kwargs.get("patience", 5)
        device = self.get_device()

        # There seems to be an issue with
        net = NBeatsNet(
            stack_types=[
                NBeatsNet.TREND_BLOCK,
                NBeatsNet.SEASONALITY_BLOCK,
                NBeatsNet.GENERIC_BLOCK,
            ],
            forecast_length=self.forecast_len,
            thetas_dims=kwargs.get("thetas_dims", [2, 8, 3]),
            nb_blocks_per_stack=kwargs.get("nb_blocks_per_stack", 3),
            backcast_length=self.forecast_len,
            hidden_layer_units=kwargs.get("hidden_layer_units", 128),
            share_weights_in_stack=False,
            device=device,
        )

        x_batch, y_batch, norm_constant = self.format_input(
            self.dataframe, self.forecast_len)

        c = len(x_batch)
        if self.in_sample:
            c -= self.forecast_len

        optimiser = optim.Adam(net.parameters())

        data = data_generator(x_batch[:c], y_batch[:c], batch_size)

        best_loss = float("inf")
        counter = 0

        for _ in tqdm(range(steps)):
            loss = self.train_100_grad_steps(data, device, net, optimiser)
            if loss < best_loss:
                best_loss, counter = loss, 0
            else:
                counter += 1
                if counter >= patience:
                    break

        self.model = net
        self.constant = norm_constant
    def __init__(self,period_to_forecast,data=None, backcast_length=None,save_checkpoint=False,path='',checkpoint_file_name='nbeats-training-checkpoint.th',mode='cpu',batch_size=None,thetas_dims=[7, 8],nb_blocks_per_stack=3,share_weights_in_stack=False,train_percent=0.8,hidden_layer_units=128,stack=None):
        if (data is None):
            print('For Prediction as no data passed')
            batch_size=np.nan
        else:
            self.data=data
            
            if(len(self.data.shape)!=2):
                raise Exception('Numpy array should be of nx1 shape')
            if self.data.shape[1]!=1:
                raise Exception('Numpy array should be of nx1 shape')
        self.forecast_length=period_to_forecast
        if backcast_length==None:
            self.backcast_length = 3 * self.forecast_length
        else:
            self.backcast_length=backcast_length
        if batch_size==None:
            self.batch_size=int(self.data.shape[0]/15)
        else:
            self.batch_size=batch_size
        self.hidden_layer_units=hidden_layer_units  
        
        if stack==None:
            self.stacks= [NBeatsNet.GENERIC_BLOCK, NBeatsNet.GENERIC_BLOCK]
        else:
            Dict = {1: NBeatsNet.GENERIC_BLOCK, 2: NBeatsNet.TREND_BLOCK , 3: NBeatsNet.SEASONALITY_BLOCK} 
            self.stacks = [Dict.get(item, item) for item in stack]

        

        self.CHECKPOINT_NAME= path+checkpoint_file_name
        self.device = torch.device(mode)       #change to gpu if gpu present for better performance
        self.thetas_dims=thetas_dims
        self.nb_blocks_per_stack=nb_blocks_per_stack
        self.train_size=train_percent
        self.share_weights_in_stack=share_weights_in_stack

        self.net = NBeatsNet(stack_types=self.stacks,
                        forecast_length=self.forecast_length,
                        thetas_dims=self.thetas_dims,
                        nb_blocks_per_stack=self.nb_blocks_per_stack,
                        backcast_length=self.backcast_length,
                        hidden_layer_units=self.hidden_layer_units,
                        share_weights_in_stack=self.share_weights_in_stack,
                        device=self.device)
        self.parameters=self.net.parameters()
        self.global_step_cl=0
        self.check_save=save_checkpoint
        self.loaded=False
        self.saved=True
Exemple #5
0
def main():
    
    num_samples, time_steps, input_dim, output_dim = 50000, 10, 1, 1

    # Definition of the model.

    model_pytorch = NBeatsPytorch(backcast_length=time_steps, forecast_length=output_dim,
                                  stack_types=(NBeatsPytorch.GENERIC_BLOCK, NBeatsPytorch.GENERIC_BLOCK),
                                  nb_blocks_per_stack=2, thetas_dim=(4, 4), share_weights_in_stack=True,
                                  hidden_layer_units=64)

    # Definition of the objective function and the optimizer.
    model_pytorch.compile_model(loss='mae', learning_rate=1e-4)

    # Definition of the data. The problem to solve is to find f such as | f(x) - y | -> 0.
    # where f = np.mean.
    x = np.random.uniform(size=(num_samples, time_steps, input_dim))
    y = np.mean(x, axis=1, keepdims=True)

    # Split data into training and testing datasets.
    c = num_samples // 10
    x_train, y_train, x_test, y_test = x[c:], y[c:], x[:c], y[:c]

    # Train the model.
    print('Pytorch training...')
    model_pytorch.fit(x_train, y_train, validation_data=(x_test, y_test), epochs=20, batch_size=128)
    
    # Save the model for later.
    model_pytorch.save('n_beats_pytorch.th')

    # Predict on the testing set.
    
    predictions_pytorch = model_pytorch.predict(x_test)
    
    print(predictions_pytorch.shape)

    # Load the model.
    
    model_pytorch_2 = NBeatsPytorch.load('n_beats_pytorch.th')

    
    np.testing.assert_almost_equal(predictions_pytorch, model_pytorch_2.predict(x_test))
Exemple #6
0
    def load(self, path=""):
        for d in self.classes:
            checkpoint = d + "_" + self.checkpoint_name_BASE + f'bl{self.nb_blocks_per_stack}-f{self.forecast_length}-b{self.backcast_length}-btch{self.batch_size}-h{self.hidden_layer_units}.th'
            if path:
                checkpoint = path + checkpoint

            print(
                f'From where to load models: {path}, and the current checkpoint is: {checkpoint}'
            )
            if not os.path.exists(checkpoint):
                print("-" * 50)
                print(
                    "* /n * /n * /n * /n * /n * /t THERE IS NO CHECKPOINT LIKE THIS! * /n * /n * /n * /n * "
                )
                print("-" * 50)
                raise ValueError("wrong checkpoint name!")
            else:
                print("-" * 50)
                print(
                    "*/n*/n*/n*/n*/n*/t CHECKPOINT IS VALID WOOOHOO ! */n*/n*/n*/n*"
                )
                print("-" * 50)

            net = NBeatsNet(
                stack_types=[NBeatsNet.GENERIC_BLOCK, NBeatsNet.GENERIC_BLOCK],
                forecast_length=self.forecast_length,
                thetas_dims=self.thetas,
                nb_blocks_per_stack=self.nb_blocks_per_stack,
                backcast_length=self.backcast_length,
                hidden_layer_units=self.hidden_layer_units,
                share_weights_in_stack=False,
                device=self.device)
            optimiser = optim.Adam(net.parameters())

            naf.load(checkpoint, net, optimiser)
            net.to(self.device)
            self.nets[d] = net
Exemple #7
0
def nbeats_dataframe(df,
                     forecast_length,
                     in_sample,
                     device,
                     train_portion=0.75):

    backcast_length = 1 * forecast_length

    df = df.values  # just keep np array here for simplicity.
    norm_constant = np.max(df)
    df = df / norm_constant  # small leak to the test set here.

    x_train_batch, y = [], []

    for i in range(backcast_length + 1,
                   len(df) - forecast_length +
                   1):  # 25% to 75% so 50% #### Watch out I had to plus one.
        x_train_batch.append(df[i - backcast_length:i])
        y.append(df[i:i + forecast_length])

    x_train_batch = np.array(x_train_batch)[..., 0]
    y = np.array(y)[..., 0]

    # x_train, y_train = x_train_batch, y

    net = NBeatsNet(
        stack_types=[NBeatsNet.GENERIC_BLOCK, NBeatsNet.GENERIC_BLOCK],
        forecast_length=forecast_length,
        thetas_dims=[7, 8],
        nb_blocks_per_stack=3,
        backcast_length=backcast_length,
        hidden_layer_units=128,
        share_weights_in_stack=False,
        device=device,
    )

    if in_sample == True:
        c = int(len(x_train_batch) * train_portion)
        x_train, y_train = x_train_batch[:c], y[:c]
        x_test, y_test = x_train_batch[c:], y[c:]

        return x_train, y_train, x_test, y_test, net, norm_constant

    else:
        c = int(len(x_train_batch) * 1)
        x_train, y_train = x_train_batch[:c], y[:c]

        return x_train, y_train, net, norm_constant
Exemple #8
0
forecast_length = 500
backcast_length = 3 * forecast_length
batch_size = 256

for folder_name in dirs:
    name = folder_name.split("/")[-1]
    checkpoint_name = name + "_" + checkpoint_name_BASE

    if os.path.isfile(checkpoint_name):
        continue

    net = NBeatsNet(
        stack_types=[NBeatsNet.GENERIC_BLOCK, NBeatsNet.GENERIC_BLOCK],
        forecast_length=forecast_length,
        thetas_dims=[7, 8],
        nb_blocks_per_stack=3,
        backcast_length=backcast_length,
        hidden_layer_units=128,
        share_weights_in_stack=False,
        device=device)
    optimiser = optim.Adam(net.parameters())

    test_losses = []
    actual_class_dir = data_dir + "/" + name + "/"
    for (_, dirs, files) in os.walk(actual_class_dir):
        iteration = 0
        for file in files:

            if 'mat' in file:
                continue
class NBeats:   #UNIVARIATE DATA TO BE PASSED AS NUMPY ARRAY
    def __init__(self,period_to_forecast,data=None, backcast_length=None,save_checkpoint=False,path='',checkpoint_file_name='nbeats-training-checkpoint.th',mode='cpu',batch_size=None,thetas_dims=[7, 8],nb_blocks_per_stack=3,share_weights_in_stack=False,train_percent=0.8,hidden_layer_units=128,stack=None):
        if (data is None):
            print('For Prediction as no data passed')
            batch_size=np.nan
        else:
            self.data=data
            
            if(len(self.data.shape)!=2):
                raise Exception('Numpy array should be of nx1 shape')
            if self.data.shape[1]!=1:
                raise Exception('Numpy array should be of nx1 shape')
        self.forecast_length=period_to_forecast
        if backcast_length==None:
            self.backcast_length = 3 * self.forecast_length
        else:
            self.backcast_length=backcast_length
        if batch_size==None:
            self.batch_size=int(self.data.shape[0]/15)
        else:
            self.batch_size=batch_size
        self.hidden_layer_units=hidden_layer_units  
        
        if stack==None:
            self.stacks= [NBeatsNet.GENERIC_BLOCK, NBeatsNet.GENERIC_BLOCK]
        else:
            Dict = {1: NBeatsNet.GENERIC_BLOCK, 2: NBeatsNet.TREND_BLOCK , 3: NBeatsNet.SEASONALITY_BLOCK} 
            self.stacks = [Dict.get(item, item) for item in stack]

        

        self.CHECKPOINT_NAME= path+checkpoint_file_name
        self.device = torch.device(mode)       #change to gpu if gpu present for better performance
        self.thetas_dims=thetas_dims
        self.nb_blocks_per_stack=nb_blocks_per_stack
        self.train_size=train_percent
        self.share_weights_in_stack=share_weights_in_stack

        self.net = NBeatsNet(stack_types=self.stacks,
                        forecast_length=self.forecast_length,
                        thetas_dims=self.thetas_dims,
                        nb_blocks_per_stack=self.nb_blocks_per_stack,
                        backcast_length=self.backcast_length,
                        hidden_layer_units=self.hidden_layer_units,
                        share_weights_in_stack=self.share_weights_in_stack,
                        device=self.device)
        self.parameters=self.net.parameters()
        self.global_step_cl=0
        self.check_save=save_checkpoint
        self.loaded=False
        self.saved=True
        
        
        
    def plot_scatter(self,*args, **kwargs):
        plt.plot(*args, **kwargs)
        plt.scatter(*args, **kwargs)
    
    
    def data_generator(self,x_full, y_full, bs):
        def split(arr, size):
            arrays = []
            while len(arr) > size:
                slice_ = arr[:size]
                arrays.append(slice_)
                arr = arr[size:]
            arrays.append(arr)
            return arrays
    
        while True:
            for rr in split((x_full, y_full), bs):
                yield rr    
                
    def loader(self,model, optimiser):
        if os.path.exists(self.CHECKPOINT_NAME):
            checkpoint = torch.load(self.CHECKPOINT_NAME)
            
            model.load_state_dict(checkpoint['model_state_dict'])
            optimiser.load_state_dict(checkpoint['optimizer_state_dict'])
            grad_step = checkpoint['grad_step']
            if self.loaded:
                self.norm_constant=checkpoint['norm_constant']
            return grad_step
        return 0
    
    def saver(self,model, optimiser, grad_step):
        if self.saved:
                 torch.save({
                'norm_constant':self.norm_constant, 
                'grad_step': grad_step,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimiser.state_dict(),
            }, self.CHECKPOINT_NAME) 
        else:            
            torch.save({
                'grad_step': grad_step,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimiser.state_dict(),
            }, self.CHECKPOINT_NAME) 
    
    
    def load(self,file=None,optimiser=None):
         if file==None:
            raise Exception('Empty File Name')
         elif file=='':
            raise Exception('Empty File Name') 
         else:
           if optimiser==None:
            self.optimiser = optim.Adam(self.net.parameters())
           else:
               self.optimiser = optimiser
           self.CHECKPOINT_NAME=file
           self.loaded=True
           self.global_step_cl=self.loader(self.net,self.optimiser)
           self.CHECKPOINT_NAME='nbeats-training-checkpoint.th'
        
            
        
    def save(self,file=None):
        if file==None:
            raise Exception('Empty File Name')
        elif file=='':
            raise Exception('Empty File Name')            
        else:
            self.CHECKPOINT_NAME=file
            self.saver(self.net,self.optimiser,self.global_step_cl)
            self.saved=True

           
    def train_100_grad_steps(self,data, test_losses):
        if not self.loaded:
            global_step = self.loader(self.net, self.optimiser)
            self.loaded=False
            self.global_step_cl=global_step
        for x_train_batch, y_train_batch in data:
            self.global_step_cl += 1
            self.optimiser.zero_grad()
            self.net.train()
            _, forecast = self.net(torch.tensor(x_train_batch, dtype=torch.float).to(self.device))
            loss = F.mse_loss(forecast, torch.tensor(y_train_batch, dtype=torch.float).to(self.device))
            loss.backward()
            self.optimiser.step()
            if self.verbose:
                if self.global_step_cl % 30 == 0:
                    print(f'grad_step = {str(self.global_step_cl).zfill(6)}, tr_loss = {loss.item():.6f}, te_loss = {test_losses[-1]:.6f}')
            if self.global_step_cl > 0 and self.global_step_cl % 100 == 0:
                with torch.no_grad():
                   self.saver(self.net, self.optimiser, self.global_step_cl)
                break
        return forecast   
    
    def eval_test(self, test_losses, x_test, y_test):
        self.net.eval()
        _, forecast = self.net(torch.tensor(x_test, dtype=torch.float))
        test_losses.append(F.mse_loss(forecast, torch.tensor(y_test, dtype=torch.float)).item())
        p = forecast.detach().numpy()
        if self.plot:
            subplots = [221, 222, 223, 224]
            plt.figure(1)
            for plot_id, i in enumerate(np.random.choice(range(len(x_test)), size=4, replace=False)):
                ff, xx, yy = p[i] * self.norm_constant, x_test[i] * self.norm_constant, y_test[i] * self.norm_constant
                plt.subplot(subplots[plot_id])
                plt.grid()
                self.plot_scatter(range(0, self.backcast_length), xx, color='b')
                self.plot_scatter(range(self.backcast_length, self.backcast_length + self.forecast_length), yy, color='g')
                self.plot_scatter(range(self.backcast_length, self.backcast_length + self.forecast_length), ff, color='r')
            plt.show()
        
        
    def fit(self,epoch=25,optimiser=None,plot=False,verbose=True):
        self.plot=plot
        self.verbose=verbose

        self.epoch=epoch
        self.norm_constant = np.max(self.data)
        self.data = self.data / self.norm_constant 
        x_train_batch, y,self.x_forecast = [], [] ,[]
        for i in range(self.backcast_length, len(self.data) - self.forecast_length):
            x_train_batch.append(self.data[i - self.backcast_length:i])
            y.append(self.data[i:i + self.forecast_length])
        i+=self.forecast_length
        self.x_forecast.append(self.data[i+1 - self.backcast_length:i+1])
        x_train_batch = np.array(x_train_batch)[..., 0]
        self.x_forecast = np.array(self.x_forecast)[..., 0]

        y = np.array(y)[..., 0]
        i+=self.forecast_length

        if optimiser==None:
            self.optimiser = optim.Adam(self.net.parameters())
        else:
            self.optimiser = optimiser
        
        c = int(len(x_train_batch) * self.train_size)
        x_train, y_train = x_train_batch[:c], y[:c]
        x_test, y_test = x_train_batch[c:], y[c:]

        train_data = self.data_generator(x_train, y_train, self.batch_size)
        test_losses = []
        for i in range(self.epoch):
            self.eval_test(test_losses, x_test, y_test)
            self.train_100_grad_steps(train_data, test_losses)
        if self.check_save:
            pass
        else:
            if os.path.exists(self.CHECKPOINT_NAME):
                os.remove(self.CHECKPOINT_NAME)
            
    def predict(self,predict_data=None):        
        if (predict_data is None):
            _, forecasted_values = self.net(torch.tensor(self.x_forecast, dtype=torch.float))
            forecasted_values= forecasted_values.detach().numpy()
            forecasted_values = forecasted_values * self.norm_constant
        else:
            if (predict_data.shape[0]!=self.backcast_length):
                raise Exception('Numpy array for prediction input should be of backcast_length: {} x 1 shape'.format(self.backcast_length))
            else:
                predict_data=predict_data/self.norm_constant
                predict_data= np.reshape(predict_data, (self.backcast_length, 1)) 
                
                _, forecasted_values = self.net(torch.tensor(predict_data.T, dtype=torch.float))
                forecasted_values= forecasted_values.detach().numpy()
                forecasted_values = forecasted_values * self.norm_constant
                
        return forecasted_values.T
Exemple #10
0
def train_and_score_model(state,
                          ili_data,
                          horizon=12,
                          lookback=120,
                          split=0.7,
                          batch_size=10,
                          model_dir='saved_models',
                          covid_19_onset='2019-11-01',
                          device=torch.device('cpu'),
                          report_interval=100,
                          num_training_intervals=20):
    '''traing a state flu model'''

    # check state data exists
    supported_states = ili_data.STATE.unique().tolist()
    assert state in supported_states, f'{state} not present in the ili_data'

    # establish checkpoint
    os.makedirs(model_dir, exist_ok=True)
    checkpoint_path = f'{model_dir}/nbeats-training-checkpoint.th'

    if os.path.isfile(checkpoint_path):
        os.remove(checkpoint_path)

    # map inputs to commonly used variables
    forecast_length = horizon  # num weeks in forecast horizon
    backcast_length = lookback
    norm_constant = 1

    # get state data and eliminate COVID-19 period
    state_data_df = ili_data.loc['2000-01-01':covid_19_onset].reset_index()
    state_data_df.set_index(['STATE', 'DATE'], inplace=True)
    state_data = state_data_df.loc[state].values

    # create n-beats training data
    x_train_batch, ylist = [], []
    for i in range(backcast_length, len(state_data) - forecast_length):
        x_train_batch.append(state_data[i - backcast_length:i])
        ylist.append(state_data[i:i + forecast_length])

    x_train_batch = np.array(x_train_batch)[..., 0]
    ylist = np.array(ylist)[..., 0]

    cut = int(len(x_train_batch) * split)
    x_train, y_train = x_train_batch[:cut], ylist[:cut]
    x_test, y_test = x_train_batch[cut:], ylist[cut:]

    # create n-beats model
    net = NBeatsNet(
        stack_types=[NBeatsNet.GENERIC_BLOCK, NBeatsNet.GENERIC_BLOCK],
        forecast_length=forecast_length,
        thetas_dims=[7, 8],
        nb_blocks_per_stack=3,
        backcast_length=backcast_length,
        hidden_layer_units=512,  #128,
        share_weights_in_stack=False,
        device=device)

    # set optimizer
    optimiser = optim.Adam(net.parameters())

    # create data generator
    data = data_generator(x_train, y_train, batch_size)

    # train model
    test_losses = []
    training_losses = []

    for i in range(num_training_intervals):
        eval_test(backcast_length, forecast_length, net, norm_constant,
                  test_losses, x_test, y_test)

        training_losses, test_losses = train_100_grad_steps(
            data, device, net, optimiser, training_losses, test_losses,
            report_interval, checkpoint_path)

    _, forecast = net(torch.tensor(x_test, dtype=torch.float))
    y_pred = forecast.detach().numpy()

    return y_test, y_pred, checkpoint_path
def main(name, copy, device='cpu'):
    xtrain = np.loadtxt('xtrain_{}.txt'.format(name))
    ytrain = np.loadtxt('ytrain_{}.txt'.format(name))
    xtest = np.loadtxt('xtest_{}.txt'.format(name))
    ytest = np.loadtxt('ytest_{}.txt'.format(name))
    print('-------we got the DATA dude : it works ------------')

    #Reminder of the hyperparameter
    backcast_length = 100
    forecast_length = 100

    #Definition of the seasonality  model :
    thetas_dim1 = 4,
    stack_types1 = NBeatsNet.SEASONALITY_BLOCK,
    model1 = NBeatsNet(device=torch.device(device),
                       backcast_length=backcast_length,
                       forecast_length=forecast_length,
                       stack_types=stack_types1,
                       nb_blocks_per_stack=2,
                       thetas_dim=thetas_dim1,
                       share_weights_in_stack=True,
                       hidden_layer_units=64)

    model1.compile_model(loss='mae', learning_rate=1e-5)
    plt.figure(figsize=(10, 10))
    #model training
    model1.fit(xtrain,
               ytrain,
               validation_data=(xtest, ytest),
               epochs=100,
               batch_size=10)

    model1.save('nbeats_test_seasonality.h5')

    predictions = model1.predict(xtest)
    print(predictions.shape)

    #Definition of the generic  model :
    thetas_dim2 = 4,
    stack_types2 = NBeatsNet.GENERIC_BLOCK,
    model2 = NBeatsNet(device=torch.device(device),
                       backcast_length=backcast_length,
                       forecast_length=forecast_length,
                       stack_types=stack_types2,
                       nb_blocks_per_stack=2,
                       thetas_dim=thetas_dim2,
                       share_weights_in_stack=True,
                       hidden_layer_units=64)

    model2.compile_model(loss='mae', learning_rate=1e-5)

    #model training
    model2.fit(xtrain,
               ytrain,
               validation_data=(xtest, ytest),
               epochs=100,
               batch_size=10)

    model2.save('nbeats_test_seasonality.h5')

    predictions = model2.predict(xtest)
    print(predictions.shape)

    #Definition of the trend model :
    thetas_dim3 = 4,
    stack_types3 = NBeatsNet.TREND_BLOCK,
    model3 = NBeatsNet(device=torch.device(device),
                       backcast_length=backcast_length,
                       forecast_length=forecast_length,
                       stack_types=stack_types3,
                       nb_blocks_per_stack=2,
                       thetas_dim=thetas_dim3,
                       share_weights_in_stack=True,
                       hidden_layer_units=64)

    model3.compile_model(loss='mae', learning_rate=1e-5)

    #model training

    model3.fit(xtrain,
               ytrain,
               validation_data=(xtest, ytest),
               epochs=100,
               batch_size=10)

    model3.save('nbeats_test_seasonality.h5')

    predictions = model3.predict(xtest)
    print(predictions.shape)
    plt.show()