def training(hyperp, options, filepaths, data_dict, prior_dict): #=== GPU Settings ===# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" if options.distributed_training == 0: os.environ["CUDA_VISIBLE_DEVICES"] = options.which_gpu if options.distributed_training == 1: os.environ["CUDA_VISIBLE_DEVICES"] = options.dist_which_gpus gpus = tf.config.experimental.list_physical_devices('GPU') #=== Construct Validation Set and Batches ===# input_and_latent_train, input_and_latent_val, input_and_latent_test,\ num_batches_train, num_batches_val, num_batches_test\ = form_train_val_test_tf_batches( data_dict["state_obs_train"], data_dict["parameter_train"], data_dict["state_obs_test"], data_dict["parameter_test"], hyperp.batch_size, options.random_seed) #=== Data and Latent Dimensions of Autoencoder ===# input_dimensions = data_dict["obs_dimensions"] latent_dimensions = options.parameter_dimensions #=== Neural Network Regularizers ===# kernel_initializer = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.05) bias_initializer = 'zeros' kernel_initializer_iaf = tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.05) bias_initializer_iaf = 'zeros' #=== Non-distributed Training ===# if options.distributed_training == 0: #=== Neural Network ===# nn = VAEIAF(hyperp, options, input_dimensions, latent_dimensions, kernel_initializer, bias_initializer, kernel_initializer_iaf, bias_initializer_iaf, positivity_constraint_log_exp) #=== Optimizer ===# optimizer = tf.keras.optimizers.Adam() #=== Training ===# optimize(hyperp, options, filepaths, nn, optimizer, input_and_latent_train, input_and_latent_val, input_and_latent_test, input_dimensions, latent_dimensions, num_batches_train, data_dict["noise_regularization_matrix"], prior_dict["prior_mean"], prior_dict["prior_covariance_cholesky_inverse"]) #=== Distributed Training ===# if options.distributed_training == 1: dist_strategy = tf.distribute.MirroredStrategy() with dist_strategy.scope(): #=== Neural Network ===# nn = VAEIAF(hyperp, options, input_dimensions, latent_dimensions, kernel_initializer, bias_initializer, positivity_constraint_log_exp) #=== Optimizer ===# optimizer = tf.keras.optimizers.Adam() #=== Training ===# optimize_distributed(dist_strategy, hyperp, options, filepaths, nn, optimizer, input_and_latent_train, input_and_latent_val, input_and_latent_test, input_dimensions, latent_dimensions, num_batches_train, data_dict["noise_regularization_matrix"], prior_dict["prior_mean"], prior_dict["prior_covariance_cholesky_inverse"])
def predict_and_plot(hyperp, options, filepaths): #=== Load Observation Indices ===# if options.obs_type == 'full': obs_dimensions = options.parameter_dimensions if options.obs_type == 'obs': obs_dimensions = options.num_obs_points print('Loading Boundary Indices') df_obs_indices = pd.read_csv(filepaths.project.obs_indices + '.csv') obs_indices = df_obs_indices.to_numpy() #=== Data and Latent Dimensions of Autoencoder ===# input_dimensions = obs_dimensions latent_dimensions = options.parameter_dimensions #=== Prepare Data ===# data = DataHandler(hyperp, options, filepaths, options.parameter_dimensions, obs_dimensions) data.load_data_test() if options.add_noise == 1: data.add_noise_qoi_test() parameter_test = data.poi_test state_obs_test = data.qoi_test #=== Load Trained Neural Network ===# nn = VAEIAF(hyperp, options, input_dimensions, latent_dimensions, None, None, None, None, positivity_constraint_log_exp) nn.load_weights(filepaths.trained_nn) #=== Selecting Samples ===# sample_number = 105 parameter_test_sample = np.expand_dims(parameter_test[sample_number, :], 0) state_obs_test_sample = np.expand_dims(state_obs_test[sample_number, :], 0) #=== Predictions ===# parameter_pred_sample, _ = nn.iaf_chain_posterior( nn.encoder(state_obs_test_sample)) state_obs_pred_sample = nn.decoder(parameter_test_sample) parameter_pred_sample = parameter_pred_sample.numpy().flatten() state_obs_pred_sample = state_obs_pred_sample.numpy().flatten() #=== Plotting Prediction ===# print('================================') print(' Plotting Predictions ') print('================================') #=== Load Mesh ===# nodes, elements, _, _, _, _, _, _ = load_mesh(filepaths.project) #=== Plot FEM Functions ===# plot_fem_function(filepaths.figures_savefile_name_parameter_test, 'True Parameter', 7.0, nodes, elements, parameter_test_sample) plot_fem_function(filepaths.figures_savefile_name_parameter_pred, 'Parameter Prediction', 7.0, nodes, elements, parameter_pred_sample) if options.obs_type == 'full': plot_fem_function(filepaths.figures_savefile_name_state_test, 'True State', 2.6, nodes, elements, state_obs_test_sample) plot_fem_function(filepaths.figures_savefile_name_state_pred, 'State Prediction', 2.6, nodes, elements, state_obs_pred_sample) print('Predictions plotted')