fix data generation of MAG240M
@@ -24,6 +24,8 @@ class MAG240M(object):
def prepare_data(self):
dataset = MAG240MDataset(self.data_dir)
+ total_num_nodes = dataset.num_papers + dataset.num_authors + dataset.num_institutions
+
graph_file_list = []
paper_edge_path = f'{dataset.dir}/paper_to_paper_symmetric_pgl_split'
graph_file_list.append(paper_edge_path)
@@ -39,7 +41,7 @@ class MAG240M(object):
edge_index = np.vstack((edge_index, edges_new))
edge_types = np.full([edge_index.shape[0], ], 0, dtype='int32')
- graph = Graph(edge_index, num_nodes=dataset.num_papers, edge_feat={'edge_type': edge_types})
+ graph = Graph(edge_index, num_nodes=total_num_nodes, edge_feat={'edge_type': edge_types})
graph.adj_dst_index
graph.dump(paper_edge_path)
log.info(f'Done! [{time.perf_counter() - t:.2f}s]')
@@ -61,7 +63,7 @@ class MAG240M(object):
edge_types = np.full(row.shape, 1, dtype='int32')
edge_index = np.stack([row, col], axis=1)
- graph = Graph(edge_index, edge_feat={'edge_type': edge_types})
+ graph = Graph(edge_index, num_nodes=total_num_nodes, edge_feat={'edge_type': edge_types})
graph.adj_dst_index
graph.dump(author_edge_path)
log.info(f'Done! finish author_edge [{time.perf_counter() - t:.2f}s]')
@@ -84,7 +86,7 @@ class MAG240M(object):
edge_types = np.full(row.shape, 2, dtype='int32')
edge_index = np.stack([col, row], axis=1)
- graph = Graph(edge_index, edge_feat={'edge_type': edge_types})
+ graph = Graph(edge_index, num_nodes=total_num_nodes, edge_feat={'edge_type': edge_types})
graph.adj_dst_index
graph.dump(author_edge_path)
log.info(f'Done! finish author_edge [{time.perf_counter() - t:.2f}s]')
@@ -110,7 +112,7 @@ class MAG240M(object):
edge_types = np.full(row.shape, 3, dtype='int32')
edge_index = np.stack([row, col], axis=1)
- graph = Graph(edge_index, edge_feat={'edge_type': edge_types})
+ graph = Graph(edge_index, num_nodes=total_num_nodes, edge_feat={'edge_type': edge_types})
graph.adj_dst_index
graph.dump(institution_edge_path)
log.info(f'Done! finish institution_edge [{time.perf_counter() - t:.2f}s]')
@@ -136,7 +138,7 @@ class MAG240M(object):
edge_types = np.full(row.shape, 4, dtype='int32')
edge_index = np.stack([col, row], axis=1)
- graph = Graph(edge_index, edge_feat={'edge_type': edge_types})
+ graph = Graph(edge_index, num_nodes=total_num_nodes, edge_feat={'edge_type': edge_types})
graph.adj_dst_index
graph.dump(institution_edge_path)
log.info(f'Done! finish institution_edge [{time.perf_counter() - t:.2f}s]')
