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]
