edge_splitter.py

!pip install stellargraph
import stellargraph as sg
from stellargraph.data import EdgeSplitter

G = nx.read_gpickle('usair.gpickle')

edge_splitter_test = EdgeSplitter(G)
G_test, edge_ids_test, edge_labels_test = edge_splitter_test.train_test_split(
    p=0.1, method="global", keep_connected=True)

edge_splitter_train = EdgeSplitter(G_test)
G_train, edge_ids_train, edge_labels_train = edge_splitter_train.train_test_split(
    p=0.99, method="global", keep_connected=False)

edge_ids_test_pos = edge_ids_test[0:int(len(edge_labels_test)/2)]
edge_ids_test_neg = edge_ids_test[int(len(edge_labels_test)/2):]

edge_ids_train_pos = edge_ids_train[0:int(len(edge_labels_train)/2)]
edge_ids_train_neg = edge_ids_train[int(len(edge_labels_train)/2):]

train = {'edge': torch.tensor(edge_ids_train_pos)}
test = {'edge': torch.tensor(edge_ids_test_pos),'edge_neg':torch.tensor(edge_ids_test_neg)}
my_split = {'train': train, 'test': test}

len(edge_ids_test_pos),len(edge_ids_test_neg),len(edge_ids_train_pos),len(edge_ids_train_neg)

usair = nx.read_gpickle('usair.gpickle')
usair = from_networkx(usair)
usair['x'] = torch.tensor(np.zeros((332,1)))

train_dataset = eval('SEALDataset')(
    '/train', 
    usair, 
    my_split, 
    num_hops=args.num_hops, 
    percent=args.train_percent, 
    split='train', 
    use_coalesce=use_coalesce, 
    node_label=args.node_label, 
    ratio_per_hop=args.ratio_per_hop, 
    max_nodes_per_hop=args.max_nodes_per_hop,) 

test_dataset = eval('SEALDataset')(
    '/test', 
    usair, 
    my_split, 
    num_hops=args.num_hops, 
    percent=args.test_percent, 
    split='test', 
    use_coalesce=use_coalesce, 
    node_label=args.node_label, 
    ratio_per_hop=args.ratio_per_hop, 
    max_nodes_per_hop=args.max_nodes_per_hop,)