(self, graph, embedding_model_creator=Word2Vec, return_dict=False)
| 52 | self.iteration = iteration |
| 53 | |
| 54 | def forward(self, graph, embedding_model_creator=Word2Vec, return_dict=False): |
| 55 | nx_g = graph.to_networkx() |
| 56 | self.G = nx_g |
| 57 | walks = self._simulate_walks(self.walk_length, self.walk_num) |
| 58 | walks = [[str(node) for node in walk] for walk in walks] |
| 59 | print("training word2vec...") |
| 60 | model = embedding_model_creator( |
| 61 | walks, |
| 62 | vector_size=self.dimension, |
| 63 | window=self.window_size, |
| 64 | min_count=0, |
| 65 | sg=1, |
| 66 | workers=self.worker, |
| 67 | epochs=self.iteration, |
| 68 | ) |
| 69 | id2node = dict([(vid, node) for vid, node in enumerate(nx_g.nodes())]) |
| 70 | embeddings = np.asarray([model.wv[str(id2node[i])] for i in range(len(id2node))]) |
| 71 | |
| 72 | if return_dict: |
| 73 | features_matrix = dict() |
| 74 | for vid, node in enumerate(nx_g.nodes()): |
| 75 | features_matrix[node] = embeddings[vid] |
| 76 | else: |
| 77 | features_matrix = np.zeros((graph.num_nodes, embeddings.shape[1])) |
| 78 | nx_nodes = nx_g.nodes() |
| 79 | features_matrix[nx_nodes] = embeddings[np.arange(graph.num_nodes)] |
| 80 | return features_matrix |
| 81 | |
| 82 | def _walk(self, start_node, walk_length): |
| 83 | # Simulate a random walk starting from start node. |
nothing calls this directly
no test coverage detected