forked from palash1992/GEM
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test.py
51 lines (43 loc) · 2.26 KB
/
test.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
import matplotlib.pyplot as plt
from time import time
from gem.utils import graph_util, plot_util
from gem.evaluation import visualize_embedding as viz
from gem.evaluation import evaluate_graph_reconstruction as gr
from gem.embedding.gf import GraphFactorization
from gem.embedding.hope import HOPE
from gem.embedding.lap import LaplacianEigenmaps
from gem.embedding.lle import LocallyLinearEmbedding
from gem.embedding.node2vec import node2vec
from gem.embedding.sdne import SDNE
# File that contains the edges. Format: source target
# Optionally, you can add weights as third column: source target weight
edge_f = 'gem/data/TEST_50M.edgelist'
# Specify whether the edges are directed
isDirected = True
# Load graph
G = graph_util.loadGraphFromEdgeListTxt(edge_f, directed=isDirected)
G = G.to_directed()
models = []
# You can comment out the methods you don't want to run
models.append(GraphFactorization(2, 50000, 1*10**-4, 1.0))
models.append(HOPE(4, 0.01))
models.append(LaplacianEigenmaps(2))
models.append(LocallyLinearEmbedding(2))
models.append(node2vec(2, 1, 80, 10, 10, 1, 1))
models.append(SDNE(d=2, beta=5, alpha=1e-5, nu1=1e-6, nu2=1e-6, K=3,n_units=[50, 15,], rho=0.3, n_iter=50, xeta=0.01,n_batch=500,
modelfile=['./intermediate/enc_model.json', './intermediate/dec_model.json'],
weightfile=['./intermediate/enc_weights.hdf5', './intermediate/dec_weights.hdf5']))
for embedding in models:
print ('Num nodes: %d, num edges: %d' % (G.number_of_nodes(), G.number_of_edges()))
t1 = time()
# Learn embedding - accepts a networkx graph or file with edge list
Y, t = embedding.learn_embedding(graph=G, edge_f=None, is_weighted=True, no_python=True)
print (embedding._method_name+':\n\tTraining time: %f' % (time() - t1))
# Evaluate on graph reconstruction
MAP, prec_curv = gr.evaluateStaticGraphReconstruction(G, embedding, Y, None)
#---------------------------------------------------------------------------------
print(("\tMAP: {} \t preccision curve: {}\n\n\n\n"+'-'*100).format(MAP,prec_curv))
#---------------------------------------------------------------------------------
# Visualize
viz.plot_embedding2D(embedding.get_embedding(), di_graph=G, node_colors=None)
plt.show()