train.py

# -*- coding: utf-8 -*-
"""Train.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1ZiDKWFQHoVB5TGKtB00aidQ7jdYGXsz_

# Space Weather

Notebook to reproduce/improve forecasts obtained by paper by Siciliano et at. *Forecasting SYM‐H Index: A Comparison Between Long Short‐Term Memory and Convolutional Neural Networks* ([paper](https://agupubs.onlinelibrary.wiley.com/doi/epdf/10.1029/2020SW002589)).
"""
import sys

sys.path.append('lib')

from trainUtils import *
from trainModels import *



# ===============================================================
#  main()
# ===============================================================
def main():

    from optparse import OptionParser
    parser = OptionParser(usage = "usage: %prog arguments", version="%prog")
    parser.add_option("--config",           dest="config",                              help="configuration file (default: %default)") 
    parser.add_option("--dnn",              dest="dnn",                                 help="Type of DNN to use (default: %default)")
    parser.add_option("--optuna",           dest="optuna",     action = "store_true",   help="Run Optimization (default: %default)")
    parser.add_option("--plot",             dest="plot",       action = "store_true",   help="Plots dnn nominal (default: %default)")
    parser.add_option("--weighed",          dest="weighed",    action = "store_true",   help="adds weigh to sample values before peaks during training (default: %default")
    parser.add_option("--dropout",          dest="dropout",    action = "store_true",   help="toggles toggle, overrides bootstrap (default: %default)")
    parser.add_option("--derivative",       dest="derivative", action = "store_true",   help="toggles transformation of the target variable into its time derivative (default: %default)")
    parser.add_option("--cv",               dest="cv",         action = "store_true",   help="toggles cross validation (default: %default)")
    parser.add_option("--aug",              dest="aug",        action = "store_true",   help="toggles data augmentation with time translation (default: %default)")
    parser.add_option("--RMSE",             dest="RMSE",                                help="0 to calc RMSE through iterations, 1 through time (default: %default)")
    parser.add_option("--DOBIG",            dest="DOBIG",      action = "store_true",   help="Use p=0.31 for dropout (default: %default)")

    parser.set_defaults(config="config/config.ini",  dnn="LSTM",  optuna=False, plot=False, weighed = False, dropout = False, derivative = False, cv = False, aug = False, RMSE = '0', DOBIG = False)
    (options,args) = parser.parse_args()

    from datetime import datetime
    print( "\n Starting Processing at : " + datetime.now().strftime('%H:%M:%S'))


    
    # ###############################################################
    # Read the configuration file and check for needed parameters
    # ###############################################################
    parameters = readConfig(options.config)

    # Type 'ACE' or 'EBRE' or 'ACE_EBRE'
    ml_data = parameters['Dataset']
    ML_DATA = ml_data
    ml_data_index = ['ACE','ACE_EBRE','EBRE'].index(ml_data)
    # ML_DATA saves original name, but functionally EBRE and ACE_EBRE are to be treated the same under ml_data
    if ML_DATA == 'EBRE':
        ml_data = 'ACE_EBRE'

    dnn_name = eval(parameters['DNN_models'])
    dnn_type = dnn_name[options.dnn]

    batch_size = int(parameters['batch_size'])
    epochs = int(parameters['epochs'])

    step = int(parameters['step']) # 5 mins for training only
    lf =  int(parameters['lf']) # 60 min
    
    # Number of repetitions for systematic evaluations
    n_syst = int(parameters['n_syst'])
    
    # Weather to calculate RMSE through time or through RT iterations
    RMSE_type = int(options.RMSE)
    RMSE_prefix = ''
    if RMSE_type != 0:
        RMSE_prefix = 'RMSEalt_'

    Dropout = options.dropout
    # Use big p for dropout toggle
      # big p was an experiment on using an equivalent p (in terms of hyperparameter optimization) that was bigger
    DOBIG = options.DOBIG

    # Weather to weigh the high value y values for training or not
      # this means giving higher weight in training to values of the target variable that are big. Not used ultimately
    weigh_peaks = options.weighed
    if weigh_peaks:
        weigh_prefix = 'WEIGHED_'
    else:
        weigh_prefix = ''

    # ###############################################################
    # Read the input data
    # ###############################################################

    all_points=True
    if dnn_type == 2 :
      all_points = False

    # Base Dir
    base_dir = os.getcwd()

    # Input directory
    drive_dir = os.path.join(base_dir, ml_data)

    # Output directory
    output_dir = os.path.join(drive_dir, 'Models/')
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)


    """# Load Train, Validation and Test

    Create train and validation data frames
    """

    selected_features = []
    dataset_types = ['Train', 'Validation', 'Test']

    ini_data = {}

    data = {}
    stat_data = {}

    file_name =os.path.join(drive_dir, 'all_data.gzip')
    all_data = pd.read_parquet(file_name)


    # Dataframe vectors
    for dset in dataset_types :
      # initial data
      file_name =os.path.join(drive_dir, dset+'.gzip')
      ini_data[dset] = pd.read_parquet(file_name)

      file_name =os.path.join(drive_dir, dset+'.json')
      with open(file_name) as json_file:
          stat_data[dset] = json.load(json_file)

      selected_features = stat_data[dset]['features'] # this may be modified later
      # The exc list is the index of variables we want to exclude
      exc = []
      if len(exc) > 0 and ML_DATA == 'ACE_EBRE':
          selected_features_ = [selected_features[:a-i]+selected_features[a+1-i:] for i,a in enumerate(exc)]
          selected_features = selected_features_[0]

    print(selected_features)
    print(all_data.columns)

    # This part reduces the variables to only the three components of the magnetic field obtained from EBRE
    if ML_DATA == 'EBRE':
        all_data = all_data.drop(selected_features[5:]+['B**2','BY**2'],axis=1)
        for dset in dataset_types:
            ini_data[dset] = ini_data[dset].drop(selected_features[5:]+['B**2','BY**2'],axis=1)
        selected_features = selected_features[:3]
        print(selected_features)
    print(ini_data)

    print('--------------------------------------------')

    predict_v = ['SYM/H nT','EBRX-P','EBRX-P'][ml_data_index] # Designates target variable according to ml_data

    """""""""DERIVATIVE"""""""""
    # Ultimately not computed, for future use

    # These lines compute the first order derivative of the predict variable with respect to time according to function finite_diff in trainUtils.py
    predict_time_derivative = options.derivative # Trigger
    derivative_hp_index = 0 # this chooses among the hyperparameter sets (0 meaning not chosen, 1, meaning chosen)
    derivative_prefix = ''
    if predict_time_derivative:
        derivative_prefix = "DIFF_"
        derivative_hp_index = 1
        k_d = 20 # derivative window
        all_data[predict_v] = (finite_diff(np.array(all_data[predict_v]), k = k_d))/5. # set up at 'predict_v', this substitutes original target variable
        aa, bb = stat_data['Train']['breaks'][0], stat_data['Train']['breaks'][1] # aa and bb are for plotting purposes inside the function
        ini_data['Train'][predict_v] = (finite_diff(np.array(ini_data['Train'][predict_v]), k = k_d, plot_d = True, aa = aa, bb = bb))/5.
        ini_data['Test'][predict_v] = (finite_diff(np.array(ini_data['Test'][predict_v]), k = k_d))/5.
        ini_data['Validation'][predict_v] = (finite_diff(np.array(ini_data['Validation'][predict_v]), k = k_d))/5.

    """""""""Crossing up the splits"""""""""
    # This section could include cross validation. For now, it just mixes up the storms between test, train and validation splits
    # Training after this crossing was done only once

    cv = options.cv
    CV_prefix = ''
    # This segment produces a list with the new tags to be placed on the all_data dataframe
    # Idea is to switch all test storms to train, 3 validation storms to train, and turn 17 train storms to test and 3 to validation
    if cv:
      print('hola a hacer cv')
      CV_prefix = 'cv_'
      dtype_new = []
      dataset_new = []
      for i in range(len(all_data)):
        if all_data['Type_x'][i] == 'Train' and int(all_data['Dataset_x'][i][1:]) < 18:
          dtype_new.append('Test')
          dataset_new.append(all_data['Dataset_x'][i])
        elif all_data['Type_x'][i] == 'Train' and int(all_data['Dataset_x'][i][1:]) > 17:
          dtype_new.append('Validation')
          dataset_new.append('V'+str(int(all_data['Dataset_x'][i][1:])-15))
        elif all_data['Type_x'][i] == 'Validation' and int(all_data['Dataset_x'][i][1:]) > 2:
          dtype_new.append('Train')
          dataset_new.append('T'+str(int(all_data['Dataset_x'][i][1:])+15))
        elif all_data['Type_x'][i] == 'Test':
          dtype_new.append('Train')
          dataset_new.append(all_data['Dataset_x'][i])
        else:
          dtype_new.append('Validation')
          dataset_new.append(all_data['Dataset_x'][i])
  
        # This verifies the process
        #if i==0:
          #print(i,dataset_new[i],dtype_new[i],all_data['Dataset_x'][i],all_data['Type_x'][i])
        #if i>0 and dataset_new[i] != dataset_new[i-1]:
          #print(i,dataset_new[i],dtype_new[i],all_data['Dataset_x'][i],all_data['Type_x'][i])

      all_data['dataset_new'] = dataset_new
      all_data['dtype_new'] = dtype_new

      #Also, all_data tags needs to be replaced
      all_data['Type'] = dtype_new
      all_data['Dataset'] = dataset_new
      all_data['Type_x'] = dtype_new
      all_data['Dataset_x'] = dataset_new

      # This chooses only the columns that are also in ini_data, this needs to not replace all_data, so new name
      ini_full = pd.concat([ini_data[i] for i in ['Train','Test','Validation']]) # might be overkill?
      all_data_selcols = all_data[list(ini_full.columns)+['dataset_new','dtype_new']]

      # final_dataset will replace ini_data

      dset_order = [["T{}".format(i) for i in range(1,21)],["T{}".format(i) for i in range(1,18)],["V{}".format(i) for i in range(1,6)]]
      dtype_order = ['Train', 'Test', 'Validation']

      final_dataset = {}
      for k,i in enumerate(dtype_order):
        final_dataset[i] = pd.concat([all_data_selcols[all_data_selcols.dataset_new.eq(j) & all_data_selcols.dtype_new.eq(i)] for j in dset_order[k]])

      ini_data = {}
      ini_data = final_dataset

      # stat_data_new replaces stat_data[dtype]['breaks']
      stat_data_new = {}

      for m,k in enumerate(dtype_order):
        f_k = final_dataset[k]
        bb = [len(f_k[f_k.dataset_new.eq(dset_order[m][i])]) for i in range(len(dset_order[m]))]
        cc = [sum(bb[:i]) for i in range(0,len(bb)+1)]
        stat_data_new[k+'Breaks'] = cc

      for i in dtype_order:
        stat_data[i]['breaks'] = stat_data_new[i+'Breaks']


    if cv:
      trainBreaks, testBreaks, validationBreaks = stat_data_new['TrainBreaks'], stat_data_new['TestBreaks'], stat_data_new['ValidationBreaks']
    else: # either way, these new variables need to be assigned
      trainBreaks, testBreaks, validationBreaks = stat_data['Train']['breaks'], stat_data['Test']['breaks'], stat_data['Validation']['breaks']

    """""""""STANDARIZATION CORRECTION"""""""""

    # The "_norm" files were incorrectly standarized with respect to their own dataset, they need to be standardized wrt the training set
    print(selected_features,ml_data)
    print(ini_data['Train'])

    # In addition, I'm adding these physically motivated variables to the EBRE+ACE case:
      # Bz_neg_switch is a number that indicates if Bz is negative or not
      # B**2 acc is a number that aggregates the value of B**2 only when BZ is negative
    if ML_DATA == 'ACE_EBRE':
      tr_full = ini_data['Train'].reset_index().iloc[:,1:]
      te_full = ini_data['Test'].reset_index().iloc[:,1:]
      va_full = ini_data['Validation'].reset_index().iloc[:,1:]

      tr_full['Bz_neg_switch'] = np.array([1 if x<0 else 0 for x in tr_full['BZ nT (GSM)']])
      te_full['Bz_neg_switch'] = np.array([1 if x<0 else 0 for x in te_full['BZ nT (GSM)']])
      va_full['Bz_neg_switch'] = np.array([1 if x<0 else 0 for x in va_full['BZ nT (GSM)']])

      def ACCB2(dataset):
        B_acc = [0]
        for x,y in enumerate(dataset['Bz_neg_switch']):
          if y == 1:
            B_acc.append(B_acc[-1]+dataset['B**2'].iloc[x])
          else:
            B_acc.append(B_acc[-1])
        return(np.array(B_acc[1:]))

      tr_full['B**2_acc'] = ACCB2(tr_full)
      te_full['B**2_acc'] = ACCB2(te_full)
      va_full['B**2_acc'] = ACCB2(va_full)

      selected_features = selected_features[:3]+['Bz_neg_switch','B**2_acc']

      tr = tr_full[selected_features]
      te = te_full[selected_features]
      va = va_full[selected_features]

      # standarization (incorrectly named normalize as function)
      train_means, train_stds, data['Train'] = normalize(tr)
      _, _, data['Test'] = normalize(te, data_mean = train_means, data_std = train_stds)
      _, _, data['Validation'] = normalize(va, train_means, train_stds)

      # now, same needs to be done for the all_data dataFrame.
        # However, the storms in that dataframe are identified in the Type_x and Dataset_x variables
      size = [20,5,17]
      dataset = ['Train','Validation','Test']
      ids = [] # produces names T1,...T20,V1,...,V5,T1,...,T17
      for x in range(3):
        for y in range(size[x]):
          ids.append([dataset[x],'{}{}'.format(dataset[x][0],y+1)])

      z_switch = np.array([])
      b_agg = np.array([])
      for x,y in enumerate(ids):
        df = all_data[all_data.Type_x.eq(y[0]) & all_data.Dataset_x.eq(y[1])]

        z_switch_i = np.array([1 if x<0 else 0 for x in df['BZ nT (GSM)']])
        df['Bz_neg_switch'] = z_switch_i
        b_agg_i = ACCB2(df)

        if x == 0:
          z_switch = z_switch_i
          b_agg = b_agg_i
        else:
          z_switch = np.concatenate((z_switch,z_switch_i))
          b_agg = np.concatenate((b_agg,b_agg_i))

      all_data['Bz_neg_switch'] = z_switch
      all_data['B**2_acc'] = b_agg

      print("after standardization selected_features: ",selected_features)
    else: # stardardization for ACE only dataframes
      tr = ini_data['Train'].reset_index().iloc[:,1:]
      te = ini_data['Test'].reset_index().iloc[:,1:]
      va = ini_data['Validation'].reset_index().iloc[:,1:]

      train_means, train_stds, data['Train'] = normalize(tr)
      _, _, data['Test'] = normalize(te, data_mean = train_means, data_std = train_stds)
      _, _, data['Validation'] = normalize(va, train_means, train_stds)
#    train_stds[0] = 1 # temporary diagnostics line

    print(data['Train'].head())
    print(data['Test'].head())
    print(data['Validation'].head())


    """""""""OPTUNA"""""""""
    """# Hyperparameter Optimization
        
    Optimization of hyperparameters using "Optuna" framework
    """
        
    """Run Optuna"""
    print("Lookforward: ", lf)
    if options.optuna:
        N_TRIALS = int(parameters['N_TRIALS']) # around 100 usually
        print( "N_TRIALS: ",  N_TRIALS)

        # Configuration parameters
        save_dir = 'saved_models_optuna'
        TIMEOUT = min(N_TRIALS * 1000, 47*3600)
        BEST_VALUE = np.inf
        BEST_HISTORY = None
        
        if N_TRIALS > 0:
            # Define Early Stopping
            es_callback = keras.callbacks.EarlyStopping(monitor="val_mse", min_delta=0.002, patience=5)
            
            # Create Optuna Study
            study = optuna.create_study(
                    sampler=optuna.samplers.TPESampler(multivariate=True),
                    study_name="first_global",
                    direction="minimize",
                    pruner=optuna.pruners.MedianPruner(
                        n_startup_trials=10, n_warmup_steps=10, interval_steps=5
                        ),
                    )
            
            # Run Optimization. Lambda call is to include all input variables for the various functions used to create the datasets
            study.optimize(lambda trial: objective(trial, options.config, data,stat_data,dnn_type, selected_features, es_callback, weigh_peaks),
                           n_trials=N_TRIALS, timeout=TIMEOUT, show_progress_bar=True, catch=())
            
            
            study.trials_dataframe().sort_values(by="value").to_csv('optuna_trials_global_'+options.dnn+'.csv')
            
            optuna.visualization.matplotlib.plot_contour(study)
            plt.savefig(str(lf) + RMSE_prefix + CV_prefix + derivative_prefix+'_'+ML_DATA+"{}hparams_contour_{}_{}.png".format(weigh_prefix,options.dnn,N_TRIALS))
            
            optuna.importance.get_param_importances(study)
            optuna.visualization.matplotlib.plot_param_importances(study)
            plt.savefig(str(lf) + RMSE_prefix + CV_prefix + derivative_prefix+'_'+ML_DATA+"{}hparam_importances_{}_{}.png".format(weigh_prefix,options.dnn,N_TRIALS))

            best_parameters = study.best_params
            display(study.trials_dataframe())
            print("\n",best_parameters)

            study.trials_dataframe().sort_values(by="value").to_csv(output_dir+'optuna_trials_global_'+options.dnn+'.csv')
            
            lb =  int(best_parameters['lb'])
            lr=best_parameters['lr']
            n_layers=best_parameters['n_layers']
            n_D=best_parameters['n_unit']
            
            print("lb: ",lb, "lr: ", lr, "n_D: ", n_D, "n_layers: ", n_layers)

            
    else: # if Optuna was not used, use pre-defined hyperparameter values
        print("No optuna")
        # Reading from config 
        # Look back
        lb =  int(eval(parameters['lb'])[ml_data_index][derivative_hp_index]) # 72 == 360 min
        # Training parameters
        lr = float(eval(parameters['lr'])[ml_data_index][derivative_hp_index]) # 0.00005
        n_D = int(eval(parameters['n_D'])[ml_data_index][derivative_hp_index])
        n_layers = int(eval(parameters['n_layers'])[ml_data_index][derivative_hp_index])
      
        if Dropout: # this could be included somehow in config, but here it is placed for now
            if not DOBIG:
                lb, lr, n_layers, n_D = 40, 0.000734, 1, 162
            else:
                lb, lr, n_layers, n_D = 360, 0.000771, 4, 40

        # Hyper-parameters for look forwards 30min, 90min, 120min, 150min, ... , 300min, respectively
        lb_lf = [150,40,195,40,75,75,40,40,40]
        lr_lf = [5.829493395977575e-05, 1.4095949551427182e-05, 2.4374952158965335e-05,
                 1.311089599011205e-05, 3.732906745895542e-05, 1.529362160868869e-05,
                 1.5824377523090944e-05, 2.5219948598180393e-05, 1.0933260846326537e-05]
        n_layers_lf = [7, 7, 19, 0, 1, 12, 9, 9, 5]
        n_d_lf      = [181, 136, 57, 455, 52, 133, 199, 74, 150]
        if lf != 12: # note the jump from 30 to 90. Only reassigning if not predicting in 1 hour.
            lf_index = [1,3,4,5,6,7,8,9,10].index(int(lf)//6)
            print(lf_index)
            for lf_i in [6,18,24,30,36,42,48,54,60]:
                print(lf_i)
                lfff = [1,3,4,5,6,7,8,9,10].index(int(lf_i)//6)
                print(lfff)
                print(lb_lf[lfff], lr_lf[lfff], n_layers_lf[lfff], n_d_lf[lfff])
            lb, lr, n_layers, n_D = lb_lf[lf_index], lr_lf[lf_index], n_layers_lf[lf_index], n_d_lf[lf_index]
        print("lb: ",lb, "lr: ", lr, "n_D: ", n_D, "n_layers: ", n_layers)

    # x_train and y_train will be overwritten in run systematics because of the bootstrapped sets
    x_train, y_train = create_dataset(data['Train'], dnn_type, lb, lf, selected_features,
                                      len(selected_features),interface=trainBreaks,
                                      step=step, predict_time_derivative = predict_time_derivative)
    x_val, y_val = create_dataset(data['Validation'], dnn_type, lb, lf, selected_features,
                                  len(selected_features),interface=validationBreaks,
                                  step=step, predict_time_derivative = predict_time_derivative)
    x_test, y_test = create_dataset(data['Test'], dnn_type, lb, lf, selected_features,
                                    len(selected_features),interface=testBreaks,
                                    step=1, predict_time_derivative = predict_time_derivative)

    print("test shapes--------------------\n",x_test.shape,y_test.shape)
    
    
    """# Run Systematics
    Run several times same configuration for systematics and check best value after preparing bootstrapped (or dropout-ed) train sets
    """

    # Choosing stochastic regularization method between bootstrap and dropout
    # bootstrap trains n_syst different models and predicts once for each, whereas dropout trains one model and applies dropout to model.predict during inference n_syst times
    block_bootstrap = True

    # Block bootstrap, with each chunk being a different storm
    from random import randrange
    # bs_indices is the array with the indices of bootstrap-sampled training storms (20 available storms)
    if block_bootstrap:
        bs_indices = np.array([[randrange(20) for i in range(20)] for j in range(n_syst)])
    elif Dropout:
        bs_indices = np.array([[i for i in range(20)] for j in range(n_syst)])
    print(bs_indices)

    # Keys is a list of names "T1,..,T17", and key_index will be a list with the index in all_data where each storm starts
    keys = all_data.loc[all_data['Type'] == 'Test']['Dataset'].unique()
    keys_train = all_data.loc[all_data['Type'] == 'Train']['Dataset'].unique()
    key_index = [0]
    key_train_index = [0]

    for dset_key in keys:
      key_index.append(key_index[-1]+len(all_data[all_data.Type.eq('Test') & all_data.Dataset.eq(dset_key)]))
    for dset_key in keys_train:
      key_train_index.append(key_train_index[-1]+len(all_data[all_data.Type.eq('Train') & all_data.Dataset.eq(dset_key)]))
    if cv:
      key_train_index = trainBreaks

    # data_train_bs is the data for the bootstrapped samples, both x and y, to be called when creating the training datasets
    data_train_bs = []
    for i, j in enumerate(bs_indices):
      bs_set = pd.DataFrame()
      c = 0
      for k in j:
        a, b = key_train_index[k], key_train_index[k+1]
        df = data['Train'].iloc[a:b]
        c += len(df)
        bs_set = bs_set.append(df, ignore_index = True)
      data_train_bs.append(bs_set)

    # these arrays store the values of each run
    mse = np.zeros(n_syst)
    r2  = np.zeros(n_syst)
    y_predict = []
    y_predict_storm = []
    # RMSE_type was an experiment on the computation of RMSE among all storms
      # RMSE_type = 0 means that the RMSE of each set of predictions is computed first, and then averaged among all RMSE_predict
    if RMSE_type == 0:
      rmse_storm = np.zeros((len(keys),n_syst))
      r2_storm = np.zeros((len(keys),n_syst))
    else:
      rmse_storm = []
      r2_storm   = []

    if Dropout: # re-tests are done in a loop later inside the following i loop for dropout to use the predict function
        N_SYST = 1 # only need to train one model
        concrete_dropout = True # applies dropout to each dense layer
        if DOBIG:
            dropout_p = 0.316118
            suffix = '_withDropoutBig'
        else:
            dropout_p = 0.012794 # result from optuna
            suffix = '_withDropout'
    else:
        N_SYST = n_syst
        dropout_p = 0
        concrete_dropout = False
        suffix = '_withBootStrap'
    
    # Variables for feature importance
    fi_results = []
    fi_results2 = []

    # Model predictions start
    for i in range(N_SYST):
      print(N_SYST,i)
      if True: # if no .h5 models built already
          # x_train until now because it waits for bootstrap config for data_train_bs[i]
          x_train, y_train = create_dataset(data_train_bs[i], dnn_type, lb, lf, selected_features, len(selected_features),
                                            interface=key_train_index, step=step, predict_time_derivative = predict_time_derivative)

          # Feature importance setup. Done once and during training, F2 was ultimately the approach that worked (shuffling all but one variable each)
          FI = False
          FI1 = False and FI
          FI2 = False and FI
          FI_size = 15
          FI_total = len(selected_features)*FI_size
          if i<FI_total and FI:
              print("FI",FI,FI1,FI2,i,selected_features[i//FI_size])
              save_col = np.array([x_train[:,:,zzz].copy() for zzz in range(len(selected_features))])
              if FI1:
                  np.random.shuffle(x_test[:,:,i//FI_size])
              for shuffle in range(len(selected_features)):
                  if shuffle != i//FI_size and FI2:
                      print(shuffle, FI2, selected_features[shuffle])
                      np.random.shuffle(x_train[:,:,shuffle])
          
          # Model training starts
          es_callback = keras.callbacks.EarlyStopping(monitor="val_mse", min_delta=0, patience=5)
          lr_callback = keras.callbacks.LearningRateScheduler(custom_lr_cycle) # Cyclical Learning Rate (triangular, centered around lr, stepsize 10)
          model = build_dnn(dnn_type,x_train.shape[1:], y_train.shape[1],  lr=lr, nlayers=n_layers, nD=n_D, dropout_p = dropout_p, concrete_dropout = concrete_dropout)

          # build a sample_weight array if weigh_peaks. This array will priorize cost losses during training to values up to 500 minutes before the lowest peak of a storm
          sample_weight = np.ones(y_train.shape) # this stays as an array of ones if weigh_peaks is not called and so nothing changes during training
          if weigh_peaks:
              y_train_abs = np.absolute(y_train)
              max_ytrain = max(y_train_abs)[0]
              for p in range(len(y_train_abs)):
                  if 2*y_train_abs[p][0] > max_ytrain and 0 <= i-100: # 100*5 mins seems to be the timescale before half-peak
                      for pp in range(100): # this for loop is redundant, but cheap, so whatever. It ensures all values in between are weighed
                          sample_weight[p-pp][0] = 5 # why 5? because the number 5 is 5 times greater than 1...

          # fitting, with CLR and weights array as set up earlier
          print("before training shapes: ", x_train.shape, y_train.shape, sample_weight.shape)
          history = model.fit(x_train, y_train,
                              batch_size=batch_size,
                              epochs=epochs,
                              callbacks=[es_callback, lr_callback],#, modelckpt_callback],
                              verbose=0,
                              validation_data=(x_val, y_val),
                              sample_weight = sample_weight[:,0])

          print(ML_DATA+weigh_prefix+str(n_syst)+suffix+'.h5')
          # Saving .h5 files:
          if not Dropout:
              model.save('BS_models/{}'.format(options.dnn)+ML_DATA+weigh_prefix+'_{}outOf{}_lf{}'.format(i,n_syst,lf)+suffix+'.h5')
          else:
              model.save('DO_models/{}'.format(options.dnn)+ML_DATA+weigh_prefix+'_{}of{}_lf{}'.format(i,n_syst,lf)+suffix+'.h5')

          # Reseting x_train dataframe if doing feature importance:
          if i<60 and FI:
              for zzz in range(len(selected_features)):
                  x_train[:,:,zzz] = save_col[zzz]

      else:
          if not Dropout:
              model = models.load_model('BS_models/'+ML_DATA+weigh_prefix+'_{}outOf200_'.format(i)+suffix+'.h5', custom_objects = {'MonteCarloLSTM': MonteCarloLSTM})
          else:
              model = models.load_model('DO_models/'+ML_DATA+weigh_prefix+'_{}of{}_'.format(i,n_syst)+suffix+'.h5', custom_objects = {'MonteCarloLSTM': MonteCarloLSTM})

      # prediction and info lists for each run is different for each SRT, so an if is called
      if not Dropout:
        if False: # trigger feature importance calcs  
          ######## alternative cheap feature importance computation (post-training)
          COLS = selected_features
          print(' Computing LSTM feature importance...')
  
          fi_instance = [] # 5 element list: baseline + each feature
          fi_instance2 = []
        
          # Algorithm adapted from https://www.kaggle.com/code/cdeotte/lstm-feature-importance
          # COMPUTE BASELINE (NO SHUFFLE)
          oof_preds = model.predict(x_test, verbose=0).squeeze() 
          baseline_mae = np.mean(np.abs( oof_preds-y_test ))
          print(baseline_mae)
          fi_instance.append(baseline_mae)
          fi_instance2.append(baseline_mae)
  
          for k in range(len(COLS)):
              
              # SHUFFLE FEATURE K
              save_cols = [x_test[:,:,zzz].copy() for zzz in range(len(selected_features))]
              np.random.shuffle(x_test[:,:,k])
              x_test2 = x_test.copy()
              for shuffling in range(len(COLS)):
                  if shuffling != k:
                      np.random.shuffle(x_test2[:,:,shuffling])
              mae_list = []
              # COMPUTE OOF MAE WITH FEATURE K SHUFFLED
              oof_preds = model.predict(x_test, verbose=0).squeeze() 
              mae_i = np.mean(np.abs( oof_preds-y_test ))
              mae_list.append(mae_i)
              mae = np.mean(mae_list)
              fi_instance.append(mae)
              fi_instance2.append(np.mean(np.mean(np.abs(model.predict(x_test2,verbose=0).squeeze() - y_test))))
              for zzz in range(len(selected_features)):
                  x_test[:,:,zzz] = save_cols[zzz]
                  x_test2[:,:,zzz] = save_cols[zzz]

          print(fi_instance)
          print(fi_instance2)
          fi_results.append(fi_instance)
          fi_results2.append(fi_instance2)
          print('-------------------------------')

      ########## storing info for each storm
        test_breaks = [0] # supposed to be redundant with testBreaks
        y_test_storm = [] # test values for each storm (to be saved in .csv)
        y_predict_storm_i = [] # prediction values for each storm and each bootstrap sample (to be saved in .csv)
        for j in range(len(keys)):
          a, b = testBreaks[j-1], testBreaks[j]
          data_storm = data['Test'].iloc[a:b]
          x_test_storm, y_test_storm_i, _, _ = getStorm(all_data,stat_data,selected_features,dnn_type,lb,lf,
                                                        keys[j],dtype = 'Test',predict_time_derivative = predict_time_derivative)
          y_test_storm.append(y_test_storm_i)

          y_predict_storm_i.append(model.predict(x_test_storm))

          test_breaks.append(test_breaks[-1]+y_predict_storm_i[j].shape[0])

          if RMSE_type == 0:
            mse_i = mean_squared_error(y_predict_storm_i[j][:, 0], y_test_storm_i[:, 0])
            rmse_storm[j][i] = mse_i # later: de-standardization, this gives nT
            r2_i = r2_score(y_predict_storm_i[j][:, 0], y_test_storm_i[:, 0])
            r2_storm[j][i] = r2_i

        y_predict_storm.append(y_predict_storm_i) # this saves the iteration with 17 storms inside

      ########## storing general info of the run
        y_predict_i = model.predict(x_test)
        y_predict_i = np.reshape(y_predict_i,(y_predict_i.shape[0],1))
        y_predict.append(y_predict_i)
        mse_i = mean_squared_error(y_predict_i, y_test)
        mse[i] = mse_i
        r2_i = r2_score(y_predict_i, y_test)
        r2[i] = r2_i
        print("Run {} out of {} has rmse {} and r2_score {}".format(i,N_SYST,np.sqrt(mse_i),r2_i))
      else: # meaning: if using Dropout
        test_breaks = [0]
        y_test_storm = []
        y_predict_storm = []
        for j in range(len(keys)):
          a, b = testBreaks[j-1], testBreaks[j]
          data_storm = data['Test'].iloc[a:b]
          x_test_storm, y_test_storm_i, _, _ = getStorm(all_data,stat_data,selected_features,dnn_type,lb,lf,
                                                        keys[j],dtype = 'Test',predict_time_derivative = predict_time_derivative)
          y_test_storm.append(y_test_storm_i)

          y_predict_storm_j = np.array([model.predict(x_test_storm) for zz in range(n_syst)]) # this is where we infer n_syst sets of test values

          test_breaks.append(test_breaks[-1]+y_predict_storm_j.shape[1])
          
          if RMSE_type == 0:
              for xx in range(n_syst): # an RMSE for each instant in time. alternatively, this could be a for in iterations, not on time
                  rmse_storm[j][xx] = mean_squared_error(y_predict_storm_j[xx,:, 0], y_test_storm_i[:, 0])
                  r2_storm[j][xx] = r2_score(y_predict_storm_j[xx,:, 0], y_test_storm_i[:, 0])
          else:
              rmse_storm.append([])
              r2_storm.append([])
              T_s = len(y_predict_storm_j[0]) # total number of time instances in a storm s
              for xx in range(T_s): # an RMSE for each instant in time. alternatively, this could be a for in iterations, not on time
                  rmse_storm[j].append(mean_squared_error(y_predict_storm_j[:,xx, 0], n_syst*[y_test_storm_i[xx, 0]]))
                  r2_storm[j].append(r2_score(y_predict_storm_j[:,xx, 0], n_syst*[y_test_storm_i[xx, 0]]))
          
          y_predict_storm.append(y_predict_storm_j)
        a = [[y_predict_storm[k][zz] for k in range(len(keys))] for zz in range(n_syst)]
        y_predict_storm = a

        # storing general info of the run
        for xx in range(n_syst):
            y_predict_i = model.predict(x_test)
            y_predict_i = np.reshape(y_predict_i,(y_predict_i.shape[0],1))
            y_predict.append(y_predict_i)

            mse_i = mean_squared_error(y_predict_i, y_test)
            mse[xx] = mse_i
            r2_i = r2_score(y_predict_i, y_test)
            r2[xx] = r2_i
            print("Run {} has rmse {} and r2_score {}".format(xx,np.sqrt(mse_i),r2_i))

    # RMSE computation alternative, where RMSE_type == 1 means calculating the average of all bootstrap indices at once
    if RMSE_type != 0 and not Dropout:
      rmse_storm = []
      r2_storm   = []
      for s in range(len(keys)):
        rmse_storm.append([])
        r2_storm.append([])
        T_s = len(y_test_storm[s]) # total number of time instances in a storm s
        predicted_s = np.array([np.array([y_predict_storm[xx][s][tt][0] for tt in range(T_s)]) for xx in range(n_syst)])
        print(s,T_s, len(y_predict_storm[0][s]), len(y_predict_storm[0]))
        for t in range(T_s):
            print(s,t,predicted_s.shape)
#            print(predicted_s[:,t])
#            print('lueg',y_test_storm[s])
            mse_i = mean_squared_error(predicted_s[:,t], n_syst*[y_test_storm[s][t, 0]])
            rmse_storm[s].append(mse_i) # later: de-normalization, this gives nT
            r2_i = r2_score(predicted_s[:,t], n_syst*[y_test_storm[s][t, 0]])
            r2_storm[s].append(r2_i)
      #rmse_storm = np.array(rmse_storm)
      #r2_storm   = np.array(r2_storm)
    
    if len(mse)>0: # sometimes i turn off the run systematics section and this conditional avoids trouble
        for st in range(len(keys)):
#            st_id = 0
            # Saves the predictions in csv form, for use in colab notebook
            print("---------------saving .csv's segment---------------")
            T_s = len(y_predict_storm[0][st])
            print('steps in storm: ',T_s)
            print('number of storms total: ',len(keys),keys)
            predicted_s = np.array([[st,tt,xx,y_predict_storm[xx][st][tt][0],y_test_storm[st][tt][0]] for tt in range(T_s) for xx in range(n_syst)])
            print(predicted_s)
            np.savetxt('predictions{}/predict_storm{}{}{}{}_lf{}.csv'.format(suffix,st,suffix,options.dnn,ml_data,lf), predicted_s, delimiter=",")
            #print(predicted_s)
            print(predicted_s.shape, T_s)
#            test_s = np.array([[st,tt,y_test_storm[st][tt][0]] for tt in range(T_s)])
#            np.savetxt('testValues/test_storm{}{}{}{}_lf{}.csv'.format(st,suffix,options.dnn,ml_data,lf), test_s, delimiter = ',')
            print("-------------------end segment-----------------------")

#            y_predict_wide = np.array(y_predict_storm)[:,st]
#            y_predict_means, y_predict_std = y_predict_wide.mean(axis=0), y_predict_wide.std(axis=0)
#            #Calculating alternative MSE with y_predict_means and upper and mid 1std intervals
#            rmse_3 = np.sqrt(mean_squared_error(y_predict_means, y_test_storm[st]))*train_stds[0]
#            rmse_3_max = np.sqrt(mean_squared_error(y_predict_means + y_predict_std, y_test_storm[st]))*train_stds[0] # preliminary
#            rmse_3_min = np.sqrt(mean_squared_error(y_predict_means - y_predict_std, y_test_storm[st]))*train_stds[0] # preliminary

        if False: # trigger post-training feature importance calcs
          # Ultimately, these post-training feature importance computations were not shown, as they did not give significant differences between variables
          print("-------------Feature importance values-------------")
          for fi_i, fi_j in enumerate(fi_results):
              if fi_i%15 == 0:
                  print("variable shift")
              print(fi_j)
          print("feature importances for {}".format(COLS))
          print(np.mean(np.array(fi_results), axis = 0))
          print(np.std(np.array(fi_results),  axis = 0))
          print("-------------Feature importance values-------------")
          print("-------------Feature importance2 values-------------")
          for fi_i, fi_j in enumerate(fi_results2):
              if fi_i%15 == 0:
                  print("variable shift")
              print(fi_j)
          print("feature importances for {}".format(COLS))
          print(np.mean(np.array(fi_results2), axis = 0))
          print(np.std(np.array(fi_results2),  axis = 0))
          print("-------------Feature importance2 values-------------")
        mse_mean = np.mean(mse)
        mse_min = min(mse)
        mse_std = np.std(mse)

        for xx in range(len(keys)):
            print("rmse of storm {}".format(xx))
            print(np.sqrt(rmse_storm[xx])*train_stds[0])

        if False:
          ########### Feature Importance graphs #############
          disr = {}
          disr_means = {}
          disr_stds = {}
          COLS = ['sym','bz','b2','by2']
          COLS1 = COLS+['baseline']
          for i in range(5):
            if i != 3:
              disr_means[COLS1[i]] = np.mean(rmse_storm[:,FI_size*i:FI_size*i+FI_size])
              disr_stds[COLS1[i]] = np.std(np.mean(rmse_storm[:,FI_size*i:FI_size*i+FI_size],axis=0))
              print(i,COLS1[i],np.mean(rmse_storm[:,FI_size*i:FI_size*i+FI_size],axis=0))
            else:
              disr_means[COLS1[i]] = np.mean(rmse_storm[:,FI_size*i:FI_size*i+FI_size])
              disr_stds[COLS1[i]] = np.std(np.mean(rmse_storm[:,FI_size*i:FI_size*i+FI_size],axis=0))
              print(i,COLS1[i],np.mean(rmse_storm[:,FI_size*i:FI_size*i+FI_size],axis=0))
          print(disr_means,disr_stds)

          # DISPLAY LSTM FEATURE IMPORTANCE
          fi_results = np.array([disr_means[COLS1[-i-2]] for i in range(4)])
          fi_errors  = np.array([disr_stds[COLS1[-i-2]] for i in range(4)])

          baseline_mae = [disr_means['baseline'],disr_stds['baseline']]
          COLS = ['SYM-H','$B_z$','$B²$','$B_y^2$']
          print(fi_results,fi_errors)

          plt.figure(figsize=(10,5))
          plt.barh(np.arange(len(COLS)),fi_results,xerr=fi_errors, capsize = 5, label = 'RMSE after shuffle of other features')
          plt.yticks(np.arange(len(COLS)),[COLS[-i-1] for i in range(len(COLS))])
          #plt.xlim(0.82,0.9)
          plt.ylim((-1,len(COLS)))
          plt.plot([baseline_mae[0],baseline_mae[0]],[-1,len(COLS)+1], '--', color='orange',
                    label=f'Baseline RMSE = {baseline_mae[0]:.1f} nT')
          plt.axvspan(xmin = baseline_mae[0]-baseline_mae[1], xmax = baseline_mae[0]+baseline_mae[1],ymin = -1, ymax = len(COLS)+1,
                      color='orange', label=f'Baseline uncertainty RMSE = {baseline_mae[1]:.1f} nT', alpha = 0.25)
          plt.xlabel('RMSE (nT)',size=15)
          plt.ylabel('Feature',size=15)
          plt.legend(fontsize = 15)
          plt.savefig('feature_importance.pdf', format = 'pdf')
          ############ END Feature Importance Graph ##############

        rmse_storm_means = np.sqrt([np.mean(rmse_storm[xx]) for xx in range(len(keys))])*train_stds[0]
        rmse_storm_std   = 0.5*rmse_storm_means**(-0.5)*[np.std(rmse_storm[xx]) for xx in range(len(keys))]*train_stds[0] # error propagation of non normalized MSE
        r2_storm_means, r2_storm_std     = [np.mean(r2_storm[xx]) for xx in range(len(keys))], [np.std(r2_storm[xx]) for xx in range(len(keys))]
        print("Mean validation RMSE: {}\nMin validation RMSE: {}\nStandard Deviation validation RMSE: {}".format(np.sqrt(mse_mean),np.sqrt(mse_min),np.sqrt(mse_std)))
        print("Mean RMSE values for each storm: {} with standard deviations: {}".format(rmse_storm_means, rmse_storm_std))

        r2_mean = np.mean(r2)
        r2_std = np.std(r2)
        print("Mean validation r2: {}\nStandard Deviation validation r2: {}".format(r2_mean,r2_std))
        print("Mean r2 values for each storm: {} with standard deviation: {}".format(r2_storm_means, r2_storm_std))

        #make table for storms:
        columns = ['RMSE(nT)','RMSE error(nT)','R2','R2_error']
        stormTable = pd.DataFrame(np.array([np.around(rmse_storm_means,4),
                                            np.around(rmse_storm_std,4),
                                            np.around(r2_storm_means,4),
                                            np.around(r2_storm_std,4)]).T,
                                  columns = columns)
        stormTable_round = stormTable#.round(decimals = 4)
        print(stormTable_round)

# ===============================================================
#  __main__
# ===============================================================
if __name__ == '__main__':
    main()