def load_seq3(batch_size=100,gnn=False, subsize=None, lognorm=False): train_coeffs = np.load(ROOT_DATA_DIR+"seq_3_train_coeffs.npy") train_adjs = np.load(ROOT_DATA_DIR+"seq_3_train_adjs.npy") train_energies = np.load(ROOT_DATA_DIR+"seq_3_train_energies.npy") test_coeffs = np.load(ROOT_DATA_DIR+"seq_3_test_coeffs.npy") test_adjs = np.load(ROOT_DATA_DIR+"seq_3_test_adjs.npy") test_energies = np.load(ROOT_DATA_DIR+"seq_3_test_energies.npy") if lognorm: # shift train_coeffs += np.abs(train_coeffs.min()) + 1 test_coeffs += np.abs(train_coeffs.min()) + 1 # log train_coeffs = np.log(train_coeffs) test_coeffs = np.log(test_coeffs) if gnn: train_diracs = torch.eye(train_adjs.shape[-1]).unsqueeze(0).repeat(train_adjs.shape[0],1,1) train_tup = (torch.Tensor(train_diracs), torch.Tensor(train_adjs), torch.Tensor(train_energies)) else: train_tup = (torch.Tensor(train_coeffs), torch.Tensor(train_energies)) if gnn: test_diracs = torch.eye(test_adjs.shape[-1]).unsqueeze(0).repeat(test_adjs.shape[0],1,1) test_tup = (torch.Tensor(test_diracs), torch.Tensor(test_adjs), torch.Tensor(test_energies)) else: test_tup = (torch.Tensor(test_coeffs), torch.Tensor(test_adjs), torch.Tensor(test_energies)) ################# # SUBSET DATA #################tre if subsize != None: train_tup, _ = eval_metrics.compute_subsample(train_tup, subsize) test_tup, _ = eval_metrics.compute_subsample(test_tup, subsize) train_dataset = torch.utils.data.TensorDataset(*train_tup) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True) return train_loader, train_tup, test_tup
def load_seq4(batch_size=100, gnn=False, subsize=None): train_coeffs = np.load(SEQ4_DATA_DIR + "train_normcoeffs_0523.npy") train_adjs = np.load(SEQ4_DATA_DIR + "train_adjs_0523.npy") train_energies = np.load(SEQ4_DATA_DIR + "train_energies_0523.npy") test_coeffs = np.load(SEQ4_DATA_DIR + "test_normcoeffs_0523.npy") test_adjs = np.load(SEQ4_DATA_DIR + "test_adjs_0523.npy") test_energies = np.load(SEQ4_DATA_DIR + "test_energies_0523.npy") if gnn: train_diracs = torch.eye(train_adjs.shape[-1]).unsqueeze(0).repeat( train_adjs.shape[0], 1, 1) train_tup = (torch.Tensor(train_diracs), torch.Tensor(train_adjs), torch.Tensor(train_energies)) else: train_tup = (torch.Tensor(train_coeffs), torch.Tensor(train_energies)) if gnn: test_diracs = torch.eye(test_adjs.shape[-1]).unsqueeze(0).repeat( test_adjs.shape[0], 1, 1) test_tup = (torch.Tensor(test_diracs), torch.Tensor(test_adjs), torch.Tensor(test_energies)) else: test_tup = (torch.Tensor(test_coeffs), torch.Tensor(test_adjs), torch.Tensor(test_energies)) ################# # SUBSET DATA ################# if subsize != None: train_tup, _ = eval_metrics.compute_subsample(train_tup, subsize) test_tup, _ = eval_metrics.compute_subsample(test_tup, subsize) train_dataset = torch.utils.data.TensorDataset(*train_tup) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True) return train_loader, train_tup, test_tup
train_embed.cpu().detach().numpy()) np.save(save_dir + "test_embedding.npy", test_embed.cpu().detach().numpy()) # EVALUATION ON TEST SET # energy pred mse print("getting test set predictions") with torch.no_grad(): x_recon_test = model(test_tup[0])[0] y_pred_test = model.predict_from_data(test_tup[0]) pred_test_val = nn.MSELoss()(y_pred_test.flatten(), test_tup[-1].flatten()) print("gathering eval subsets") eval_tup_list = [ eval_metrics.compute_subsample([test_embed, test_tup[-1]], 10000)[0] for i in range(8) ] # trainer.test() print("getting smoothness vals") embed_eval_array = np.expand_dims( np.array([x[0].numpy() for x in eval_tup_list]), 0) energy_eval_array = np.array([x[1].numpy() for x in eval_tup_list]) print('embed_eval_array shape: {}'.format(embed_eval_array.shape)) print('energy_eval_array shape: {}'.format(energy_eval_array.shape)) energy_smoothness = eval_metrics.eval_over_replicates( embed_eval_array, energy_eval_array, eval_metrics.get_smoothnes_kNN, [5, 10])[0]