add data.Dictionary & data.Graph.match

Author: KiddoZhu

test/data/test_graph.py

@@ -36,6 +36,16 @@ def block_diag(self, tensors):
             y += num_col
         return result
 
+    def assert_equal(self, graph1, graph2, prompt):
+        self.assertTrue(torch.equal(graph1.adjacency.to_dense(), graph2.adjacency.to_dense()),
+                        "Incorrect edge list in %s" % prompt)
+        if hasattr(graph1, "node_feature") and hasattr(graph2, "node_feature"):
+            self.assertTrue(torch.equal(graph1.node_feature, graph2.node_feature), "Incorrect feature in %s" % prompt)
+        if hasattr(graph1, "edge_feature") and hasattr(graph2, "edge_feature"):
+            self.assertTrue(torch.equal(graph1.edge_feature, graph2.edge_feature), "Incorrect feature in %s" % prompt)
+        if hasattr(graph1, "graph_feature") and hasattr(graph2, "graph_feature"):
+            self.assertTrue(torch.equal(graph1.graph_feature, graph2.graph_feature), "Incorrect feature in %s" % prompt)
+
     def test_type_cast(self):
         dense_edge_feature = torch.zeros(self.num_node, self.num_node, self.num_feature)
         dense_edge_feature[tuple(self.edge_list.t())] = self.edge_feature
@@ -44,12 +54,8 @@ def test_type_cast(self):
                             node_feature=self.node_feature.tolist(), edge_feature=self.edge_feature.tolist())
         graph2 = data.Graph(self.edge_list.numpy(), self.edge_weight.numpy(), self.num_node,
                             node_feature=self.node_feature.numpy(), edge_feature=self.edge_feature.numpy())
-        self.assertTrue(torch.equal(graph.edge_list, graph1.edge_list), "Incorrect type cast")
-        self.assertTrue(torch.equal(graph.edge_feature, graph1.edge_feature), "Incorrect type cast")
-        self.assertTrue(torch.equal(graph1.edge_list, graph2.edge_list), "Incorrect type cast")
-        self.assertTrue(torch.equal(graph1.edge_weight, graph2.edge_weight), "Incorrect type cast")
-        self.assertTrue(torch.equal(graph1.node_feature, graph2.node_feature), "Incorrect type cast")
-        self.assertTrue(torch.equal(graph1.edge_feature, graph2.edge_feature), "Incorrect type cast")
+        self.assert_equal(graph, graph1, "type cast")
+        self.assert_equal(graph, graph2, "type cast")
 
     def test_index(self):
         graph = data.Graph(self.edge_list, self.edge_weight, self.num_node,
@@ -58,7 +64,7 @@ def test_index(self):
         index = tuple(torch.randint(self.num_node, (2,)).tolist())
         result = graph[index]
         truth = self.adjacency[index]
-        self.assertTrue(torch.equal(result, truth), "Incorrect index in single item")
+        self.assertTrue(torch.equal(result, truth), "Incorrect edge in single item")
 
         h_index = torch.randperm(self.num_node)[:self.num_node // 2]
         t_index = torch.randperm(self.num_node)[:self.num_node // 2]
@@ -71,7 +77,7 @@ def test_index(self):
         adj_truth[not_h_index, :] = 0
         adj_truth[:, not_t_index] = 0
         feat_truth = self.node_feature
-        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect index in node mask")
+        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect edge list in node mask")
         self.assertTrue(torch.equal(feat_result, feat_truth), "Incorrect feature in node mask")
 
         new_graph = graph[:, 1: -1]
@@ -80,7 +86,7 @@ def test_index(self):
         adj_truth = torch.zeros_like(self.adjacency)
         adj_truth[:, 1: -1] = self.adjacency[:, 1: -1]
         feat_truth = self.node_feature
-        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect index in slice")
+        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect edge list in slice")
         self.assertTrue(torch.equal(feat_result, feat_truth), "Incorrect feature in slice")
 
         index = torch.randperm(self.num_node)[:self.num_node // 2]
@@ -89,7 +95,7 @@ def test_index(self):
         feat_result = new_graph.node_feature
         adj_truth = self.adjacency[index][:, index]
         feat_truth = self.node_feature[index]
-        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect index in subgraph")
+        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect edge list in subgraph")
         self.assertTrue(torch.equal(feat_result, feat_truth), "Incorrect feature in subgraph")
 
     def test_device(self):
@@ -100,11 +106,7 @@ def test_device(self):
         self.assertEqual(graph1.adjacency.device.type, "cuda", "Incorrect device")
         graph2 = graph1.cpu()
         self.assertEqual(graph2.adjacency.device.type, "cpu", "Incorrect device")
-        self.assertTrue(torch.equal(graph.adjacency.to_dense(), graph2.adjacency.to_dense()),
-                        "Incorrect feature when changing device")
-        self.assertTrue(torch.equal(graph.node_feature, graph2.node_feature), "Incorrect feature when changing device")
-        self.assertTrue(torch.equal(graph.edge_feature, graph2.edge_feature), "Incorrect feature when changing device")
-        self.assertTrue(torch.equal(graph.graph_feature, graph2.graph_feature), "Incorrect feature when changing device")
+        self.assert_equal(graph, graph2, "device")
 
     def test_pack(self):
         graph = data.Graph(self.edge_list, self.edge_weight, self.num_node,
@@ -128,26 +130,15 @@ def test_pack(self):
         edge_feat_truth = torch.cat([graph.edge_feature for graph in graphs])
         graph_feat_result = packed_graph.graph_feature
         graph_feat_truth = torch.stack([graph.graph_feature for graph in graphs])
-        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect index in pack")
+        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect edge list in pack")
         self.assertTrue(torch.equal(node_feat_result, node_feat_truth), "Incorrect feature in pack")
         self.assertTrue(torch.equal(edge_feat_result, edge_feat_truth), "Incorrect feature in pack")
         self.assertTrue(torch.equal(graph_feat_result, graph_feat_truth), "Incorrect feature in pack")
 
         new_graphs = packed_graph.unpack()
         self.assertEqual(len(graphs), len(new_graphs), "Incorrect length in unpack")
         for graph, new_graph in zip(graphs, new_graphs):
-            adj_truth = graph.adjacency.to_dense()
-            adj_result = new_graph.adjacency.to_dense()
-            node_feat_truth = graph.node_feature
-            node_feat_result = new_graph.node_feature
-            edge_feat_truth = graph.edge_feature
-            edge_feat_result = new_graph.edge_feature
-            graph_feat_truth = graph.graph_feature
-            graph_feat_result = new_graph.graph_feature
-            self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect index in unpack")
-            self.assertTrue(torch.equal(node_feat_result, node_feat_truth), "Incorrect feature in unpack")
-            self.assertTrue(torch.equal(edge_feat_result, edge_feat_truth), "Incorrect feature in unpack")
-            self.assertTrue(torch.equal(graph_feat_result, graph_feat_truth), "Incorrect feature in unpack")
+            self.assert_equal(graph, new_graph, "unpack")
 
         graph = data.Graph(self.edge_list, self.edge_weight, self.num_node,
                            node_feature=self.node_feature, edge_feature=self.edge_feature)
@@ -158,28 +149,12 @@ def test_pack(self):
         for start in range(4):
             mask[start * self.num_node + start: (start + 1) * self.num_node] = 1
         packed_graph2 = packed_graph2.subgraph(mask)
-        adj_result = packed_graph2.adjacency.to_dense()
-        adj_truth = packed_graph.adjacency.to_dense()
-        node_feat_result = packed_graph2.node_feature
-        node_feat_truth = packed_graph.node_feature
-        edge_feat_result = packed_graph2.edge_feature
-        edge_feat_truth = packed_graph.edge_feature
-        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect index in subgraph")
-        self.assertTrue(torch.equal(node_feat_result, node_feat_truth), "Incorrect feature in subgraph")
-        self.assertTrue(torch.equal(edge_feat_result, edge_feat_truth), "Incorrect feature in subgraph")
+        self.assert_equal(packed_graph, packed_graph2, "subgraph")
 
         packed_graph = data.Graph.pack(graphs[::2])
         packed_graph2 = data.Graph.pack(graphs)[::2]
-        adj_result = packed_graph2.adjacency.to_dense()
-        adj_truth = packed_graph.adjacency.to_dense()
-        node_feat_result = packed_graph2.node_feature
-        node_feat_truth = packed_graph.node_feature
-        edge_feat_result = packed_graph2.edge_feature
-        edge_feat_truth = packed_graph.edge_feature
         self.assertEqual(len(packed_graph), len(packed_graph2), "Incorrect batch size in graph mask")
-        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect index in graph mask")
-        self.assertTrue(torch.equal(node_feat_result, node_feat_truth), "Incorrect feature in graph mask")
-        self.assertTrue(torch.equal(edge_feat_result, edge_feat_truth), "Incorrect feature in graph mask")
+        self.assert_equal(packed_graph, packed_graph2, "graph mask")
 
     def test_reorder(self):
         graph = data.Graph(self.edge_list, self.edge_weight, self.num_node,
@@ -192,9 +167,9 @@ def test_reorder(self):
         node_feat_truth = graph.node_feature[order]
         edge_feat_result = new_graph.edge_feature
         edge_feat_truth = graph.edge_feature
-        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect node reorder")
-        self.assertTrue(torch.equal(node_feat_result, node_feat_truth), "Incorrect node reorder")
-        self.assertTrue(torch.equal(edge_feat_result, edge_feat_truth), "Incorrect node reorder")
+        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect edge list in node reorder")
+        self.assertTrue(torch.equal(node_feat_result, node_feat_truth), "Incorrect feature in node reorder")
+        self.assertTrue(torch.equal(edge_feat_result, edge_feat_truth), "Incorrect feature in node reorder")
 
         order = torch.randperm(graph.num_edge)
         new_graph = graph.edge_mask(order)
@@ -204,9 +179,9 @@ def test_reorder(self):
         node_feat_truth = graph.node_feature
         edge_feat_result = new_graph.edge_feature
         edge_feat_truth = graph.edge_feature[order]
-        self.assertTrue(torch.equal(edge_result, edge_truth), "Incorrect edge reorder")
-        self.assertTrue(torch.equal(node_feat_result, node_feat_truth), "Incorrect edge reorder")
-        self.assertTrue(torch.equal(edge_feat_result, edge_feat_truth), "Incorrect edge reorder")
+        self.assertTrue(torch.equal(edge_result, edge_truth), "Incorrect edge list in edge reorder")
+        self.assertTrue(torch.equal(node_feat_result, node_feat_truth), "Incorrect feature in edge reorder")
+        self.assertTrue(torch.equal(edge_feat_result, edge_feat_truth), "Incorrect feature in edge reorder")
 
         graphs = []
         for start in range(4):
@@ -216,30 +191,14 @@ def test_reorder(self):
         order = torch.randperm(4)
         packed_graph = packed_graph.subbatch(order)
         packed_graph2 = data.Graph.pack([graphs[i] for i in order])
-        adj_result = packed_graph.adjacency.to_dense()
-        adj_truth = packed_graph2.adjacency.to_dense()
-        node_feat_result = packed_graph.node_feature
-        node_feat_truth = packed_graph2.node_feature
-        edge_feat_result = packed_graph.edge_feature
-        edge_feat_truth = packed_graph2.edge_feature
-        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect graph reorder")
-        self.assertTrue(torch.equal(node_feat_result, node_feat_truth), "Incorrect graph reorder")
-        self.assertTrue(torch.equal(edge_feat_result, edge_feat_truth), "Incorrect graph reorder")
+        self.assert_equal(packed_graph, packed_graph2, "graph reorder")
 
     def test_repeat(self):
         graph = data.Graph(self.edge_list, self.edge_weight, self.num_node,
                            node_feature=self.node_feature, edge_feature=self.edge_feature)
         repeat_graph = graph.repeat(5)
         true_graph = data.Graph.pack([graph] * 5)
-        adj_result = repeat_graph.adjacency.to_dense()
-        adj_truth = true_graph.adjacency.to_dense()
-        node_feat_result = repeat_graph.node_feature
-        node_feat_truth = true_graph.node_feature
-        edge_feat_result = repeat_graph.edge_feature
-        edge_feat_truth = true_graph.edge_feature
-        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect index in repeat")
-        self.assertTrue(torch.equal(node_feat_result, node_feat_truth), "Incorrect feature in repeat")
-        self.assertTrue(torch.equal(edge_feat_result, edge_feat_truth), "Incorrect feature in repeat")
+        self.assert_equal(repeat_graph, true_graph, "repeat")
 
         # special case: graphs with no edges
         graphs = [graph.edge_mask([]), graph.edge_mask([])]
@@ -249,15 +208,7 @@ def test_repeat(self):
         packed_graph = data.Graph.pack(graphs)
         repeat_graph = packed_graph.repeat(5)
         true_graph = data.Graph.pack(graphs * 5)
-        adj_result = repeat_graph.adjacency.to_dense()
-        adj_truth = true_graph.adjacency.to_dense()
-        node_feat_result = repeat_graph.node_feature
-        node_feat_truth = true_graph.node_feature
-        edge_feat_result = repeat_graph.edge_feature
-        edge_feat_truth = true_graph.edge_feature
-        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect index in repeat")
-        self.assertTrue(torch.equal(node_feat_result, node_feat_truth), "Incorrect feature in repeat")
-        self.assertTrue(torch.equal(edge_feat_result, edge_feat_truth), "Incorrect feature in repeat")
+        self.assert_equal(repeat_graph, true_graph, "repeat")
 
     def test_repeat_interleave(self):
         graph = data.Graph(self.edge_list, self.edge_weight, self.num_node,
@@ -275,15 +226,20 @@ def test_repeat_interleave(self):
         for i, graph in zip(repeats, graphs):
             true_graphs += [graph] * i
         true_graph = data.Graph.pack(true_graphs)
-        adj_result = repeat_graph.adjacency.to_dense()
-        adj_truth = true_graph.adjacency.to_dense()
-        node_feat_result = repeat_graph.node_feature
-        node_feat_truth = true_graph.node_feature
-        edge_feat_result = repeat_graph.edge_feature
-        edge_feat_truth = true_graph.edge_feature
-        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect index in repeat_interleave")
-        self.assertTrue(torch.equal(node_feat_result, node_feat_truth), "Incorrect feature in repeat_interleave")
-        self.assertTrue(torch.equal(edge_feat_result, edge_feat_truth), "Incorrect feature in repeat_interleave")
+        self.assert_equal(repeat_graph, true_graph, "repeat interleave")
+
+    def test_repeated_index(self):
+        graph = data.Graph(self.edge_list, self.edge_weight, self.num_node)
+        graphs = []
+        for start in range(4):
+            index = torch.arange(start, self.num_node)
+            graphs.append(graph.subgraph(index))
+        packed_graph = data.Graph.pack(graphs)
+        # special case: some indexes missing, not sorted
+        index = [1, 0, 2, 1, 0]
+        packed_graph = packed_graph[index]
+        packed_graph2 = data.Graph.pack([graphs[i] for i in index])
+        self.assert_equal(packed_graph, packed_graph2, "repeated index")
 
     def test_split(self):
         graph = data.Graph(self.edge_list, self.edge_weight, self.num_node)
@@ -352,11 +308,29 @@ def test_directed(self):
         graph = digraph.undirected()
         adj_result = graph.adjacency.to_dense()
         adj_truth = (digraph.adjacency + digraph.adjacency.t()).to_dense()
-        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect undirected graph")
+        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect conversion from directed to undirected")
         digraph2 = graph.directed()
         adj_result = digraph2.adjacency.to_dense()
         adj_truth = adj_truth.triu()
-        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect directed graph")
+        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect conversion from undirected to directed")
+
+    def test_match(self):
+        graph = data.Graph(self.edge_list, self.edge_weight, self.num_node)
+        index = torch.randperm(graph.num_edge)[:self.num_node]
+        edge = graph.edge_list[index]
+        mask = torch.randint(2, (len(edge), 1))
+        edge.scatter_(1, mask, -1)
+        random = torch.randint_like(edge, self.num_node)
+        edge = torch.cat([edge, random])
+        index_result, num_match_result = graph.match(edge)
+        index_results = index_result.split(num_match_result.tolist())
+        match = ((graph.edge_list.unsqueeze(0) == edge.unsqueeze(1)) | (edge.unsqueeze(1) == -1)).all(dim=-1)
+        query_index, index_truth = match.nonzero().t()
+        num_match_truth = torch.bincount(query_index, minlength=len(edge))
+        index_truths = index_truth.split(num_match_truth.tolist())
+        self.assertTrue(torch.equal(num_match_result, num_match_truth), "Incorrect edge match")
+        for index_result, index_truth in zip(index_results, index_truths):
+            self.assertTrue(torch.equal(index_result.sort()[0], index_truth.sort()[0]), "Incorrect edge match")
 
 
 if __name__ == "__main__":