-
Notifications
You must be signed in to change notification settings - Fork 7
/
MethodWLNodeColoring.py
59 lines (50 loc) · 2.09 KB
/
MethodWLNodeColoring.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
54
55
56
57
58
59
'''
Concrete MethodModule class for a specific learning MethodModule
'''
# Copyright (c) 2017 Jiawei Zhang <jwzhanggy@gmail.com>
# License: TBD
from method import method
import hashlib
class MethodWLNodeColoring(method):
data = None
max_iter = 2
node_color_dict = {}
node_neighbor_dict = {}
def setting_init(self, node_list, link_list):
for node in node_list:
self.node_color_dict[node] = 1
self.node_neighbor_dict[node] = {}
for pair in link_list:
u1, u2 = pair
if u1 not in self.node_neighbor_dict:
self.node_neighbor_dict[u1] = {}
if u2 not in self.node_neighbor_dict:
self.node_neighbor_dict[u2] = {}
self.node_neighbor_dict[u1][u2] = 1
self.node_neighbor_dict[u2][u1] = 1
def WL_recursion(self, node_list):
iteration_count = 1
while True:
new_color_dict = {}
for node in node_list:
neighbors = self.node_neighbor_dict[node]
neighbor_color_list = [self.node_color_dict[neb] for neb in neighbors]
color_string_list = [str(self.node_color_dict[node])] + sorted([str(color) for color in neighbor_color_list])
color_string = "_".join(color_string_list)
hash_object = hashlib.md5(color_string.encode())
hashing = hash_object.hexdigest()
new_color_dict[node] = hashing
color_index_dict = {k: v+1 for v, k in enumerate(sorted(set(new_color_dict.values())))}
for node in new_color_dict:
new_color_dict[node] = color_index_dict[new_color_dict[node]]
if self.node_color_dict == new_color_dict or iteration_count == self.max_iter:
return
else:
self.node_color_dict = new_color_dict
iteration_count += 1
def run(self):
node_list = self.data['idx']
link_list = self.data['edges']
self.setting_init(node_list, link_list)
self.WL_recursion(node_list)
return self.node_color_dict