CheapNode2Vec/visualization.py at master · shovalf/CheapNode2Vec · GitHub

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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
"""
Visualization Task for Evaluation. For full details please enter to the survey file in git.
Code explanation:
For this task, one should have a labeled graph: 2 files are required: Graph edges in '.edgelist' or '.txt' format and
nodes' labels in '.txt' format. For example labeled graphs you can enter to the link in the github page. You should
insert the in the appropriate place in the main function in the file 'directed_cheap_node2vec' or
'undirected_cheap_node2vec' - depends on your graph.
Go to the file 'directed_cheap_node2vec' (or undirected). In the main function change 'initial' variable to a list
that consists the percentage of nodes you want in the initial projection (for example, for pubmed2, 0.975 means 100
nodes in the initial projection). Go back to this file and run main('directed') or main('undirected'.
"""

import matplotlib.pyplot as plt
try:
    import cPickle as pickle
except:
    import pickle
import numpy as np
from node2vec import Node2Vec
from sklearn.manifold import TSNE
import pandas as pd
import matplotlib as mpl


mpl.rcParams['xtick.labelsize'] = 14
mpl.rcParams['ytick.labelsize'] = 14
mpl.rcParams['axes.titlesize'] = 20


def regular_node2vec(G, nodes):
    """
    Run regular node2vec
    :param G: Our graph
    :param nodes: nodes of the graph
    :return: a list of projections of every node as np arrays
    """
    node2vec = Node2Vec(G, dimensions=10, walk_length=80, num_walks=16, workers=2)
    model = node2vec.fit()
    projections = []
    for i in range(len(nodes)):
        projections.append(np.asarray(model.wv.get_vector(nodes[i])))
    return projections


def cheap_node2vec(final_dict_proj):
    """
    Run cheap node2vec
    :param final_dict_proj: A dictionary with keys==nodes in projection and values==projection
    :return: a list of projections of every node as np arrays
    """
    keys = list(final_dict_proj.keys())
    projections = []
    for i in range(len(keys)):
        projections.append(final_dict_proj[keys[i]])
    return projections


def read_labels(file):
    """
    Read the labels file and return the labels as a list
    :param file: a file with labels for every node
    :return: a list of the labels of every node
    """
    c = np.loadtxt(file, dtype=int)
    labels = {x: y for (x, y) in c}
    list_labels = []
    keys = list(labels.keys())
    for i in range(len(keys)):
        list_labels.append(labels[keys[i]])
    return list_labels


def visualization(projections, nodes, i, labels):
    """
    The visualization task explained in details in the pdf file attached in the github.
    :param projections: a list of projections of every node as np arrays
    :param nodes: nodes of the graph
    :param i: number of figure
    :param labels:a list of the labels of every node
    :return: tsne representation of both regular node2vec and cheap node2vec
    """
    projections = np.asarray(projections)

    names = []
    for i in range(len(nodes)):
        names.append(nodes[i])

    data = pd.DataFrame(projections)
    X = data.values

    tsne_points = TSNE(n_components=2).fit_transform(X)

    fig = plt.figure(i)
    ax = fig.add_subplot(111)
    plt.xlabel('First Dimension')
    plt.ylabel('Second Dimension')

    if len(labels) == 0:
        ax.scatter(tsne_points[:, 0], tsne_points[:, 1], cmap='tab10', alpha=0.8, s=8)
        plt.title('Regular Node2vec Representation with TSNE')
    else:
        ax.scatter(tsne_points[:, 0], tsne_points[:, 1], labels, cmap='tab10', alpha=0.8, s=8)
        plt.title('Cheap Node2vec Representation with TSNE')
    plt.show()


def main(mission):
    """
    Main function to run all task. It can be run on both directed and undirected graph, determined
    by the mission variable- 'directed' or 'undirected'
    :param mission: 'directed' or 'undirected'
    :return: tsne representation of both regular node2vec and cheap node2vec
    """
    if mission == 'directed':
        from directed_cheap_node2vec import list_dicts, G, file
        final_dict_proj = list_dicts[0]
        nodes = list(G.nodes())
        for i in range(len(nodes)):
            nodes[i] = str(nodes[i])
        if file is not None:
            labels = read_labels(file)
        else:
            labels = []
        projections_regular = regular_node2vec(G, nodes)
        projections_cheap = cheap_node2vec(final_dict_proj)
        visualization(projections_regular, nodes, 1, labels)
        visualization(projections_cheap, nodes, 2, labels)
    elif mission == 'undirected':
        from undirected_cheap_node2vec import list_dicts, G, file
        final_dict_proj = list_dicts[0]
        nodes = list(G.nodes())
        for i in range(len(nodes)):
            nodes[i] = str(nodes[i])
        if file is not None:
            labels = read_labels(file)
        else:
            labels = []
        projections_regular = regular_node2vec(G, nodes)
        projections_cheap = cheap_node2vec(final_dict_proj)
        visualization(projections_regular, nodes, 1, labels)
        visualization(projections_cheap, nodes, 2, labels)
    else:
        print("If your graph is directed please input 'directed', for an undirected graph"
              "please input 'undirected'")


# can be changed to undirected, depends on the graph
main('directed')