def test_twostage(model_ts, test_instances_eval, algoname): for test_i in test_instances_eval: #predict probability that all unobserved edges exist n = adj_train[test_i].shape[0] indices = torch.tensor(np.arange(n)) to_pred = torch.zeros(n**2, 2) to_pred[:, 1] = indices.repeat(n) for i in range(n): to_pred[i * n:(i + 1) * n, 0] = i to_pred = to_pred.long() preds = model_ts(features_train[test_i], adj_train[test_i], to_pred) preds = nn.Sigmoid()(preds).view(n, n) preds = bin_adj_train[test_i] + (1 - bin_adj_train[test_i]) * preds if args.objective == 'modularity': r = greedy_modularity_communities(preds, K) loss = loss_fn(None, r, None, None, bin_adj_all[test_i], test_object[test_i], args).item() vals[algoname + '-agglomerative'][test_instances.index(test_i)] = loss r = partition(preds, K) loss = loss_fn(None, r, None, None, bin_adj_all[test_i], test_object[test_i], args).item() vals[algoname + '-recursive'][test_instances.index(test_i)] = loss degrees = preds.sum(dim=1) preds = torch.diag(1. / degrees) @ preds mod_pred = make_modularity_matrix(preds) r = baseline_spectral(mod_pred, K) loss = loss_fn(None, r, None, None, bin_adj_all[test_i], test_object[test_i], args).item() vals[algoname + '-spectral'][test_instances.index(test_i)] = loss elif args.objective == 'kcenter': print('making dists') if args.use_igraph: dist_ts = make_dists_igraph(preds) else: dist_ts = make_all_dists(preds, 100) diameter = dist_ts[dist_ts < 100].max() dist_ts[dist_ts == 100] = diameter print(test_i) dist_ts = dist_ts.float() diameter = dist_ts.max() x = gonzalez_kcenter(dist_ts, K) loss = obj_test[test_i](x) vals[algoname + '-gonzalez'][test_instances.index(test_i)] = loss.item() x = greedy_kcenter(dist_ts, diameter, K) loss = obj_test[test_i](x) vals[algoname + '-greedy'][test_instances.index(test_i)] = loss.item()
indices = torch.tensor(np.arange(n)) to_pred = torch.zeros(n**2, 2) to_pred[:, 1] = indices.repeat(n) for i in range(n): to_pred[i * n:(i + 1) * n, 0] = i to_pred = to_pred.long() preds = model_ts(features_train, adj_train, to_pred) preds = nn.Sigmoid()(preds).view(n, n) preds = bin_adj_train + (1 - bin_adj_train) * preds if args.objective == 'modularity': r = greedy_modularity_communities(preds, K) print('agglomerative', loss_fn(None, r, None, None, bin_adj_all, test_object, args)) r = partition(preds, K) print('recursive', loss_fn(None, r, None, None, bin_adj_all, test_object, args)) degrees = preds.sum(dim=1) preds = torch.diag(1. / degrees) @ preds mod_pred = make_modularity_matrix(preds) r = baseline_spectral(mod_pred, K) print('spectral', loss_fn(None, r, None, None, bin_adj_all, test_object, args)) elif args.objective == 'kcenter': try: dist_ts = torch.load('{}_twostage_dist.pt'.format(args.dataset)) print('loaded ts dists from {}'.format( '{}_twostage_dist.pt'.format(args.dataset))) except: print('making dists')