-
Notifications
You must be signed in to change notification settings - Fork 2
/
load_graph.py
47 lines (40 loc) · 1.47 KB
/
load_graph.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
import os
import pickle
import networkx as nx
import scipy.io as sio
def read_graph(name, is_directed=False):
"""
read graph data from edge list file
:param name: name of the graph
:param is_directed: if the graph is directed or not
:return: adjacency matrix, Networkx Graph (if undirected) or DiGraph (if directed)
"""
if is_directed:
PATH = os.path.join('data', 'directed', '{}.txt'.format(name))
G = nx.read_edgelist(PATH, create_using=nx.DiGraph(), nodetype=int, data=(('weight', float),))
else:
PATH = os.path.join('data', '{}.txt'.format(name))
G = nx.read_edgelist(PATH, create_using=nx.Graph(), nodetype=int, data=(('weight', float),))
return G
def read_mat(name):
"""
read .mat file
:param name: dataset name
:return: a dict containing adjacency matrix, nx.Graph() and its node labels
"""
result = dict()
PATH = os.path.join('data', '{}.mat'.format(name))
matfile = sio.loadmat(PATH)
result['adjacency'] = matfile['network']
result['label'] = matfile['group']
result['graph'] = nx.from_scipy_sparse_matrix(result['adjacency'], create_using=nx.Graph(), edge_attribute='weight')
return result
def read_pickle(name):
"""
read .pickle file
:param name: dataset name
:return: a dict containing adjacency matrix, nx.Graph() and its node labels
"""
with open('data/{}.pickle'.format(name), 'rb') as f:
result = pickle.load(f)
return result