code: import os
  import dgl
  import torch
  import numpy as np
  import torch.nn as nn
  from model import PGNN
  from sklearn.metrics import roc_auc_score
  from utils import get_dataset, preselect_anchor
  
  import warnings
  warnings.filterwarnings('ignore')
  
  def get_loss(p, data, out, loss_func, device, get_auc=True):
      edge_mask = np.concatenate((data['positive_edges_{}'.format(p)], data['negative_edges_{}'.format(p)]), axis=-1)
  
      nodes_first = torch.index_select(out, 0, torch.from_numpy(edge_mask[0, :]).long().to(out.device))
      nodes_second = torch.index_select(out, 0, torch.from_numpy(edge_mask[1, :]).long().to(out.device))
  
      pred = torch.sum(nodes_first * nodes_second, dim=-1)
  
      label_positive = torch.ones([data['positive_edges_{}'.format(p)].shape[1], ], dtype=pred.dtype)
      label_negative = torch.zeros([data['negative_edges_{}'.format(p)].shape[1], ], dtype=pred.dtype)
      label = torch.cat((label_positive, label_negative)).to(device)
      loss = loss_func(pred, label)
  
      if get_auc:
          auc = roc_auc_score(label.flatten().cpu().numpy(), torch.sigmoid(pred).flatten().data.cpu().numpy())
          return loss, auc
      else:
          return loss
  
  def train_model(data, model, loss_func, optimizer, device, g_data):
      model.train()
      out = model(g_data)
  
      loss = get_loss('train', data, out, loss_func, device, get_auc=False)
  
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()
      optimizer.zero_grad()
  
      return g_data
  
  def eval_model(data, g_data, model, loss_func, device):
      model.eval()
      out = model(g_data)
  
      # train loss and auc
      tmp_loss, auc_train = get_loss('train', data, out, loss_func, device)
      loss_train = tmp_loss.cpu().data.numpy()
  
      # val loss and auc
      _, auc_val = get_loss('val', data, out, loss_func, device)
  
      # test loss and auc
      _, auc_test = get_loss('test', data, out, loss_func, device)
  
      return loss_train, auc_train, auc_val, auc_test
  
  def main(args):
      # The mean and standard deviation of the experiment results
      # are stored in the 'results' folder
      if not os.path.isdir('results'):
          os.mkdir('results')
  
      if torch.cuda.is_available():
          device = 'cuda:0'
      else:
          device = 'cpu'
  
      print('Learning Type: {}'.format(['Transductive', 'Inductive'][args.inductive]),
            'Task: {}'.format(args.task))
  
      results = []
  
      for repeat in range(args.repeat_num):
          data = get_dataset(args)
  
          # pre-sample anchor nodes and compute shortest distance values for all epochs
          g_list, anchor_eid_list, dist_max_list, edge_weight_list = preselect_anchor(data, args)
  
          # model
          model = PGNN(input_dim=data['feature'].shape[1]).to(device)
  
          # loss
          optimizer = torch.optim.Adam(model.parameters(), lr=1e-2, weight_decay=5e-4)
          loss_func = nn.BCEWithLogitsLoss()
  
          best_auc_val = -1
          best_auc_test = -1
  
          for epoch in range(args.epoch_num):
              if epoch == 200:
                  for param_group in optimizer.param_groups:
                      param_group['lr'] /= 10
  
              g = dgl.graph(g_list[epoch])
              g.ndata['feat'] = torch.FloatTensor(data['feature'])
              g.edata['sp_dist'] = torch.FloatTensor(edge_weight_list[epoch])
              g_data = {
                  'graph': g.to(device),
                  'anchor_eid': anchor_eid_list[epoch],
                  'dists_max': dist_max_list[epoch]
              }
  
              train_model(data, model, loss_func, optimizer, device, g_data)
  
              loss_train, auc_train, auc_val, auc_test = eval_model(
                  data, g_data, model, loss_func, device)
              if auc_val > best_auc_val:
                  best_auc_val = auc_val
                  best_auc_test = auc_test
  
              if epoch % args.epoch_log == 0:
                  print(repeat, epoch, 'Loss {:.4f}'.format(loss_train), 'Train AUC: {:.4f}'.format(auc_train),
                        'Val AUC: {:.4f}'.format(auc_val), 'Test AUC: {:.4f}'.format(auc_test),
                        'Best Val AUC: {:.4f}'.format(best_auc_val), 'Best Test AUC: {:.4f}'.format(best_auc_test))
  
          results.append(best_auc_test)
  
      results = np.array(results)
      results_mean = np.mean(results).round(6)
      results_std = np.std(results).round(6)
      print('-----------------Final-------------------')
      print(results_mean, results_std)
  
      with open('results/{}_{}_{}.txt'.format(['Transductive', 'Inductive'][args.inductive], args.task,
                                              args.k_hop_dist), 'w') as f:
          f.write('{}, {}\n'.format(results_mean, results_std))
  
  if __name__ == '__main__':
      from argparse import ArgumentParser
  
      parser = ArgumentParser()
      parser.add_argument('--task', type=str, default='link', choices=['link', 'link_pair'])
      parser.add_argument('--inductive', action='store_true',
                          help='Inductive learning or transductive learning')
      parser.add_argument('--k_hop_dist', default=-1, type=int,
                          help='K-hop shortest path distance, -1 means exact shortest path.')
  
      parser.add_argument('--epoch_num', type=int, default=2001)
      parser.add_argument('--repeat_num', type=int, default=10)
      parser.add_argument('--epoch_log', type=int, default=10)
  
      args = parser.parse_args()
      main(args)
  
apis: torch.sum torch.FloatTensor torch.cuda.is_available torch.from_numpy torch.nn.BCEWithLogitsLoss numpy.array numpy.std torch.sigmoid torch.cat numpy.mean