def test_hetero_graph_basics(self): G = generate_simple_hete_graph() hete = HeteroGraph(G) hete = HeteroGraph(node_feature=hete.node_feature, node_label=hete.node_label, edge_feature=hete.edge_feature, edge_label=hete.edge_label, edge_index=hete.edge_index, directed=True) self.assertEqual(hete.num_node_features('n1'), 10) self.assertEqual(hete.num_node_features('n2'), 12) self.assertEqual(hete.num_edge_features(('n1', 'e1', 'n1')), 8) self.assertEqual(hete.num_edge_features(('n1', 'e2', 'n2')), 12) self.assertEqual(hete.num_nodes('n1'), 4) self.assertEqual(hete.num_nodes('n2'), 5) self.assertEqual(len(hete.node_types), 2) self.assertEqual(len(hete.edge_types), 2) message_types = hete.message_types self.assertEqual(len(message_types), 7) self.assertEqual(hete.num_node_labels('n1'), 2) self.assertEqual(hete.num_node_labels('n2'), 2) self.assertEqual(hete.num_edge_labels(('n1', 'e1', 'n1')), 2) self.assertEqual(hete.num_edge_labels(('n1', 'e2', 'n2')), 2) self.assertEqual(hete.num_edges(message_types[0]), 3) self.assertEqual(len(hete.node_label_index), 2)
def test_hetero_graph_split(self): # directed G G = generate_dense_hete_graph() hete = HeteroGraph(G) hete = HeteroGraph( node_feature=hete.node_feature, node_label=hete.node_label, edge_feature=hete.edge_feature, edge_label=hete.edge_label, edge_index=hete.edge_index, ) # node hete_node = hete.split() for node_type in hete.node_label_index: num_nodes = len(hete.node_label_index[node_type]) node_0 = int(num_nodes * 0.8) node_1 = int(num_nodes * 0.1) node_2 = num_nodes - node_0 - node_1 self.assertEqual( len(hete_node[0].node_label_index[node_type]), node_0, ) self.assertEqual( len(hete_node[1].node_label_index[node_type]), node_1, ) self.assertEqual( len(hete_node[2].node_label_index[node_type]), node_2, ) # node with specified split type node_split_types = ['n1'] hete_node = hete.split(split_types=node_split_types) for node_type in hete.node_label_index: if node_type in node_split_types: num_nodes = len(hete.node_label_index[node_type]) node_0 = int(num_nodes * 0.8) node_1 = int(num_nodes * 0.1) node_2 = num_nodes - node_0 - node_1 self.assertEqual( len(hete_node[0].node_label_index[node_type]), node_0, ) self.assertEqual( len(hete_node[1].node_label_index[node_type]), node_1, ) self.assertEqual( len(hete_node[2].node_label_index[node_type]), node_2, ) else: self.assertEqual( len(hete_node[0].node_label_index[node_type]), len(hete.node_label_index[node_type]), ) self.assertEqual( len(hete_node[1].node_label_index[node_type]), len(hete.node_label_index[node_type]), ) self.assertEqual( len(hete_node[2].node_label_index[node_type]), len(hete.node_label_index[node_type]), ) # edge hete_edge = hete.split(task='edge') for edge_type in hete.edge_label_index: num_edges = int(hete.edge_label_index[edge_type].shape[1]) edge_0 = int(num_edges * 0.8) edge_1 = int(num_edges * 0.1) edge_2 = num_edges - edge_0 - edge_1 self.assertEqual( hete_edge[0].edge_label_index[edge_type].shape[1], edge_0, ) self.assertEqual( hete_edge[1].edge_label_index[edge_type].shape[1], edge_1, ) self.assertEqual( hete_edge[2].edge_label_index[edge_type].shape[1], edge_2, ) # edge with specified split type edge_split_types = [('n1', 'e1', 'n1'), ('n1', 'e2', 'n2')] hete_edge = hete.split(task='edge', split_types=edge_split_types) for edge_type in hete.edge_label_index: if edge_type in edge_split_types: num_edges = int(hete.edge_label_index[edge_type].shape[1]) edge_0 = int(num_edges * 0.8) edge_1 = int(num_edges * 0.1) edge_2 = num_edges - edge_0 - edge_1 self.assertEqual( hete_edge[0].edge_label_index[edge_type].shape[1], edge_0, ) self.assertEqual( hete_edge[1].edge_label_index[edge_type].shape[1], edge_1, ) self.assertEqual( hete_edge[2].edge_label_index[edge_type].shape[1], edge_2, ) else: self.assertEqual( hete_edge[0].edge_label_index[edge_type].shape[1], hete.edge_label_index[edge_type].shape[1], ) self.assertEqual( hete_edge[1].edge_label_index[edge_type].shape[1], hete.edge_label_index[edge_type].shape[1], ) self.assertEqual( hete_edge[2].edge_label_index[edge_type].shape[1], hete.edge_label_index[edge_type].shape[1], ) # link_pred hete_link = hete.split(task='link_pred', split_ratio=[0.5, 0.3, 0.2]) for key, val in hete.edge_label_index.items(): num_edges = val.shape[1] edge_0 = int(0.5 * num_edges) edge_1 = int(0.3 * num_edges) edge_2 = num_edges - edge_0 - edge_1 self.assertEqual(hete_link[0].edge_label[key].shape[0], edge_0) self.assertEqual(hete_link[1].edge_label[key].shape[0], edge_1) self.assertEqual(hete_link[2].edge_label[key].shape[0], edge_2) # undirected G G = generate_dense_hete_graph(directed=False) hete = HeteroGraph(G) hete = HeteroGraph(node_feature=hete.node_feature, node_label=hete.node_label, edge_feature=hete.edge_feature, edge_label=hete.edge_label, edge_index=hete.edge_index, directed=False) # node hete_node = hete.split() for node_type in hete.node_label_index: num_nodes = len(hete.node_label_index[node_type]) node_0 = int(num_nodes * 0.8) node_1 = int(num_nodes * 0.1) node_2 = num_nodes - node_0 - node_1 self.assertEqual( len(hete_node[0].node_label_index[node_type]), node_0, ) self.assertEqual( len(hete_node[1].node_label_index[node_type]), node_1, ) self.assertEqual( len(hete_node[2].node_label_index[node_type]), node_2, ) # node with specified split type node_split_types = ['n1'] hete_node = hete.split(split_types=node_split_types) for node_type in hete.node_label_index: if node_type in node_split_types: num_nodes = len(hete.node_label_index[node_type]) node_0 = int(num_nodes * 0.8) node_1 = int(num_nodes * 0.1) node_2 = num_nodes - node_0 - node_1 self.assertEqual( len(hete_node[0].node_label_index[node_type]), node_0, ) self.assertEqual( len(hete_node[1].node_label_index[node_type]), node_1, ) self.assertEqual( len(hete_node[2].node_label_index[node_type]), node_2, ) else: self.assertEqual( len(hete_node[0].node_label_index[node_type]), len(hete.node_label_index[node_type]), ) self.assertEqual( len(hete_node[1].node_label_index[node_type]), len(hete.node_label_index[node_type]), ) self.assertEqual( len(hete_node[2].node_label_index[node_type]), len(hete.node_label_index[node_type]), ) # edge hete_edge = hete.split(task='edge') for edge_type in hete.edge_label_index: num_edges = int(hete.num_edges(edge_type)) edge_0 = int(num_edges * 0.8) edge_1 = int(num_edges * 0.1) edge_2 = num_edges - edge_0 - edge_1 self.assertEqual( hete_edge[0].edge_label_index[edge_type].shape[1], edge_0, ) self.assertEqual( hete_edge[1].edge_label_index[edge_type].shape[1], edge_1, ) self.assertEqual( hete_edge[2].edge_label_index[edge_type].shape[1], edge_2, ) # edge with specified split type edge_split_types = [('n1', 'e1', 'n1'), ('n1', 'e2', 'n2')] hete_edge = hete.split(task='edge', split_types=edge_split_types) for edge_type in hete.edge_label_index: if edge_type in edge_split_types: num_edges = int(hete.num_edges(edge_type)) edge_0 = int(num_edges * 0.8) edge_1 = int(num_edges * 0.1) edge_2 = num_edges - edge_0 - edge_1 self.assertEqual( hete_edge[0].edge_label_index[edge_type].shape[1], edge_0, ) self.assertEqual( hete_edge[1].edge_label_index[edge_type].shape[1], edge_1, ) self.assertEqual( hete_edge[2].edge_label_index[edge_type].shape[1], edge_2, ) else: self.assertEqual( hete_edge[0].edge_label_index[edge_type].shape[1], hete.edge_label_index[edge_type].shape[1], ) self.assertEqual( hete_edge[1].edge_label_index[edge_type].shape[1], hete.edge_label_index[edge_type].shape[1], ) self.assertEqual( hete_edge[2].edge_label_index[edge_type].shape[1], hete.edge_label_index[edge_type].shape[1], ) # link_pred hete_link = hete.split(task='link_pred', split_ratio=[0.5, 0.3, 0.2]) for key, val in hete.edge_label_index.items(): num_edges = int(val.shape[1] / 2) edge_0 = 2 * int(0.5 * num_edges) edge_1 = 2 * int(0.3 * num_edges) edge_2 = 2 * (num_edges - int(0.5 * num_edges) - int(0.3 * num_edges)) self.assertEqual(hete_link[0].edge_label[key].shape[0], edge_0) self.assertEqual(hete_link[1].edge_label[key].shape[0], edge_1) self.assertEqual(hete_link[2].edge_label[key].shape[0], edge_2)