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()
losses_test = [] num_cluster_iter = args.num_cluster_iter ############################################################################## #MAKE AUXILIARY DATA FOR OBJECTIVES ############################################################################## if make_objectives: if args.objective == 'kcenter': try: dist_all = torch.load('{}_test_dist.pt'.format(args.dataset)) dist_train = torch.load('{}_{}_train_dist.pt'.format( args.dataset, train_pct)) diameter = dist_all.max() except: dist_all = make_all_dists(bin_adj_all, 100) diameter = dist_all[dist_all < 100].max() dist_all[dist_all == 100] = diameter torch.save(dist_all, '{}_test_dist.pt'.format(args.dataset)) dist_train = make_all_dists(bin_adj_train, 100) dist_train[dist_train == 100] = diameter torch.save(dist_train, '{}_{}_train_dist.pt'.format(args.dataset, train_pct)) obj_train = CenterObjective(dist_train, diameter, args.kcentermintemp) obj_train_hardmax = CenterObjective(dist_train, diameter, args.kcentermintemp, hardmax=True) obj_test = CenterObjective(dist_all, diameter, args.kcentertemp,
aucs[algo] = np.zeros(numtest) if args.objective == 'modularity': mods_test = [make_modularity_matrix(A) for A in bin_adj_all] mods_train = [make_modularity_matrix(A) for A in bin_adj_train] test_object = mods_test train_object = mods_train loss_fn = loss_modularity elif args.objective == 'kcenter': for i in range(num_graphs): try: dist_all.append(torch.load('{}_{}_test_dist.pt'.format(args.dataset, i))) dist_train.append(torch.load('{}_{}_{:.2f}_train_dist.pt'.format(args.dataset, i, train_pct))) diameter = dist_all[-1].max() except: dist_all_i = make_all_dists(bin_adj_all[i], 100) diameter = dist_all_i[dist_all_i < 100].max() dist_all_i[dist_all_i == 100] = diameter torch.save(dist_all_i, '{}_{}_test_dist.pt'.format(args.dataset, i)) dist_all.append(dist_all_i) dist_train_i = make_all_dists(bin_adj_train[i], 100) dist_train_i[dist_train_i == 100] = diameter torch.save(dist_train_i, '{}_{}_{:.2f}_train_dist.pt'.format(args.dataset, i, train_pct)) dist_train.append(dist_train_i) obj_train = [CenterObjective(d, diameter, args.kcentermintemp) for d in dist_train] obj_train_hardmax = [CenterObjective(d, diameter, args.kcentermintemp, hardmax=True) for d in dist_train] obj_test = [CenterObjective(d, diameter, args.kcentertemp, hardmax=True) for d in dist_all] obj_test_softmax = [CenterObjective(d, diameter, args.kcentermintemp) for d in dist_all] test_object = obj_test