code: import numpy as np
  import tensorflow as tf
  import time
  import datetime
  import os
  import sys
  import h5py
  from pathlib import Path
  import pandas as pd
  import matplotlib.pyplot as plt
  
  import evidential_deep_learning as edl
  from .util import normalize, gallery
  
  class Evidential:
      def __init__(self, model, opts, dataset="", learning_rate=1e-3, lam=0.0, epsilon=1e-2, maxi_rate=1e-4, tag="", custom_plot_folder="", custom_best_results_dat_folder=""):
          self.nll_loss_function = edl.losses.NIG_NLL
          self.reg_loss_function = edl.losses.NIG_Reg
  
          self.model = model
          self.learning_rate = learning_rate
          self.maxi_rate = maxi_rate
  
          self.optimizer = tf.optimizers.Adam(self.learning_rate)
          self.lam = tf.Variable(lam)
  
          self.epsilon = epsilon
  
          self.min_rmse = self.running_rmse = float('inf')
          self.min_nll = self.running_nll = float('inf')
          self.min_vloss = self.running_vloss = float('inf')
  
          trainer = self.__class__.__name__
          current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
          self.save_dir = os.path.join('save','{}_{}_{}_{}'.format(current_time, dataset, trainer, tag))
          Path(self.save_dir).mkdir(parents=True, exist_ok=True)
  
          self.custom_plot_folder = custom_plot_folder
          Path(self.custom_plot_folder).mkdir(parents=True, exist_ok=True)
          ''' Create 4 folders for valid/test plots for Epistemic/Aleatoric uncertainty'''
          Path(self.custom_plot_folder + '/valid_epistemic').mkdir(parents=True, exist_ok=True)
          Path(self.custom_plot_folder + '/valid_aleatoric').mkdir(parents=True, exist_ok=True)
          Path(self.custom_plot_folder + '/test_epistemic').mkdir(parents=True, exist_ok=True)
          Path(self.custom_plot_folder + '/test_aleatoric').mkdir(parents=True, exist_ok=True)
  
          self.custom_best_results_dat_folder = custom_best_results_dat_folder
          Path(self.custom_best_results_dat_folder).mkdir(parents=True, exist_ok=True)
  
          train_log_dir = os.path.join('logs', '{}_{}_{}_{}_train'.format(current_time, dataset, trainer, tag))
          self.train_summary_writer = tf.summary.create_file_writer(train_log_dir)
          val_log_dir = os.path.join('logs', '{}_{}_{}_{}_val'.format(current_time, dataset, trainer, tag))
          self.val_summary_writer = tf.summary.create_file_writer(val_log_dir)
  
      def loss_function(self, y, mu, v, alpha, beta, reduce=True, return_comps=False):
          nll_loss = self.nll_loss_function(y, mu, v, alpha, beta, reduce=reduce)
          reg_loss = self.reg_loss_function(y, mu, v, alpha, beta, reduce=reduce)
          loss = nll_loss + self.lam * (reg_loss - self.epsilon)
          # loss = nll_loss
  
          return (loss, (nll_loss, reg_loss)) if return_comps else loss
  
      @tf.function
      def run_train_step(self, x, y):
          with tf.GradientTape() as tape:
              outputs = self.model(x, training=True)
              mu, v, alpha, beta = tf.split(outputs, 4, axis=-1)
              loss, (nll_loss, reg_loss) = self.loss_function(y, mu, v, alpha, beta, return_comps=True)
  
          grads = tape.gradient(loss, self.model.trainable_variables) #compute gradient
          self.optimizer.apply_gradients(zip(grads, self.model.trainable_variables))
          self.lam = self.lam.assign_add(self.maxi_rate * (reg_loss - self.epsilon)) #update lambda
  
          return loss, nll_loss, reg_loss, mu, v, alpha, beta
  
      @tf.function
      def evaluate(self, x, y):
          outputs = self.model(x, training=False)
          mu, v, alpha, beta = tf.split(outputs, 4, axis=-1)
  
          rmse = edl.losses.RMSE(y, mu)
          loss, (nll, reg_loss) = self.loss_function(y, mu, v, alpha, beta, return_comps=True)
  
          return mu, v, alpha, beta, loss, rmse, nll, reg_loss
  
      def normalize(self, x):
          return tf.divide(tf.subtract(x, tf.reduce_min(x)),
                 tf.subtract(tf.reduce_max(x), tf.reduce_min(x)))
  
  
      def save_train_summary(self, loss, x, y, y_hat, v, alpha, beta):
          # np.random.seed(1234)
          with self.train_summary_writer.as_default():
              tf.summary.scalar('mse', tf.reduce_mean(edl.losses.MSE(y, y_hat)), step=self.iter)
              tf.summary.scalar('loss', tf.reduce_mean(self.loss_function(y, y_hat, v, alpha, beta)), step=self.iter)
              # pass
              idx = np.random.choice(int(tf.shape(x)[0]), 9)  ###  !!! give different model predictions !!!!!
              if tf.shape(x).shape==4:
                  tf.summary.image("x", [gallery(tf.gather(x,idx).numpy())], max_outputs=1, step=self.iter)
  
              if tf.shape(y).shape==4:
                  tf.summary.image("y", [gallery(tf.gather(y,idx).numpy())], max_outputs=1, step=self.iter)
                  tf.summary.image("y_hat", [gallery(tf.gather(y_hat,idx).numpy())], max_outputs=1, step=self.iter)
  
      # ''' Add additional saving for predicted y_hat and y_var '''
      # def save_train_summary_scalar(self, loss, x, y, y_hat, v, alpha, beta):
      #     with self.train_summary_writer.as_default():
      #         tf.summary.scalar('mse', tf.reduce_mean(edl.losses.MSE(y, y_hat)), step=self.iter)
      #         tf.summary.scalar('loss', tf.reduce_mean(self.loss_function(y, y_hat, v, alpha, beta)), step=self.iter)
      #         # tf.summary.scalar('y_hat_2', y_hat, step=self.iter)
      #         # var = beta / (v * (alpha - 1))
      #         # tf.summary.scalar('var', var, step=self.iter)
      #         idx = np.random.choice(int(tf.shape(x)[0]), 9)
      #         if tf.shape(x).shape == 4:
      #             tf.summary.image("x", [gallery(tf.gather(x, idx).numpy())], max_outputs=1, step=self.iter)
      #
      #         if tf.shape(y).shape == 4:
      #             tf.summary.image("y", [gallery(tf.gather(y, idx).numpy())], max_outputs=1, step=self.iter)
      #             tf.summary.image("y_hat", [gallery(tf.gather(y_hat, idx).numpy())], max_outputs=1, step=self.iter)
  
      def save_val_summary(self, loss, x, y, mu, v, alpha, beta):
          with self.val_summary_writer.as_default():
              tf.summary.scalar('mse', tf.reduce_mean(edl.losses.MSE(y, mu)), step=self.iter)
              tf.summary.scalar('loss', tf.reduce_mean(self.loss_function(y, mu, v, alpha, beta)), step=self.iter)
              idx = np.random.choice(int(tf.shape(x)[0]), 9)
  
              # var = beta / (v * (alpha - 1))
              # print('--Type of v: {}'.format(type(v)))
              # print('--Type of alpha: {}'.format(type(alpha)))
              # print('--Type of beta: {}'.format(type(beta)))
              # print('--Type of var: {}'.format(type(var)))
  
              if tf.shape(x).shape==4:
                  tf.summary.image("x", [gallery(tf.gather(x,idx).numpy())], max_outputs=1, step=self.iter)
  
              if tf.shape(y).shape==4:
                  tf.summary.image("y", [gallery(tf.gather(y,idx).numpy())], max_outputs=1, step=self.iter)
                  tf.summary.image("y_hat", [gallery(tf.gather(mu,idx).numpy())], max_outputs=1, step=self.iter)
                  var = beta/(v*(alpha-1))
                  tf.summary.image("y_var", [gallery(normalize(tf.gather(var,idx)).numpy())], max_outputs=1, step=self.iter)
  
      def get_batch(self, x, y, batch_size):
          idx = np.random.choice(x.shape[0], batch_size, replace=False)
          if isinstance(x, tf.Tensor):
              x_ = x[idx,...]
              y_ = y[idx,...]
          elif isinstance(x, np.ndarray) or isinstance(x, h5py.Dataset):
              idx = np.sort(idx)
              x_ = x[idx,...]
              y_ = y[idx,...]
  
              x_divisor = 255. if x_.dtype == np.uint8 else 1.0
              y_divisor = 255. if y_.dtype == np.uint8 else 1.0
  
              x_ = tf.convert_to_tensor(x_/x_divisor, tf.float32)
              y_ = tf.convert_to_tensor(y_/y_divisor, tf.float32)
          else:
              print("unknown dataset type {} {}".format(type(x), type(y)))
          return x_, y_
  
      def save(self, name):
          self.model.save(os.path.join(self.save_dir, "{}.h5".format(name)))
  
      def update_running(self, previous, current, alpha=0.0):
          if previous == float('inf'):
              new = current
          else:
              new = alpha*previous + (1-alpha)*current
          return new
  
  
      ''' Siyan added for testing plot '''
      def plot_scatter_with_var_from_pandas(self, x_train, y_train, x_test, y_test, mu, var, path, n_stds=3, test_bounds=[[-7, +7]], show=True):
          plt.scatter(x_train, y_train, s=1., c='#463c3c', zorder=0, label='Train (x_train vs y_train)')
          for k in np.linspace(0, n_stds, 4):
              if k == 0:
                  plt.fill_between(x_test[:, 0], (mu - k * var), (mu + k * var), alpha=0.3, edgecolor=None,
                                   facecolor='#00aeef', linewidth=0, antialiased=True, zorder=1, label='Unc.')
              else:
                  plt.fill_between(x_test[:, 0], (mu - k * var), (mu + k * var), alpha=0.3, edgecolor=None,
                                   facecolor='#00aeef', linewidth=0, antialiased=True, zorder=1)
  
          plt.plot(x_test, y_test, 'r--', zorder=2, label='True (x_test vs y_test)')
          plt.plot(x_test, mu, color='#007cab', zorder=3, label='Pred (x_test vs mu)')
          plt.gca().set_xlim(*test_bounds)
          plt.gca().set_ylim(-150, 150)
          plt.title(path)
          plt.legend()
          # print('asdfsdafdsafdsafsdafdsaf')
          # print(path)
          plt.savefig(path, transparent=True)
          if show:
              plt.show()
          plt.clf()
  
      ''' Siyan added for testing plot V2 '''
      def plot_scatter_with_var_from_pandas_v2(self, iter, x_train, y_train, x_test, y_test, mu, var, path, n_stds=3, test_bounds=[[-7, +7]], evalType='valid', uq='epistemic', show=True, save=False):
          # print(path)
          plt.scatter(x_train, y_train, s=1., c='#463c3c', zorder=0, label='Train (x_train vs y_train)')
          idxx = np.argsort(x_test.flatten())
          x_test = x_test.flatten()[idxx]
          y_test = y_test.flatten()[idxx]
          # mu = mu.numpy().flatten()[idxx]
          mu = mu.flatten()[idxx]
          var = var.flatten()[idxx]
          for k in np.linspace(0, n_stds, 4):
              if k == 0:
                  plt.fill_between(x_test, (mu - k * var), (mu + k * var), alpha=0.3, edgecolor=None,
                                   facecolor='#00aeef', linewidth=0, antialiased=True, zorder=1, label='Unc.')
              else:
                  plt.fill_between(x_test, (mu - k * var), (mu + k * var), alpha=0.3, edgecolor=None,
                                   facecolor='#00aeef', linewidth=0, antialiased=True, zorder=1)
  
          # idxx = np.argsort(x_test.flatten())
          # print(idxx)
          # plt.plot(x_test.flatten()[idxx], y_test.flatten()[idxx], 'r--', zorder=2, label='True (x_test vs y_test)')
          # plt.plot(x_test.flatten()[idxx], mu.flatten()[idxx], color='#007cab', zorder=3, label='Pred (x_test vs mu)')
  
          if evalType == 'valid':
              plt.plot(x_test, y_test, 'r--', zorder=2, label='True (x_valid vs y_valid)')
              plt.plot(x_test, mu, color='#007cab', zorder=3, label='Pred (x_valid vs mu)')
          elif evalType == 'test':
              plt.plot(x_test, y_test, 'r--', zorder=2, label='True (x_test vs y_test)')
              plt.plot(x_test, mu, color='#007cab', zorder=3, label='Pred (x_test vs mu)')
  
          plt.gca().set_xlim(*test_bounds)
          plt.gca().set_ylim(-150, 150)
          plt.title(path)
          if uq == 'epistemic':
              plt.suptitle('Epistemic uncertainty of {} data (Iters: {})\n'.format(evalType, iter), fontsize=18)
          elif uq == 'aleatoric':
              plt.suptitle('Aleatoric uncertainty of {} data (Iters: {})\n'.format(evalType, iter), fontsize=18)
  
          plt.legend()
          if save:
              plt.savefig(path)
          if show:
              plt.show()
          plt.clf()
  
  
  
      def train(self, x_train, y_train, x_test, y_test, x_valid, y_valid, y_scale, batch_size=128, iters=10000, verbose=True):
          ''' Siyan added START '''
          # np.random.seed(1234)
          test_eval_count = 0
          valid_eval_count = 0
          test_best_iter_str = '0'
          valid_best_iter_str = '0'
  
          x_valid_input = tf.convert_to_tensor(x_valid, tf.float32)
          y_valid_input = tf.convert_to_tensor(y_valid, tf.float32)
  
          ''' Siyan added END '''
          tic = time.time()
          for self.iter in range(iters):
              x_input_batch, y_input_batch = self.get_batch(x_train, y_train, batch_size)
              loss, nll_loss, reg_loss, y_hat, v, alpha, beta = self.run_train_step(x_input_batch, y_input_batch)
  
              # if self.iter % 10 == 0:
              #     # np.random.seed(1234)
              #     self.save_train_summary(loss, x_input_batch, y_input_batch, y_hat, v, alpha, beta)
              #     # self.save_train_summary_scalar(loss, x_input_batch, y_input_batch, y_hat, v, alpha, beta)
              #
              if self.iter % 10 == 0:
                  x_test_batch, y_test_batch = self.get_batch(x_test, y_test, min(100, x_test.shape[0]))
                  mu, v, alpha, beta, vloss, rmse, nll, reg_loss = self.evaluate(x_test_batch, y_test_batch)
  
                  nll += np.log(y_scale[0, 0])
                  rmse *= y_scale[0, 0]
  
                  ## validation summary
                  self.save_val_summary(vloss, x_test_batch, y_test_batch, mu, v, alpha, beta)
  
                  self.running_rmse = self.update_running(self.running_rmse, rmse.numpy())
                  if self.running_rmse < self.min_rmse:
                      self.min_rmse = self.running_rmse
                      # self.save(f"model_rmse_{self.iter}")
  
                  self.running_nll = self.update_running(self.running_nll, nll.numpy())
                  if self.running_nll < self.min_nll:
                      self.min_nll = self.running_nll
                      # self.save(f"model_nll_{self.iter}")
  
                  self.running_vloss = self.update_running(self.running_vloss, vloss.numpy())
                  if self.running_vloss < self.min_vloss:
                      self.min_vloss = self.running_vloss
                      # self.save(f"model_vloss_{self.iter}")
  
                  # if verbose: print(
                  #     "[A {}]  RMSE: {:.4f} \t NLL: {:.4f} \t loss: {:.4f} \t reg_loss: {:.4f} \t lambda: {:.2f} \t t: {:.2f} sec".format(
                  #         self.iter, self.min_rmse, self.min_nll, vloss, reg_loss.numpy().mean(), self.lam.numpy(),
                  #         time.time() - tic))
  
  
  
                  # ### Loss function for reference
                  # def loss_function(self, y, mu, v, alpha, beta, reduce=True, return_comps=False):
                  #     nll_loss = self.nll_loss_function(y, mu, v, alpha, beta, reduce=reduce)
                  #     reg_loss = self.reg_loss_function(y, mu, v, alpha, beta, reduce=reduce)
                  #     loss = nll_loss + self.lam * (reg_loss - self.epsilon)
                  #     return (loss, (nll_loss, reg_loss)) if return_comps else loss
  
  
              if self.iter % 10 == 0:
                  # self.save_train_summary(loss, x_input_batch, y_input_batch, y_hat, v, alpha, beta)
  
  
                  ''' Siyan Test START --- (for entire test data instead of batch of test data)'''
                  x_test_input = tf.convert_to_tensor(x_test, tf.float32)
                  y_test_input = tf.convert_to_tensor(y_test, tf.float32)
                  tmp_outputs_1 = self.model(x_test_input, training=False)
                  test_mu, test_v, test_alpha, test_beta = tf.split(tmp_outputs_1, 4, axis=1)
                  test_rmse = edl.losses.RMSE(y_test_input, test_mu)
  
                  ### Calculate loss for all training data instead of a batch of training data
                  test_loss, (test_nll_loss, test_reg_loss) = self.loss_function(y_test_input, test_mu, test_v, test_alpha, test_beta, return_comps=True)
                  # print(test_reg_loss)
                  if test_eval_count == 0:
                      tmp_test_loss = test_loss
                  else:
                      if test_loss < tmp_test_loss:
                          tmp_test_loss = test_loss
                          print("[{}] Test loss: {:.6f} \t RMSE: {:.4f} \t NLL: {:.4f} \t reg_loss: {:.4f} \t lambda: {:.2f}".
                                format(self.iter, test_loss, test_rmse.numpy(), test_nll_loss.numpy(), test_reg_loss.numpy(), self.lam.numpy()))
                          ### update the DataFrame
                          test_results_df = pd.DataFrame({
                              'test_x': list(x_test.flatten()),
                              'test_y': list(y_test.flatten()),
                              'test_mu': list(test_mu.numpy().flatten()),
                              'test_v': list(test_v.numpy().flatten()),
                              'test_alpha': list(test_alpha.numpy().flatten()),
                              'test_beta': list(test_beta.numpy().flatten()),
                          })
                          test_best_iter_str = str(self.iter)
                  test_eval_count += 1
                  ''' Siyan Test END --- (for entire test data instead of batch of test data)'''
  
                  ''' Siyan Test START --- (for entire validation data instead of batch of validation data)'''
                  tmp_outputs_2 = self.model(x_valid_input, training=False)
                  valid_mu, valid_v, valid_alpha, valid_beta = tf.split(tmp_outputs_2, 4, axis=1)
                  valid_rmse = edl.losses.RMSE(y_valid_input, valid_mu)
  
                  valid_loss, (valid_nll_loss, valid_reg_loss) = self.loss_function(y_valid_input, valid_mu, valid_v,
                                                                                    valid_alpha, valid_beta,
                                                                                    return_comps=True)
  
                  # print("[{}] Valid loss: {:.6f} \t RMSE: {:.4f} \t NLL: {:.4f} \t reg_loss: {:.4f} \t lambda: {:.2f}".
                  #       format(self.iter, valid_loss, valid_rmse.numpy(), valid_nll_loss.numpy(), valid_reg_loss.numpy(),
                  #              self.lam.numpy()))
  
                  if valid_eval_count == 0:
                      tmp_valid_loss = valid_loss
                      # tmp_valid_rmse = valid_rmse
                  else:
                      if valid_loss < tmp_valid_loss:
                      # if valid_rmse < tmp_valid_rmse:
                          tmp_valid_loss = valid_loss
                          # tmp_valid_rmse = valid_rmse
                          print("[{}] Valid loss: {:.6f} \t RMSE: {:.4f} \t NLL: {:.4f} \t reg_loss: {:.4f} \t lambda: {:.2f}".
                                format(self.iter, valid_loss, valid_rmse.numpy(), valid_nll_loss.numpy(), valid_reg_loss.numpy(), self.lam.numpy()))
                          ### update the DataFrame
                          valid_results_df = pd.DataFrame({
                              'valid_x': list(x_valid.flatten()),
                              'valid_y': list(y_valid.flatten()),
                              'valid_mu': list(valid_mu.numpy().flatten()),
                              'valid_v': list(valid_v.numpy().flatten()),
                              'valid_alpha': list(valid_alpha.numpy().flatten()),
                              'valid_beta': list(valid_beta.numpy().flatten()),
                          })
                          valid_best_iter_str = str(self.iter)
  
                  # valid_results_df = pd.DataFrame({
                  #     'valid_x': list(x_valid.flatten()),
                  #     'valid_y': list(y_valid.flatten()),
                  #     'valid_mu': list(valid_mu.numpy().flatten()),
                  #     'valid_v': list(valid_v.numpy().flatten()),
                  #     'valid_alpha': list(valid_alpha.numpy().flatten()),
                  #     'valid_beta': list(va.lid_beta.numpy().flatten()),
                  # })
                  valid_eval_count += 1
  
                  ''' Siyan Test END --- (for entire validation data instead of batch of validation data)'''
          test_results_df.to_csv(self.custom_best_results_dat_folder + "/best_test_results_iter_"+test_best_iter_str+".dat", sep=" ")
          print('--- Saved testing results to '+self.custom_best_results_dat_folder+'/best_test_results_iter_'+test_best_iter_str+".dat")
          valid_results_df.to_csv(self.custom_best_results_dat_folder + "/best_valid_results_iter_"+valid_best_iter_str+".dat", sep=" ")
          print('--- Saved validation results to '+self.custom_best_results_dat_folder +'/best_valid_results_iter_'+valid_best_iter_str+".dat")
  
          ''' Test/validation results calculation and plotting:
          (1) Best train results; (2) Best validation results. And compare to the original results
          '''
          load_test_df = pd.read_csv(self.custom_best_results_dat_folder + "/best_test_results_iter_"+test_best_iter_str+".dat", sep=" ")
          load_valid_df = pd.read_csv(self.custom_best_results_dat_folder + "/best_valid_results_iter_"+valid_best_iter_str+".dat", sep=" ")
  
          epistemic_test = np.sqrt(load_test_df['test_beta'].values / (load_test_df['test_v'].values * (load_test_df['test_alpha'].values - 1)))
          epistemic_test = np.minimum(epistemic_test, 1e3)  # clip the unc for vis
  
          print(load_test_df)
          epistemic_valid = np.sqrt(load_valid_df['valid_beta'].values / (load_valid_df['valid_v'].values * (load_valid_df['valid_alpha'].values - 1)))
          epistemic_valid = np.minimum(epistemic_valid, 1e3)  # clip the unc for vis
  
          valid_bounds = [[-7, +7]]
          self.plot_scatter_with_var_from_pandas(x_train, y_train, x_valid, y_valid,
                                                 load_test_df['test_mu'].values, epistemic_test, path=self.custom_plot_folder+"/test_plot.pdf", n_stds=3, test_bounds=valid_bounds, show=True)
          self.plot_scatter_with_var_from_pandas(x_train, y_train, x_valid, y_valid,
                                                 load_valid_df['valid_mu'].values, epistemic_valid, path=self.custom_plot_folder+"/valid_plot.pdf", n_stds=3, test_bounds=valid_bounds, show=True)
  
  
          return self.model, self.min_rmse, self.min_nll
  
  
      # def train(self, x_train, y_train, x_test, y_test, x_valid, y_valid, y_scale, batch_size=128, iters=10000, verbose=True):
      def train_v2(self, x_train, y_train, x_valid, y_valid, x_test, y_test, y_scale, batch_size=128, iters=10000, verbose=True):
          ''' Siyan added START '''
          # np.random.seed(1234)
          valid_eval_count = 0
          valid_best_iter_str = '0'
          test_eval_count = 0
          test_best_iter_str = '0'
  
          x_test_input = tf.convert_to_tensor(x_test, tf.float32)
          y_test_input = tf.convert_to_tensor(y_test, tf.float32)
  
          ''' Siyan added END '''
          tic = time.time()
          for self.iter in range(iters):
              x_input_batch, y_input_batch = self.get_batch(x_train, y_train, batch_size)
              loss, nll_loss, reg_loss, y_hat, v, alpha, beta = self.run_train_step(x_input_batch, y_input_batch)
  
              if self.iter % 10 == 0:
                  # np.random.seed(1234)
                  self.save_train_summary(loss, x_input_batch, y_input_batch, y_hat, v, alpha, beta)
                  # self.save_train_summary_scalar(loss, x_input_batch, y_input_batch, y_hat, v, alpha, beta)
  
              if self.iter % 10 == 0:
                  x_valid_batch, y_valid_batch = self.get_batch(x_valid, y_valid, min(100, x_valid.shape[0]))
                  mu, v, alpha, beta, vloss, rmse, nll, reg_loss = self.evaluate(x_valid_batch, y_valid_batch)
  
                  nll += np.log(y_scale[0, 0])
                  rmse *= y_scale[0, 0]
  
                  ## validation summary
                  self.save_val_summary(vloss, x_valid_batch, y_valid_batch, mu, v, alpha, beta)
  
                  self.running_rmse = self.update_running(self.running_rmse, rmse.numpy())
                  if self.running_rmse < self.min_rmse:
                      self.min_rmse = self.running_rmse
                      # self.save(f"model_rmse_{self.iter}")
  
                  self.running_nll = self.update_running(self.running_nll, nll.numpy())
                  if self.running_nll < self.min_nll:
                      self.min_nll = self.running_nll
                      # self.save(f"model_nll_{self.iter}")
  
                  self.running_vloss = self.update_running(self.running_vloss, vloss.numpy())
                  if self.running_vloss < self.min_vloss:
                      self.min_vloss = self.running_vloss
                      # self.save(f"model_vloss_{self.iter}")
  
                  # if verbose: print(
                  #     "[A {}]  RMSE: {:.4f} \t NLL: {:.4f} \t loss: {:.4f} \t reg_loss: {:.4f} \t lambda: {:.2f} \t t: {:.2f} sec".format(
                  #         self.iter, self.min_rmse, self.min_nll, vloss, reg_loss.numpy().mean(), self.lam.numpy(),
                  #         time.time() - tic))
  
                  # ### Loss function for reference
                  # def loss_function(self, y, mu, v, alpha, beta, reduce=True, return_comps=False):
                  #     nll_loss = self.nll_loss_function(y, mu, v, alpha, beta, reduce=reduce)
                  #     reg_loss = self.reg_loss_function(y, mu, v, alpha, beta, reduce=reduce)
                  #     loss = nll_loss + self.lam * (reg_loss - self.epsilon)
                  #     return (loss, (nll_loss, reg_loss)) if return_comps else loss
  
  
              if self.iter % 10 == 0 or self.iter < 50 or self.iter >= iters - 50:
                  # self.save_train_summary(loss, x_input_batch, y_input_batch, y_hat, v, alpha, beta)
  
                  ''' Siyan Test START --- (for all validation dat)'''
                  x_valid_input = tf.convert_to_tensor(x_valid, tf.float32)
                  y_valid_input = tf.convert_to_tensor(y_valid, tf.float32)
                  tmp_outputs_1 = self.model(x_valid_input, training=False)
                  valid_mu, valid_v, valid_alpha, valid_beta = tf.split(tmp_outputs_1, 4, axis=1)
                  valid_rmse = edl.losses.RMSE(y_valid_input, valid_mu)
  
                  ### Calculate loss for batch validation data
                  valid_loss, (valid_nll_loss, valid_reg_loss) = self.loss_function(y_valid_input, valid_mu, valid_v, valid_alpha, valid_beta, return_comps=True)
  
  
                  ### Calculate Epistemic, Aleatoric uncertainty for validation data directly and plot
                  ''' epistemic (Var[mu]) = beta/ (v(alpha - 1)) '''
                  epistemic_valid = np.sqrt(valid_beta / (valid_v * (valid_alpha - 1)))
                  aleatoric_sigma_valid = np.sqrt((epistemic_valid ** 2) * valid_v)
                  epistemic_valid = np.minimum(epistemic_valid, 1e3)  # clip the unc for vis
                  aleatoric_sigma_valid = np.minimum(aleatoric_sigma_valid, 1e3)  # clip the unc for vis
                  test_bounds = [[-7, +7]]
  
                  ### x_train, y_train: np_arr (size, 1)
                  ### x_valid, y_valid: np_arr (size, 1)
                  ### valid_mu: tf tensor (size, 1)
                  ### epistemic_valid: np_arr (size, 1)
                  valid_mu_np = valid_mu.numpy()
                  self.plot_scatter_with_var_from_pandas_v2(self.iter, x_train, y_train, x_valid, y_valid, valid_mu_np, epistemic_valid,
                                                            path=self.custom_plot_folder+"/valid_epistemic"+"/valid_plot_iter_{}".format(self.iter)+".png",
                                                            n_stds=3, test_bounds=test_bounds, evalType='valid', uq='epistemic', show=False, save=True)
                  self.plot_scatter_with_var_from_pandas_v2(self.iter, x_train, y_train, x_valid, y_valid, valid_mu_np, aleatoric_sigma_valid,
                                                            path=self.custom_plot_folder+"/valid_aleatoric"+"/valid_plot_iter_{}".format(self.iter)+".png",
                                                            n_stds=3, test_bounds=test_bounds, evalType='valid', uq='aleatoric', show=False, save=True)
  
                  # # print(valid_reg_loss)
                  if valid_eval_count == 0:
                      tmp_valid_loss = valid_loss
                  else:
                      if valid_loss < tmp_valid_loss:
                          tmp_valid_loss = valid_loss
                          print("[{}] Valid loss: {:.6f} \t RMSE: {:.4f} \t NLL: {:.4f} \t reg_loss: {:.4f} \t lambda: {:.2f}".
                                format(self.iter, valid_loss, valid_rmse.numpy(), valid_nll_loss.numpy(), valid_reg_loss.numpy(), self.lam.numpy()))
                          ### update the DataFrame
                          valid_results_df = pd.DataFrame({
                              'valid_x': list(x_valid.flatten()),
                              'valid_y': list(y_valid.flatten()),
                              'valid_mu': list(valid_mu.numpy().flatten()),
                              'valid_v': list(valid_v.numpy().flatten()),
                              'valid_alpha': list(valid_alpha.numpy().flatten()),
                              'valid_beta': list(valid_beta.numpy().flatten()),
                          })
                          valid_best_iter_str = str(self.iter)
                  valid_eval_count += 1
                  ''' Siyan Test END --- (for all validation data)'''
  
                  ''' Siyan Test START --- (for entire testing data)'''
                  tmp_outputs_2 = self.model(x_test_input, training=False)
                  test_mu, test_v, test_alpha, test_beta = tf.split(tmp_outputs_2, 4, axis=1)
                  test_rmse = edl.losses.RMSE(y_test_input, test_mu)
  
                  test_loss, (test_nll_loss, test_reg_loss) = self.loss_function(y_test_input, test_mu, test_v,
                                                                                    test_alpha, test_beta,
                                                                                    return_comps=True)
  
                  ### Calculate Epistemic, Aleatoric uncertainty for test data directly and plot
                  ''' epistemic (Var[mu]) = beta/ (v(alpha - 1)) '''
                  epistemic_test = np.sqrt(test_beta / (test_v * (test_alpha - 1)))
                  aleatoric_sigma_test = np.sqrt((epistemic_test ** 2) * test_v)
                  epistemic_test = np.minimum(epistemic_test, 1e3)  # clip the unc for vis
                  aleatoric_sigma_test = np.minimum(aleatoric_sigma_test, 1e3)  # clip the unc for vis
                  test_bounds = [[-7, +7]]
  
                  test_mu_np = test_mu.numpy()
                  self.plot_scatter_with_var_from_pandas_v2(self.iter, x_train, y_train, x_test, y_test, test_mu_np, epistemic_test,
                                                            path=self.custom_plot_folder+"/test_epistemic"+"/test_plot_iter_{}".format(self.iter)+".png",
                                                            n_stds=3, test_bounds=test_bounds, evalType='test', uq='epistemic', show=False, save=True)
                  self.plot_scatter_with_var_from_pandas_v2(self.iter, x_train, y_train, x_test, y_test, test_mu_np, aleatoric_sigma_test,
                                                            path=self.custom_plot_folder+"/test_aleatoric"+"/test_plot_iter_{}".format(self.iter)+".png",
                                                            n_stds=3, test_bounds=test_bounds, evalType='test', uq='aleatoric', show=False, save=True)
  
                  # print("[{}] Valid loss: {:.6f} \t RMSE: {:.4f} \t NLL: {:.4f} \t reg_loss: {:.4f} \t lambda: {:.2f}".
                  #       format(self.iter, valid_loss, valid_rmse.numpy(), valid_nll_loss.numpy(), valid_reg_loss.numpy(),
                  #              self.lam.numpy()))
  
                  # test_mu_np = test_mu.numpy()
                  # self.plot_scatter_with_var_from_pandas_v2(self.iter, x_train, y_train, x_test, y_test, test_mu_np, epistemic_test,
                  #                                           path=self.custom_plot_folder+"/test_plot_iter_{}".format(self.iter)+".png",
                  #                                           n_stds=3, test_bounds=test_bounds, evalType='test', uq='epistemic', show=False, save=True)
  
                  if test_eval_count == 0:
                      tmp_test_loss = test_loss
                      # tmp_valid_rmse = valid_rmse
                  else:
                      if test_loss < tmp_test_loss:
                      # if valid_rmse < tmp_valid_rmse:
                          tmp_test_loss = test_loss
                          # tmp_valid_rmse = valid_rmse
                          print("[{}] Test loss: {:.6f} \t RMSE: {:.4f} \t NLL: {:.4f} \t reg_loss: {:.4f} \t lambda: {:.2f}".
                                format(self.iter, test_loss, test_rmse.numpy(), test_nll_loss.numpy(), test_reg_loss.numpy(), self.lam.numpy()))
                          ### update the DataFrame
                          test_results_df = pd.DataFrame({
                              'test_x': list(x_test.flatten()),
                              'test_y': list(y_test.flatten()),
                              'test_mu': list(test_mu.numpy().flatten()),
                              'test_v': list(test_v.numpy().flatten()),
                              'test_alpha': list(test_alpha.numpy().flatten()),
                              'test_beta': list(test_beta.numpy().flatten()),
                          })
                          test_best_iter_str = str(self.iter)
                  test_eval_count += 1
  
                  ''' Siyan Test END --- (for entire test data)'''
  
          valid_results_df.to_csv(self.custom_best_results_dat_folder + "/best_valid_results_iter_"+valid_best_iter_str+".dat", sep=" ")
          print('--- Saved validation results to '+self.custom_best_results_dat_folder +'/best_valid_results_iter_'+valid_best_iter_str+".dat")
          test_results_df.to_csv(self.custom_best_results_dat_folder + "/best_test_results_iter_"+test_best_iter_str+".dat", sep=" ")
          print('--- Saved testing results to '+self.custom_best_results_dat_folder+'/best_test_results_iter_'+test_best_iter_str+".dat")
  
  
          ''' Validation/test results calculation and plotting:
          (1) Best validation results; (2) Best testing results. And compare to the original results
          '''
          load_valid_df = pd.read_csv(self.custom_best_results_dat_folder + "/best_valid_results_iter_"+valid_best_iter_str+".dat", sep=" ")
          load_test_df = pd.read_csv(self.custom_best_results_dat_folder + "/best_test_results_iter_"+test_best_iter_str+".dat", sep=" ")
  
          ''' epistemic (Var[mu]) = beta/ (v(alpha - 1)) '''
          epistemic_valid = np.sqrt(load_valid_df['valid_beta'].values / (load_valid_df['valid_v'].values * (load_valid_df['valid_alpha'].values - 1)))
          epistemic_valid = np.minimum(epistemic_valid, 1e3)  # clip the unc for vis
  
          epistemic_test = np.sqrt(load_test_df['test_beta'].values / (load_test_df['test_v'].values * (load_test_df['test_alpha'].values - 1)))
          epistemic_test = np.minimum(epistemic_test, 1e3)  # clip the unc for vis
  
          ''' aleatoric (E[sigma^2]) = beta / (alpha - 1)'''
          # aleatoric_valid_sigma = np.sqrt()
  
          # print(load_test_df)
  
          # valid_bounds = [[-4, +4]]
          test_bounds = [[-7, +7]]
  
          # print(x_train.shape)
          # print(y_train.shape)
          # print(x_valid.shape)
          # print(y_valid.shape)
          # print(load_valid_df['valid_mu'].values.shape)
          # print(epistemic_valid.shape)
  
  
          # self.plot_scatter_with_var_from_pandas_v2(x_train, y_train, x_valid, y_valid,
          #                                        load_valid_df['valid_mu'].values, epistemic_valid, path=self.custom_plot_folder+"/valid_plot.pdf", n_stds=3, test_bounds=test_bounds, show=True)
  
          #
          # self.plot_scatter_with_var_from_pandas_v2(x_train, y_train, x_test, y_test,
          #                                        load_test_df['test_mu'].values, epistemic_test, path=self.custom_plot_folder+"/test_plot.pdf", n_stds=3, test_bounds=test_bounds, show=True)
  
          return self.model, self.min_rmse, self.min_nll
   import numpy as np
  import tensorflow as tf
  import tensorflow_probability as tfp
  import time
  import datetime
  import os
  import sys
  import h5py
  from pathlib import Path
  import pandas as pd
  import matplotlib.pyplot as plt
  
  import evidential_deep_learning as edl
  from .util import normalize, gallery
  
  class BBBP:
      def __init__(self, model, opts, dataset="", learning_rate=1e-3, tag="", custom_plot_folder="", custom_best_results_dat_folder=""):
          self.loss_function = edl.losses.MSE
  
          self.model = model
  
          self.optimizer = tf.optimizers.Adam(learning_rate)
  
          self.min_rmse = float('inf')
          self.min_nll = float('inf')
          self.min_vloss = float('inf')
  
          trainer = self.__class__.__name__
          current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
          self.save_dir = os.path.join('save','{}_{}_{}_{}'.format(current_time, dataset, trainer, tag))
          Path(self.save_dir).mkdir(parents=True, exist_ok=True)
  
          self.custom_plot_folder = custom_plot_folder
          Path(self.custom_plot_folder).mkdir(parents=True, exist_ok=True)
  
          self.custom_best_results_dat_folder = custom_best_results_dat_folder
          Path(self.custom_best_results_dat_folder).mkdir(parents=True, exist_ok=True)
  
  
          train_log_dir = os.path.join('logs', '{}_{}_{}_{}_train'.format(current_time, dataset, trainer, tag))
          self.train_summary_writer = tf.summary.create_file_writer(train_log_dir)
          val_log_dir = os.path.join('logs', '{}_{}_{}_{}_val'.format(current_time, dataset, trainer, tag))
          self.val_summary_writer = tf.summary.create_file_writer(val_log_dir)
  
      @tf.function
      def run_train_step(self, x, y):
          with tf.GradientTape() as tape:
              y_hat = self.model(x, training=True) #forward pass
              loss = self.loss_function(y, y_hat)
              loss += tf.reduce_mean(self.model.losses)
  
          grads = tape.gradient(loss, self.model.variables) #compute gradient
          self.optimizer.apply_gradients(zip(grads, self.model.variables))
          return loss, y_hat
  
      @tf.function
      def evaluate(self, x, y):
          preds = tf.stack([self.model(x, training=True) for _ in range(5)], axis=0) #forward pass
          mu, var = tf.nn.moments(preds, axes=0)
          rmse = edl.losses.RMSE(y, mu)
          nll = edl.losses.Gaussian_NLL(y, mu, tf.sqrt(var))
          loss = self.loss_function(y, mu)
  
          return mu, var, loss, rmse, nll
  
      @tf.function
      def siyan_bbbp_evaluate(self, x_input):
          preds = tf.stack([self.model(x_input, training=True) for _ in range(15)], axis=0)  # forward pass
          mean_mu = tf.reduce_mean(preds, axis=0)
          epistemic = tf.math.reduce_std(preds, axis=0)
          return mean_mu, epistemic
  
      def save_train_summary(self, loss, x, y, y_hat):
          with self.train_summary_writer.as_default():
              # tf.summary.scalar('loss', tf.reduce_mean(loss), step=self.iter)
              tf.summary.scalar('mse', tf.reduce_mean(edl.losses.MSE(y, y_hat)), step=self.iter)
              idx = np.random.choice(int(tf.shape(x)[0]), 9)
              if tf.shape(x).shape==4:
                  tf.summary.image("x", [gallery(tf.gather(x,idx).numpy())], max_outputs=1, step=self.iter)
  
              if tf.shape(y).shape==4:
                  tf.summary.image("y", [gallery(tf.gather(y,idx).numpy())], max_outputs=1, step=self.iter)
                  tf.summary.image("y_hat", [gallery(tf.gather(y_hat,idx).numpy())], max_outputs=1, step=self.iter)
  
      def save_val_summary(self, loss, x, y, mu, var):
          with self.val_summary_writer.as_default():
              tf.summary.scalar('loss', tf.reduce_mean(self.loss_function(y, mu)), step=self.iter)
              tf.summary.scalar('mse', tf.reduce_mean(edl.losses.MSE(y, mu)), step=self.iter)
              idx = np.random.choice(int(tf.shape(x)[0]), 9)
              if tf.shape(x).shape==4:
                  tf.summary.image("x", [gallery(tf.gather(x,idx).numpy())], max_outputs=1, step=self.iter)
  
              if tf.shape(y).shape==4:
                  tf.summary.image("y", [gallery(tf.gather(y,idx).numpy())], max_outputs=1, step=self.iter)
                  tf.summary.image("y_hat", [gallery(tf.gather(mu,idx).numpy())], max_outputs=1, step=self.iter)
                  tf.summary.image("y_var", [gallery(normalize(tf.gather(var,idx)).numpy())], max_outputs=1, step=self.iter)
  
      def get_batch(self, x, y, batch_size):
          idx = np.random.choice(x.shape[0], batch_size, replace=False)
          if isinstance(x, tf.Tensor):
              x_ = x[idx,...]
              y_ = y[idx,...]
          elif isinstance(x, np.ndarray) or isinstance(x, h5py.Dataset):
              idx = np.sort(idx)
              x_ = x[idx,...]
              y_ = y[idx,...]
  
              x_divisor = 255. if x_.dtype == np.uint8 else 1.0
              y_divisor = 255. if y_.dtype == np.uint8 else 1.0
  
              x_ = tf.convert_to_tensor(x_/x_divisor, tf.float32)
              y_ = tf.convert_to_tensor(y_/y_divisor, tf.float32)
          else:
              print("unknown dataset type {} {}".format(type(x), type(y)))
          return x_, y_
  
      def save(self, name):
          self.model.save(os.path.join(self.save_dir, "{}.h5".format(name)))
          # pass
  
      ''' Siyan added for testing plot '''
      def plot_scatter_with_var_from_pandas(self, x_train, y_train, x_test, y_test, mu, var, path, n_stds=3, test_bounds=[[-7, +7]], show=True):
          plt.scatter(x_train, y_train, s=1., c='#463c3c', zorder=0, label='Train (x_train vs y_train)')
          for k in np.linspace(0, n_stds, 4):
              if k == 0:
                  plt.fill_between(x_test[:, 0], (mu - k * var), (mu + k * var), alpha=0.3, edgecolor=None,
                                   facecolor='#00aeef', linewidth=0, antialiased=True, zorder=1, label='Unc.')
              else:
                  plt.fill_between(x_test[:, 0], (mu - k * var), (mu + k * var), alpha=0.3, edgecolor=None,
                                   facecolor='#00aeef', linewidth=0, antialiased=True, zorder=1)
  
          plt.plot(x_test, y_test, 'r--', zorder=2, label='True (x_test vs y_test)')
          plt.plot(x_test, mu, color='#007cab', zorder=3, label='Pred (x_test vs mu)')
          plt.gca().set_xlim(*test_bounds)
          plt.gca().set_ylim(-150, 150)
          plt.title(path)
          plt.legend()
          plt.savefig(path, transparent=True)
          if show:
              plt.show()
          plt.clf()
  
      def train(self, x_train, y_train, x_test, y_test, x_valid, y_valid, y_scale, batch_size=128, iters=10000, verbose=True):
          ''' Siyan added START '''
          # np.random.seed(1234)
          test_eval_count = 0
          valid_eval_count = 0
          test_best_iter_str = '0'
          valid_best_iter_str = '0'
  
          x_valid_input = tf.convert_to_tensor(x_valid, tf.float32)
          y_valid_input = tf.convert_to_tensor(y_valid, tf.float32)
  
          tic = time.time()
          for self.iter in range(iters):
              x_input_batch, y_input_batch = self.get_batch(x_train, y_train, batch_size)
              loss, y_hat = self.run_train_step(x_input_batch, y_input_batch)
  
              if self.iter % 10 == 0:
                  self.save_train_summary(loss, x_input_batch, y_input_batch, y_hat)
  
              if self.iter % 100 == 0:
                  x_test_batch, y_test_batch = self.get_batch(x_test, y_test, min(100, x_test.shape[0]))
                  mu, var, vloss, rmse, nll = self.evaluate(x_test_batch, y_test_batch)
                  nll += np.log(y_scale[0,0])
                  rmse *= y_scale[0,0]
  
                  self.save_val_summary(vloss, x_test_batch, y_test_batch, mu, var)
  
                  if rmse.numpy() < self.min_rmse:
                      self.min_rmse = rmse.numpy()
                      print("SAVING")
                      self.save("model_rmse")
  
                  if nll.numpy() < self.min_nll:
                      self.min_nll = nll.numpy()
                      self.save("model_nll")
  
                  if vloss.numpy() < self.min_vloss:
                      self.min_vloss = vloss.numpy()
                      self.save("model_vloss")
  
                  # if verbose: print("[{}] \t RMSE: {:.4f} \t NLL: {:.4f} \t train_loss: {:.4f} \t t: {:.2f} sec".format(self.iter, self.min_rmse, self.min_nll, vloss, time.time()-tic))
                  # tic = time.time()
  
              ''' Siyan Test START --- (for entire test data instead of batch of test data)'''
              if self.iter % 10 == 0:
                  x_test_input = tf.convert_to_tensor(x_test, tf.float32)
                  y_test_input = tf.convert_to_tensor(y_test, tf.float32)
                  test_mu, test_var, test_vloss, test_rmse, test_nll = self.evaluate(x_test_input, y_test_input)
                  if test_eval_count == 0:
                      tmp_test_loss = test_vloss
                  else:
                      if test_vloss < tmp_test_loss:
                          tmp_test_loss = test_vloss
                          print("[{}] Test loss: {:.6f} \t RMSE: {:.4f} \t NLL: {:.4f} ".
                                format(self.iter, test_vloss, test_rmse.numpy(), test_nll.numpy()))
  
                          test_mean_mu, test_epistemic = self.siyan_bbbp_evaluate(x_test_input)
                          ### update the DataFrame
                          test_results_df = pd.DataFrame({
                              'test_x': list(x_test.flatten()),
                              'test_y': list(y_test.flatten()),
                              'test_mean_mu': list(test_mean_mu.numpy().flatten()),
                              'test_epistemic': list(test_epistemic.numpy().flatten())
                              # 'test_mu': list(test_mu.numpy().flatten()),
                              # 'test_var': list(test_var.numpy().flatten())
                          })
                          test_best_iter_str = str(self.iter)
                  test_eval_count += 1
                  ''' Siyan Test END --- (for entire test data instead of batch of test data)'''
  
                  ''' Siyan Test START --- (for entire validation data instead of batch of validation data)'''
                  valid_mu, valid_var, valid_vloss, valid_rmse, valid_nll = self.evaluate(x_valid_input, y_valid_input)
                  if valid_eval_count == 0:
                      tmp_valid_loss = valid_vloss
                  else:
                      if valid_vloss < tmp_valid_loss:
                          tmp_valid_loss = valid_vloss
                          print("[{}] Validation loss: {:.6f} \t RMSE: {:.4f} \t NLL: {:.4f} ".
                                format(self.iter, valid_vloss, valid_rmse.numpy(), valid_nll.numpy()))
  
                          valid_mean_mu, valid_epistemic = self.siyan_bbbp_evaluate(x_valid_input)
                          ### update the DataFrame
                          valid_results_df = pd.DataFrame({
                              'valid_x': list(x_valid.flatten()),
                              'valid_y': list(y_valid.flatten()),
                              'valid_mean_mu': list(valid_mean_mu.numpy().flatten()),
                              'valid_epistemic': list(valid_epistemic.numpy().flatten())
                              # 'valid_mean_var': list(valid_mean_var.numpy().flatten()),
                              # 'valid_reduce_std_mu': list(valid_reduce_std_mu.numpy().flatten()),
                              # 'valid_reduce_mean_var': list(valid_reduce_mean_var.numpy().flatten())
                          })
                          valid_best_iter_str = str(self.iter)
  
                  valid_eval_count += 1
                  ''' Siyan Test END --- (for entire validation data instead of batch of validation data)'''
  
          test_results_df.to_csv(self.custom_best_results_dat_folder + "/best_test_results_iter_"+test_best_iter_str+".dat", sep=" ")
          print('--- Saved testing results to '+self.custom_best_results_dat_folder+'/best_test_results_iter_'+test_best_iter_str+".dat")
          valid_results_df.to_csv(self.custom_best_results_dat_folder + "/best_valid_results_iter_"+valid_best_iter_str+".dat", sep=" ")
          print('--- Saved validation results to '+self.custom_best_results_dat_folder +'/best_valid_results_iter_'+valid_best_iter_str+".dat")
  
          ''' Test/validation results calculation and plotting:
          (1) Best train results; (2) Best validation results. And compare to the original results
          '''
          load_test_df = pd.read_csv(self.custom_best_results_dat_folder + "/best_test_results_iter_"+test_best_iter_str+".dat", sep=" ")
          load_valid_df = pd.read_csv(self.custom_best_results_dat_folder + "/best_valid_results_iter_"+valid_best_iter_str+".dat", sep=" ")
  
          valid_bounds = [[-7, +7]]
          self.plot_scatter_with_var_from_pandas(x_train, y_train, x_valid, y_valid,
                                                 load_test_df['test_mean_mu'].values, load_test_df['test_epistemic'].values, path=self.custom_plot_folder+"/test_plot.pdf", n_stds=3, test_bounds=valid_bounds, show=True)
          self.plot_scatter_with_var_from_pandas(x_train, y_train, x_valid, y_valid,
                                                 load_valid_df['valid_mean_mu'].values, load_valid_df['valid_epistemic'], path=self.custom_plot_folder+"/valid_plot.pdf", n_stds=3, test_bounds=valid_bounds, show=True)
  
  
  
          return self.model, self.min_rmse, self.min_nll
   # -*- coding: utf-8 -*-
  """
  Run method, and save results.
  Run as:
      python main.py --dataset <ds> --method <met>
      where dataset name should be in UCI_Datasets folder
      and method is piven, qd, deep-ens, mid or only-rmse.
  """
  import argparse
  import json
  import datetime
  import tensorflow as tf
  # import tensorflow.compat.v1 as tf
  import scipy.stats as stats
  import itertools
  import os
  import random
  import numpy as np
  
  
  import data_loader
  from DataGen import DataGenerator
  from DeepNetPI import TfNetwork
  from utils import *
  from sklearn.model_selection import train_test_split
  
  
  start_time = datetime.datetime.now()
  
  parser = argparse.ArgumentParser()
  parser.add_argument('--dataset', type=str, default='flight_delay', metavar='',
                      help='dataset name, flight_delay')
  parser.add_argument('--method', type=str, help='piven, qd, mid, only-rmse, deep-ens', required=True)
  args = parser.parse_args()
  method = args.method
  
  #############  added code start  ############# 
  print(args)
  print(args.dataset)
  original_data_path = '../flight_delay_data/' ## flight delay data
  results_path = './Results/piven/'+args.dataset + '_PIVEN_UCI.txt'
  
  
  seed = 12345
  neurons = [100]
  lambda_in = 15.0
  sigma_in = 0.2
  
  random.seed(seed)
  np.random.seed(seed)
  tf.compat.v1.random.set_random_seed(seed)
  # tf.random.set_random_seed(seed)
  h_size = neurons
  
  
  # n_runs = params['n_runs']  # number of runs
  n_runs = 1
  n_epoch = 500  # number epochs to train for
  # h_size = params['h_size']  # number of hidden units in network: [50]=layer_1 of 50, [8,4]=layer_1 of 8, layer_2 of 4
  l_rate = 0.01  # learning rate of optimizer
  decay_rate = 0.99  # learning rate decay
  soften = 160.0 # soften param in the loss
  patience = -1  # patience
  n_ensemble = 1  # number of individual NNs in ensemble  # 5
  alpha = 0.05  # data points captured = (1 - alpha)
  train_prop = 0.9  # % of data to use as training
  in_ddof = 1 if n_runs > 1 else 0  # this is for results over runs only
  is_early_stop = patience != -1
  
  if args.dataset == 'YearPredictionMSD':
      n_batch = 1000  # batch size
      out_biases = [5., -5.]
  else:
      n_batch = 100  # batch size
      out_biases = [2., -2.]  # chose biases for output layer (for deep_ens is overwritten to 0,1)
  
  
  results_runs = []
  run = 0
  fail_times = 0
  for run in range(0, n_runs):
  
  
      ''' ######### flight delay data loading ###############'''
      xTrain, yTrain, test_data_list  = data_loader.load_flight_delays(original_data_path)
  
      # y_train = np.reshape(y_train, (-1, 1))
      # y_test = np.reshape(y_test, (-1, 1))
  
      '''choose the train/test dataset '''
      x_train = xTrain
      y_train = yTrain
      y_train = y_train.reshape(-1, 1)
      # y_scale = yTrain_scale
      test_idx = 0  # [0, 1, 2, 3] for test 1,2,3,4
      X_test = test_data_list[test_idx][0]
      y_test = test_data_list[test_idx][1]
      y_test = y_test.reshape(-1, 1)
  
  
      X_val = X_test
      y_val = y_test.reshape(-1, 1)
  
  
      y_pred_all = []
      y_pred_all_train = []
  
      i = 0
      while i < n_ensemble:
          is_failed_run = False
  
          tf.reset_default_graph()
          sess = tf.Session()
  
          print(f'\nrun number {run+1} of {n_runs} -- ensemble number {i+1} of {n_ensemble}')
  
          # create network
          NN = TfNetwork(x_size=x_train.shape[1],
                         y_size=2,
                         h_size=h_size,
                         alpha=alpha,
                         soften=soften,
                         lambda_in=lambda_in,
                         sigma_in=sigma_in,
                         out_biases=out_biases,
                         method=method,
                         patience=patience,
                         dataset=args.dataset,
                         rnd_seed=seed)
  
          # train
          NN.train(sess, x_train, y_train, X_val, y_val,
                   n_epoch=n_epoch,
                   l_rate=l_rate,
                   decay_rate=decay_rate,
                   is_early_stop=is_early_stop,
                   n_batch=n_batch)
  
          # predict
          y_loss, y_pred = NN.predict(sess, X_test=X_test, y_test=y_test)
  
          # prediction for training data
          y_loss_train, y_pred_train = NN.predict(sess, X_test=x_train, y_test=y_train)
  
          # check whether the run failed or not
          if np.abs(y_loss) > 20. and fail_times < 1: # jump out of some endless failures
              # if False:
              is_failed_run = True
              fail_times+=1
              print('\n\n### one messed up! repeating ensemble ### failed {}/5 times!'.format(fail_times))
              continue  # without saving result
          else:
              i += 1  # continue to next
  
          # save prediction
          y_pred_all.append(y_pred)
          y_pred_all_train.append(y_pred_train)
          sess.close()
  
      y_pred_all = np.array(y_pred_all)
      y_pred_all_train = np.array(y_pred_all_train)
  
      if method == 'deep-ens':
          y_pred_gauss_mid_all = y_pred_all[:, :, 0]
          # occasionally may get -ves for std dev so need to do max
          y_pred_gauss_dev_all = np.sqrt(np.maximum(np.log(1. + np.exp(y_pred_all[:, :, 1])), 10e-6))
          y_pred_gauss_mid, y_pred_gauss_dev, y_pred_U, \
          y_pred_L = gauss_to_pi(y_pred_gauss_mid_all, y_pred_gauss_dev_all)
  
      else:
          ### for test data
          y_pred_gauss_mid, y_pred_gauss_dev, y_pred_U, y_pred_L, y_pred_v = pi_to_gauss(y_pred_all, method=method)
          ### for training data
          y_pred_gauss_mid_train, y_pred_gauss_dev_train, y_pred_U_train, y_pred_L_train, y_pred_v_train = pi_to_gauss(y_pred_all_train, method=method)
      
  
      ### calculate the confidence scores 
      # for train
      y_U_cap_train = y_pred_U_train > y_train.reshape(-1)
      y_L_cap_train = y_pred_L_train < y_train.reshape(-1)
      MPIW_array_train = y_pred_U_train - y_pred_L_train      
      MPIW_train = np.mean(MPIW_array_train)
  
      MPIW_array_test = y_pred_U - y_pred_L
  
      confidence_arr_test = [min(MPIW_train/test_width, 1.0) for test_width in MPIW_array_test]
      confidence_arr_train = [min(MPIW_train/train_width, 1.0) for train_width in MPIW_array_train]
  
      print('----------- OOD analysis --- confidence scores ----------------')
      print('--- Train conf_scores MEAN: {}, STD: {}'.format(np.mean(confidence_arr_train), np.std(confidence_arr_train)))
      print('--- Test: {} rank: {} conf_scores MEAN: {}, STD: {}'.format(test_idx+1, test_idx+1, np.mean(confidence_arr_test), np.std(confidence_arr_test)))
  
      dist_arr_train = np.sqrt(np.sum(x_train ** 2.0, axis=1))
      dist_arr_test = np.sqrt(np.sum(X_val ** 2.0, axis=1))
  
      confidence_arr_train = np.array(confidence_arr_train)
      confidence_arr_test = np.array(confidence_arr_test)
  
      PIVEN_OOD_train_np = np.hstack((dist_arr_train.reshape(-1, 1), confidence_arr_train.reshape(-1, 1)))
      PIVEN_OOD_test_np = np.hstack((dist_arr_test.reshape(-1, 1), confidence_arr_test.reshape(-1, 1)))
  
      np.savetxt('PIVEN_OOD_flight_delay_'+ str(test_idx+1) +'_train_np.txt', PIVEN_OOD_train_np, delimiter=',')
      np.savetxt('PIVEN_OOD_flight_delay_'+ str(test_idx+1) +'_test_np.txt', PIVEN_OOD_test_np, delimiter=',')
  
      # # work out metrics
      # y_U_cap = y_pred_U > y_test.reshape(-1)
      # y_L_cap = y_pred_L < y_test.reshape(-1)
  
      # y_all_cap = y_U_cap * y_L_cap
      # PICP = np.sum(y_all_cap) / y_L_cap.shape[0]
      # MPIW = np.mean(y_pred_U - y_pred_L)
      # y_pred_mid = np.mean((y_pred_U, y_pred_L), axis=0)
      # # MSE = np.mean(np.square(Gen.scale_c * (y_pred_mid - y_test[:, 0])))
      # # RMSE = np.sqrt(MSE)
  
      # if method == 'qd' or method == 'deep-ens':
      #     RMSE_ELI = 0.0  # RMSE_PIVEN
      # else:
      #     if method == 'piven':
      #         y_piven = y_pred_v * y_pred_U + (1 - y_pred_v) * y_pred_L
  
      #     elif method == 'mid':
      #         y_piven = 0.5 * y_pred_U + 0.5 * y_pred_L
  
      #     elif method == 'only-rmse':
      #         y_piven = y_pred_v
  
      #     MSE_ELI = np.mean(np.square(Gen.scale_c * (y_piven - y_test[:, 0])))
      #     RMSE_ELI = np.sqrt(MSE_ELI) # RMSE_PIVEN
  
      # CWC = np_QD_loss(y_test, y_pred_L, y_pred_U, alpha, lambda_in)  # from qd paper.
      # neg_log_like = gauss_neg_log_like(y_test, y_pred_gauss_mid, y_pred_gauss_dev, Gen.scale_c)
      # residuals = y_pred_mid - y_test[:, 0]
      # shapiro_W, shapiro_p = stats.shapiro(residuals[:])
      # results_runs.append((PICP, MPIW, CWC, RMSE, RMSE_ELI, neg_log_like, shapiro_W, shapiro_p))
  
  # # summarize results
  # results_path = f"./Results/{method}/"
  # results_path += f"{params['dataset']}-{start_time.strftime('%d-%m-%H-%M')}-{method}.csv"
  
  # results = np.array(results_runs)
  # results_to_csv(results_path, results, params, n_runs, n_ensemble, in_ddof)
  
  # timing info
  end_time = datetime.datetime.now()
  total_time = end_time - start_time
  print('\n\nminutes taken:', round(total_time.total_seconds() / 60, 3),
        '\nstart_time:', start_time.strftime('%H:%M:%S'),
        'end_time:', end_time.strftime('%H:%M:%S'))
  
apis: tensorflow.reduce_max numpy.log tensorflow.convert_to_tensor tensorflow.Variable matplotlib.pyplot.plot tensorflow.split matplotlib.pyplot.savefig matplotlib.pyplot.gca numpy.random.choice matplotlib.pyplot.title tensorflow.GradientTape tensorflow.reduce_min numpy.linspace matplotlib.pyplot.scatter numpy.minimum tensorflow.shape pandas.read_csv matplotlib.pyplot.clf numpy.sort matplotlib.pyplot.fill_between matplotlib.pyplot.legend matplotlib.pyplot.show numpy.sqrt tensorflow.gather tensorflow.optimizers.Adam tensorflow.summary.create_file_writer
apis: tensorflow.math.reduce_std numpy.log tensorflow.convert_to_tensor matplotlib.pyplot.plot matplotlib.pyplot.savefig matplotlib.pyplot.gca numpy.random.choice matplotlib.pyplot.title tensorflow.GradientTape numpy.linspace matplotlib.pyplot.scatter tensorflow.shape pandas.read_csv matplotlib.pyplot.clf numpy.sort matplotlib.pyplot.fill_between tensorflow.nn.moments matplotlib.pyplot.legend tensorflow.sqrt tensorflow.reduce_mean matplotlib.pyplot.show tensorflow.gather tensorflow.optimizers.Adam tensorflow.summary.create_file_writer
apis: numpy.sum numpy.random.seed numpy.abs numpy.exp tensorflow.Session numpy.array numpy.std tensorflow.reset_default_graph tensorflow.compat.v1.random.set_random_seed numpy.mean