forked from GraphSAINT/GraphSAINT
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathconvert.py
53 lines (47 loc) · 1.38 KB
/
convert.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
import scipy.sparse as sp
import numpy as np
import networkx as nx
import sys
import json
import os
from networkx.readwrite import json_graph
dataset_str=sys.argv[1]
baseline_str='data.ignore/'+dataset_str+'/'
dataset_str='data/'+dataset_str+'/'
if not os.path.exists(baseline_str[:-1]):
os.mkdir(baseline_str[:-1])
# G.json
adj_full=sp.load_npz(dataset_str+'adj_full.npz')
G=nx.from_scipy_sparse_matrix(adj_full)
print('nx: finish load graph')
data=json_graph.node_link_data(G)
role=json.load(open(dataset_str+'role.json','r'))
te=set(role['te'])
va=set(role['va'])
for node in data['nodes']:
node['test']=False
node['val']=False
if node['id'] in te:
node['test']=True
elif node['id'] in va:
node['val']=True
for edge in data['links']:
del edge['weight']
edge['target']=int(edge['target'])
with open(baseline_str+'G.json','w') as f:
json.dump(data,f)
# id_map.json
id_map={}
for i in range(G.number_of_nodes()):
id_map[str(i)]=i
with open(baseline_str+'id_map.json','w') as f:
json.dump(id_map,f)
# feats.npy
feats=np.load(dataset_str+'feats.npy')
np.save(baseline_str+'feats.npy',feats)
# class_map.json
class_map=json.load(open(dataset_str+'class_map.json','r'))
for k,v in class_map.items():
class_map[k]=v
with open(baseline_str+'class_map.json','w') as f:
json.dump(class_map,f)