MCPcopy
hub / github.com/THUDM/CogDL / forward

Method forward

cogdl/models/emb/node2vec.py:72–103  ·  view source on GitHub ↗
(self, graph, return_dict=False)

Source from the content-addressed store, hash-verified

70 self.q = q
71
72 def forward(self, graph, return_dict=False):
73 G = graph.to_networkx()
74 self.G = G
75 is_directed = nx.is_directed(self.G)
76 for i, j in G.edges():
77 G[i][j]["weight"] = G[i][j].get("weight", 1.0)
78 if not is_directed:
79 G[j][i]["weight"] = G[j][i].get("weight", 1.0)
80 self._preprocess_transition_probs()
81 walks = self._simulate_walks(self.walk_num, self.walk_length)
82 walks = [[str(node) for node in walk] for walk in walks]
83 model = Word2Vec(
84 walks,
85 vector_size=self.dimension,
86 window=self.window_size,
87 min_count=0,
88 sg=1,
89 workers=self.worker,
90 epochs=self.iteration,
91 )
92 id2node = dict([(vid, node) for vid, node in enumerate(G.nodes())])
93 embeddings = np.asarray([model.wv[str(id2node[i])] for i in range(len(id2node))])
94
95 if return_dict:
96 features_matrix = dict()
97 for vid, node in enumerate(G.nodes()):
98 features_matrix[node] = embeddings[vid]
99 else:
100 features_matrix = np.zeros((graph.num_nodes, embeddings.shape[1]))
101 nx_nodes = G.nodes()
102 features_matrix[nx_nodes] = embeddings[np.arange(graph.num_nodes)]
103 return features_matrix
104
105 def _node2vec_walk(self, walk_length, start_node):
106 # Simulate a random walk starting from start node.

Callers

nothing calls this directly

Calls 5

_simulate_walksMethod · 0.95
nodesMethod · 0.80
to_networkxMethod · 0.45
getMethod · 0.45

Tested by

no test coverage detected