text: # nohup python -u pred_STTransformer.py > pred_STTransformer.log 2>&1 &
  import sys
  import os
  import shutil
  import math
  import numpy as np
  import pandas as pd
  import scipy.sparse as ss
  from sklearn.preprocessing import StandardScaler, MinMaxScaler
  from datetime import datetime
  import time
  import torch
  import torch.nn as nn
  import torch.nn.init as init
  import torch.nn.functional as F
  from torchsummary import summary
  from model.Double_C_STTN import *
  import Metrics
  
  from Param import *
  import Utils
  
  import os
  
  cpu_num = 1
  os.environ['OMP_NUM_THREADS'] = str(cpu_num)
  os.environ['OPENBLAS_NUM_THREADS'] = str(cpu_num)
  os.environ['MKL_NUM_THREADS'] = str(cpu_num)
  os.environ['VECLIB_MAXIMUM_THREADS'] = str(cpu_num)
  os.environ['NUMEXPR_NUM_THREADS'] = str(cpu_num)
  torch.set_num_threads(cpu_num)
  
  DEVICE = 'cuda:1'
  BATCHSIZE = 32
  EPOCH = 200
  
  
  def getXSYS(data, mode):
      TRAIN_NUM = int(data.shape[0] * TRAINRATIO)
      XS, YS = [], []
      if mode == 'TRAIN':
          for i in range(TRAIN_NUM - TIMESTEP_OUT - TIMESTEP_IN + 1):
              x = data[i:i + TIMESTEP_IN, :]
              y = data[i + TIMESTEP_IN:i + TIMESTEP_IN + TIMESTEP_OUT, :]
              XS.append(x), YS.append(y)
      elif mode == 'TEST':
          for i in range(TRAIN_NUM - TIMESTEP_IN, data.shape[0] - TIMESTEP_OUT - TIMESTEP_IN + 1):
              x = data[i:i + TIMESTEP_IN, :]
              y = data[i + TIMESTEP_IN:i + TIMESTEP_IN + TIMESTEP_OUT, :]
              XS.append(x), YS.append(y)
      XS, YS = np.array(XS), np.array(YS)
      # XS, YS = XS[:, :, :, np.newaxis], YS[:, :, :, np.newaxis]
      XS = XS.transpose(0, 3, 2, 1)
      YS = YS.transpose(0, 3, 1, 2)
      YS = YS[:,0:1,:,:]
  
      return XS, YS
  
  
  def getModel(name):
      #     ks, kt, bs, T, n, p = 3, 3, [[CHANNEL, 16, 64], [64, 16, 64]], TIMESTEP_IN, N_NODE, 0
      #     A = pd.read_csv(ADJPATH).values
      #     W = weight_matrix(A)
      #     L = scaled_laplacian(W)
      #     Lk = cheb_poly(L, ks)
      #     Lk = torch.Tensor(Lk.astype(np.float32)).to(device)
      #     model = STGCN(ks, kt, bs, T, n, Lk, p).to(device)
      #     return model
  
      ### Adjacency Matrix Import
      #     adj_mx = pd.read_csv('../PEMSD7M/W_228.csv')#.iloc[:,1:]
      #     adj_mx = np.array(adj_mx)
      #     A = adj_mx
      #     A = torch.Tensor(A)
  
      A = pd.read_csv(ADJPATH).values
      A = torch.Tensor(A)
  
      ### Training Hyparameter
      in_channels = 2  # Channels of input
      embed_size = 32  # Dimension of hidden embedding features
      time_num = 288
      num_layers = 2  # Number of ST Block
      T_dim = 12  # Input length, should be the same as prepareData.py
      output_T_dim = 12  # Output Expected length
      heads = 4  # Number of Heads in MultiHeadAttention
      cheb_K = 2  # Order for Chebyshev Polynomials (Eq 2)
      forward_expansion = 4  # Dimension of Feed Forward Network: embed_size --> embed_size * forward_expansion --> embed_size
      dropout = 0
  
      ### Construct Network
      model = STTransformer(
          A,
          in_channels,
          embed_size,
          time_num,
          num_layers,
          T_dim,
          output_T_dim,
          heads,
          cheb_K,
          forward_expansion,
          dropout).to(DEVICE)
      return model
  
  
  def evaluateModel(model, criterion, data_iter):
      model.eval()
      l_sum, n = 0.0, 0
      with torch.no_grad():
          for x, y in data_iter:
              y_pred = model(x)
              l = criterion(y_pred, y)
              l_sum += l.item() * y.shape[0]
              n += y.shape[0]
          return l_sum / n
  
  
  def predictModel(model, data_iter):
      YS_pred = []
      model.eval()
      with torch.no_grad():
          for x, y in data_iter:
              YS_pred_batch = model(x)
              YS_pred_batch = YS_pred_batch.cpu().numpy()
              YS_pred.append(YS_pred_batch)
          YS_pred = np.vstack(YS_pred)
      return YS_pred
  
  
  def trainModel(name, mode, XS, YS):
      print('Model Training Started ...', time.ctime())
      print('TIMESTEP_IN, TIMESTEP_OUT', TIMESTEP_IN, TIMESTEP_OUT)
      model = getModel(name)
      summary(model, (2, N_NODE, TIMESTEP_IN), device=device)
      XS_torch, YS_torch = torch.Tensor(XS).to(device), torch.Tensor(YS).to(device)
      trainval_data = torch.utils.data.TensorDataset(XS_torch, YS_torch)
      trainval_size = len(trainval_data)
      train_size = int(trainval_size * (1 - TRAINVALSPLIT))
      print('XS_torch.shape:  ', XS_torch.shape)
      print('YS_torch.shape:  ', YS_torch.shape)
      train_data = torch.utils.data.Subset(trainval_data, list(range(0, train_size)))
      val_data = torch.utils.data.Subset(trainval_data, list(range(train_size, trainval_size)))
      train_iter = torch.utils.data.DataLoader(train_data, BATCHSIZE, shuffle=True)
      val_iter = torch.utils.data.DataLoader(val_data, BATCHSIZE, shuffle=True)
  
      min_val_loss = np.inf
      wait = 0
  
      print('LOSS is :', LOSS)
      if LOSS == "MaskMAE":
          criterion = Utils.masked_mae
      if LOSS == 'MSE':
          criterion = nn.MSELoss()
      if LOSS == 'MAE':
          criterion = nn.L1Loss()
      if OPTIMIZER == 'RMSprop':
          optimizer = torch.optim.RMSprop(model.parameters(), lr=LEARN)
      if OPTIMIZER == 'Adam':
          optimizer = torch.optim.Adam(model.parameters(), lr=LEARN)
  
      for epoch in range(EPOCH):
          starttime = datetime.now()
          loss_sum, n = 0.0, 0
          model.train()
          for x, y in train_iter:
              optimizer.zero_grad()
              y_pred = model(x)
              loss = criterion(y_pred, y)
              loss.backward()
              optimizer.step()
              loss_sum += loss.item() * y.shape[0]
              n += y.shape[0]
          train_loss = loss_sum / n
          val_loss = evaluateModel(model, criterion, val_iter)
          if val_loss < min_val_loss:
              wait = 0
              min_val_loss = val_loss
              torch.save(model.state_dict(), PATH + '/' + name + '.pt')
          else:
              wait += 1
              if wait == PATIENCE:
                  print('Early stopping at epoch: %d' % epoch)
                  break
          endtime = datetime.now()
          epoch_time = (endtime - starttime).seconds
          print("epoch", epoch, "time used:", epoch_time, " seconds ", "train loss:", train_loss, "validation loss:",
                val_loss)
          with open(PATH + '/' + name + '_log.txt', 'a') as f:
              f.write("%s, %d, %s, %d, %s, %s, %.10f, %s, %.10f\n" % (
              "epoch", epoch, "time used", epoch_time, "seconds", "train loss", train_loss, "validation loss:", val_loss))
  
      torch_score = evaluateModel(model, criterion, train_iter)
      YS_pred = predictModel(model, torch.utils.data.DataLoader(trainval_data, BATCHSIZE, shuffle=False))
      print('YS.shape, YS_pred.shape,', YS.shape, YS_pred.shape)
      YS, YS_pred = scaler.inverse_transform(np.squeeze(YS)), scaler.inverse_transform(np.squeeze(YS_pred))
      print('YS.shape, YS_pred.shape,', YS.shape, YS_pred.shape)
      MSE, RMSE, MAE, MAPE = Metrics.evaluate(YS, YS_pred)
      with open(PATH + '/' + name + '_prediction_scores.txt', 'a') as f:
          f.write("%s, %s, Torch MSE, %.10e, %.10f\n" % (name, mode, torch_score, torch_score))
          f.write("%s, %s, MSE, RMSE, MAE, MAPE, %.10f, %.10f, %.10f, %.10f\n" % (name, mode, MSE, RMSE, MAE, MAPE))
      print('*' * 40)
      print("%s, %s, Torch MSE, %.10e, %.10f" % (name, mode, torch_score, torch_score))
      print("%s, %s, MSE, RMSE, MAE, MAPE, %.10f, %.10f, %.10f, %.10f" % (name, mode, MSE, RMSE, MAE, MAPE))
      print('Model Training Ended ...', time.ctime())
  
  
  def testModel(name, mode, XS, YS):
      if LOSS == "MaskMAE":
          criterion = Utils.masked_mae
      if LOSS == 'MSE':
          criterion = nn.MSELoss()
      if LOSS == 'MAE':
          criterion = nn.L1Loss()
      print('Model Testing Started ...', time.ctime())
      print('TIMESTEP_IN, TIMESTEP_OUT', TIMESTEP_IN, TIMESTEP_OUT)
      XS_torch, YS_torch = torch.Tensor(XS).to(device), torch.Tensor(YS).to(device)
      test_data = torch.utils.data.TensorDataset(XS_torch, YS_torch)
      test_iter = torch.utils.data.DataLoader(test_data, BATCHSIZE, shuffle=False)
      model = getModel(name)
      model.load_state_dict(torch.load(PATH + '/' + name + '.pt'))
  
      torch_score = evaluateModel(model, criterion, test_iter)
      YS_pred = predictModel(model, test_iter)
      print('YS.shape, YS_pred.shape,', YS.shape, YS_pred.shape)
      YS, YS_pred = scaler.inverse_transform(np.squeeze(YS)), scaler.inverse_transform(np.squeeze(YS_pred))
      print('YS.shape, YS_pred.shape,', YS.shape, YS_pred.shape)
      np.save(PATH + '/' + MODELNAME + '_prediction.npy', YS_pred)
      np.save(PATH + '/' + MODELNAME + '_groundtruth.npy', YS)
      MSE, RMSE, MAE, MAPE = Metrics.evaluate(YS, YS_pred)
      print('*' * 40)
      print("%s, %s, Torch MSE, %.10e, %.10f" % (name, mode, torch_score, torch_score))
      f = open(PATH + '/' + name + '_prediction_scores.txt', 'a')
      f.write("%s, %s, Torch MSE, %.10e, %.10f\n" % (name, mode, torch_score, torch_score))
      print(
          "all train steps, %s, %s, MSE, RMSE, MAE, MAPE, %.10f, %.10f, %.10f, %.10f" % (name, mode, MSE, RMSE, MAE, MAPE))
      f.write("all train steps, %s, %s, MSE, RMSE, MAE, MAPE, %.10f, %.10f, %.10f, %.10f\n" % (
      name, mode, MSE, RMSE, MAE, MAPE))
      for i in range(TIMESTEP_OUT):
          MSE, RMSE, MAE, MAPE = Metrics.evaluate(YS[:, i, :], YS_pred[:, i, :])
          print("%d step, %s, %s, MSE, RMSE, MAE, MAPE, %.10f, %.10f, %.10f, %.10f" % (
          i + 1, name, mode, MSE, RMSE, MAE, MAPE))
          f.write("%d step, %s, %s, MSE, RMSE, MAE, MAPE, %.10f, %.10f, %.10f, %.10f\n" % (
          i + 1, name, mode, MSE, RMSE, MAE, MAPE))
      f.close()
      print('Model Testing Ended ...', time.ctime())
  
  
  ################# Parameter Setting #######################
  MODELNAME = 'Double_C_STTN'
  KEYWORD = 'pred_' + DATANAME + '_' + MODELNAME + '_' + datetime.now().strftime("%y%m%d%H%M")
  PATH = '../save/' + KEYWORD
  torch.manual_seed(100)
  torch.cuda.manual_seed(100)
  np.random.seed(100)
  # torch.backends.cudnn.deterministic = True
  ###########################################################
  GPU = 1
  device = torch.device("cuda:{}".format(GPU)) if torch.cuda.is_available() else torch.device("cpu")
  ###########################################################
  data = np.load(C_2_PATH)
  
  data = data.transpose(2,0,1)
  scaler = StandardScaler()
  
  data[0] = scaler.fit_transform(data[0])
  
  data = data.transpose(1,2,0)
  print('data.shape', data.shape)
  
  
  ###########################################################
  def main():
      if not os.path.exists(PATH):
          os.makedirs(PATH)
      currentPython = sys.argv[0]
      shutil.copy2(currentPython, PATH)
      shutil.copy2('model/STTN.py', PATH)
      shutil.copy2('Param.py', PATH)
  
  
      print(KEYWORD, 'training started', time.ctime())
      trainXS, trainYS = getXSYS(data, 'TRAIN')
      print('TRAIN XS.shape YS,shape', trainXS.shape, trainYS.shape)
      trainModel(MODELNAME, 'train', trainXS, trainYS)
  
      print(KEYWORD, 'testing started', time.ctime())
      testXS, testYS = getXSYS(data, 'TEST')
      print('TEST XS.shape, YS.shape', testXS.shape, testYS.shape)
      testModel(MODELNAME, 'test', testXS, testYS)
  
  
  if __name__ == '__main__':
      main()
  
  
  
  
  