add node/edge/graph reference & change node/edge to atom/bond

Author: KiddoZhu

README.md

@@ -95,7 +95,7 @@ Molecules are also supported in TorchDrug. You can get the desired molecule
 properties without any domain knowledge.
 
 ```python
-mol = data.Molecule.from_smiles("CCOC(=O)N", node_feature="default", edge_feature="default")
+mol = data.Molecule.from_smiles("CCOC(=O)N", atom_feature="default", bond_feature="default")
 print(mol.node_feature)
 print(mol.atom_type)
 print(mol.to_scaffold())

doc/source/notes/graph.rst

@@ -9,7 +9,7 @@ Create a Graph
 
 To begin with, let's create a graph.
 
-.. code-block:: python
+.. code:: python
 
     import torch
     from torchdrug import data
@@ -25,15 +25,15 @@ This will plot a ring graph like the following.
     :width: 33%
 
 Internally, the graph is stored as a sparse edge list to save memory footprint. For
-an intuitive comparison, a `scale-free graph`_ mayr have 1 million nodes and 10 million
+an intuitive comparison, a `scale-free graph`_ may have 1 million nodes and 10 million
 edges. The dense version takes about 4TB, while the sparse version only requires 120MB.
 
 .. _scale-free graph:
     https://en.wikipedia.org/wiki/Scale-free_network
 
 Here are some commonly used properties of the graph.
 
-.. code-block:: python
+.. code:: python
 
     print(graph.num_node)
     print(graph.num_edge)
@@ -45,7 +45,7 @@ molecules have bond types like ``single bound``, while knowledge graphs have rel
 like ``consists of``. To construct such a relational graph, we can pass the edge type
 as a third variable in the edge list.
 
-.. code-block:: python
+.. code:: python
 
     triplet_list = [[0, 1, 0], [1, 2, 1], [2, 3, 0], [3, 4, 1], [4, 5, 0], [5, 0, 1]]
     graph = data.Graph(triplet_list, num_node=6, num_relation=2)
@@ -62,7 +62,7 @@ corresponds to an edge from node :math:`i` to node :math:`j`. The relational gra
 uses a 3D adjacency matrix :math:`A`, where non-zero :math:`A_{i,j,k}` denotes an
 edge from node :math:`i` to node :math:`j` with edge type :math:`k`.
 
-.. code-block:: python
+.. code:: python
 
     adjacency = torch.zeros(6, 6)
     adjacency[edge_list] = 1
@@ -78,7 +78,7 @@ For example, the following code creates a benzene molecule.
 .. _SMILES:
     https://en.wikipedia.org/wiki/Simplified_molecular-input_line-entry_system
 
-.. code-block:: python
+.. code:: python
 
     mol = data.Molecule.from_smiles("C1=CC=CC=C1")
     mol.visualize()
@@ -90,7 +90,7 @@ For example, the following code creates a benzene molecule.
 Once the graph is created, we can transfer it between CPU and GPUs, just like
 :class:`torch.Tensor`.
 
-.. code-block:: python
+.. code:: python
 
     graph = graph.cuda()
     print(graph.device)
@@ -109,7 +109,7 @@ during any graph operation.
 
 Here we specify some features during the construction of the molecule graph.
 
-.. code-block:: python
+.. code:: python
 
     mol = data.Molecule.from_smiles("C1=CC=CC=C1", node_feature="default",
                                     edge_feature="default", graph_feature="ecfp")
@@ -122,14 +122,15 @@ We may also want to define our own attributes. This only requires to wrap the
 assignment lines with a context manager. The following example defines edge importance
 as the reciprocal of node degrees.
 
-.. code-block:: python
+.. code:: python
 
     node_in, node_out = mol.edge_list.t()[:2]
     with mol.edge():
         mol.edge_importance = 1 / graph.degree_in[node_in] + 1 / graph.degree_out[node_out]
 
 We can use ``mol.node()`` and ``mol.graph()`` for node- and graph-level attributes
-respectively.
+respectively. Attributes may also be a reference to node/edge/graph indexes. See
+:doc:`reference` for more details.
 
 Note in order to support batching and masking, attributes should always have the same
 length as their corresponding components. This means the size of the first dimension of
@@ -142,7 +143,7 @@ Modern deep learning frameworks employs batched operations to accelerate computa
 In TorchDrug, we can easily batch same kind of graphs with **arbitary sizes**. Here
 is an example of creating a batch of 4 graphs.
 
-.. code-block:: python
+.. code:: python
 
     graphs = [graph, graph, graph, graph]
     batch = data.Graph.pack(graphs)
@@ -170,7 +171,7 @@ where :math:`A_i` is the adjacency of :math:`i`-th graph.
 To get a single graph from the batch, use the conventional index or
 :meth:`PackedGraph.unpack <torchdrug.data.PackedGraph.unpack>`.
 
-.. code-block:: python
+.. code:: python
 
     graph = batch[1]
     graphs = batch.unpack()
@@ -186,7 +187,7 @@ Subgraph and Masking
 The graph data structure also provides a bunch of slicing operations to create subgraphs
 or masked graphs in a sparse manner. Some typical operations include
 
-.. code-block:: python
+.. code:: python
 
     g1 = graph.subgraph([1, 2, 3, 4])
     g1.visualize()
@@ -220,7 +221,7 @@ isolated nodes.
 The same operations can also be applied to batches. In this case, we need to convert
 the index of a single graph into the index in a batch.
 
-.. code-block:: python
+.. code:: python
 
     graph_ids = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1, 1, 1])
     node_ids = torch.tensor([1, 2, 3, 4, 0, 1, 2, 3, 4, 5])
@@ -232,7 +233,7 @@ the index of a single graph into the index in a batch.
 
 We can also pick a subset of graphs in a batch.
 
-.. code-block:: python
+.. code:: python
 
     batch = batch[[0, 1]]
     batch.visualize()

doc/source/notes/index.rst

@@ -7,3 +7,4 @@ Notes
     variadic
     layer
     model
+    reference

doc/source/notes/layer.rst

@@ -105,7 +105,7 @@ representations as a graph representation. First, we readout the mean of node
 representations. Second, we broadcast the mean representation to each node to compute
 the difference. Finally, we readout the mean of the squared difference as the variance.
 
-.. code-block:: python
+.. code:: python
 
     from torch import nn
     from torch_scatter import scatter_mean

doc/source/notes/model.rst

@@ -69,7 +69,7 @@ distributed, module-centric manner.
 We compute the variational regularization loss, and add it to the global loss and the
 global metric.
 
-.. code-block::
+.. code::
 
         def reparameterize(self, mu, log_sigma):
             if self.training:

doc/source/notes/reference.rst

@@ -0,0 +1,88 @@
+Deal with References
+====================
+
+As we show in :doc:`graph`, custom graph attributes will be automatically processed
+in any graph operation. However, some attributes may refer to node/edge/graph indexes,
+and their values need to be modified when the indexes change. TorchDrug provides a
+mechanism to deal with such cases.
+
+Inverse Edge Index
+------------------
+
+A typical example of reference is a mapping from each edge to its inverse edge.
+We first prepare an undirected graph with the indexes of inverse edges.
+
+.. code:: python
+
+    import torch
+    from torchdrug import data
+
+    edge_list = [[0, 1], [1, 0], [1, 2], [2, 1], [2, 0], [0, 2]]
+    inv_edge_index = [1, 0, 3, 2, 5, 4]
+    graph = data.Graph(edge_list, num_node=3)
+
+.. image:: ../../../asset/graph/inverse_edge.png
+    :align: center
+    :width: 33%
+
+If we assign the indexes as an edge attribute and apply an edge mask operation,
+the result is not desired. The edges are masked out correctly, but the values of
+inverse indexes are wrong.
+
+.. code:: python
+    with graph.edge():
+        graph.inv_edge_index = torch.tensor(inv_edge_index)
+    g1 = graph.edge_mask([0, 2, 3])
+
+.. image:: ../../../asset/graph/wrong_reference.png
+    :align: center
+    :width: 33%
+
+Instead, we need to explicitly tell TorchDrug that the attribute ``graph.inv_edge_index``
+is a reference to edge indexes. This is done by an additional context manager
+``graph.edge_reference()``. Now we get the correct inverse indexes. Note that missing
+references will be set to ``-1``. In this case, the inverse index of ``0`` is ``-1``,
+since the corresponding inverse edge has been masked out.
+
+.. code:: python
+
+    with graph.edge(), graph.edge_reference():
+        graph.inv_edge_index = torch.tensor(inv_edge_index)
+    g2 = graph.edge_mask([0, 2, 3])
+
+.. image:: ../../../asset/graph/correct_reference.png
+    :align: center
+    :width: 33%
+
+We can use ``graph.node_reference()`` and ``graph.graph_reference()`` for references
+to nodes and graphs respectively.
+
+Use Cases in Proteins
+---------------------
+
+In :class:`data.Protein`, the mapping ``atom2residue`` is implemented as
+references. The intuition is that references enable flexible indexing on either atoms
+or residues, while maintaining the correspondence between two views.
+
+The following example shows how to track a specific residue with ``atom2residue`` in
+the atom view. For a protein, we first create a mask for atoms in a glutamine (GLN).
+
+.. code:: python
+
+    protein = data.Protein.from_sequence("KALKQMLDMG")
+    is_glutamine = protein.residue_type[protein.atom2residue] == protein.residue2id["GLN"]
+    with protein.node():
+        protein.is_glutamine = is_glutamine
+
+We then apply a mask to the protein residue sequence. In the output protein,
+``atom2residue`` is able to map the masked atoms back to the glutamine residue.
+
+.. code:: python
+
+    p1 = protein[3:6]
+    residue_type = p1.residue_type[p1.atom2residue[p1.is_glutamine]]
+    print([p1.id2residue[r] for r in residue_type.tolist()])
+
+.. code:: bash
+
+    ['GLN', 'GLN', 'GLN', 'GLN', 'GLN', 'GLN', 'GLN', 'GLN', 'GLN']

test/data/test_graph.py

@@ -326,12 +326,64 @@ def test_match(self):
         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))
+        num_match_truth = query_index.bincount(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")
 
+    def test_reference(self):
+        node_out = torch.arange(1, self.num_node)
+        node_in = (node_out - 1) // 2
+        edge_list = torch.stack([node_in, node_out], dim=-1)
+        tree = data.Graph(edge_list, num_node=self.num_node)
+        with tree.node(), tree.node_reference():
+            tree.dad = (torch.arange(self.num_node) - 1) // 2
+
+        mask = torch.arange(1, self.num_node)
+        graph = tree.subgraph(mask)
+        degree_in_result = graph.dad[graph.dad != -1].bincount(minlength=graph.num_node)
+        is_root_result = graph.dad == -1
+        node_in, node_out = graph.edge_list.t()
+        degree_in_truth = node_in.bincount(minlength=graph.num_node)
+        is_root_truth = node_out.bincount(minlength=graph.num_node) == 0
+        self.assertTrue(torch.equal(degree_in_result, degree_in_truth), "Incorrect node reference")
+        self.assertTrue(torch.equal(is_root_result, is_root_truth), "Incorrect node reference")
+
+        packed_graph = tree.repeat(4)
+        packed_graph2 = data.Graph.pack([tree] * 4)
+        self.assert_equal(packed_graph, packed_graph2, "node reference")
+
+        # special case: 0 repetition
+        repeats = [2, 0, 1, 2]
+        trees = []
+        for start in range(4):
+            index = torch.arange(start, self.num_node)
+            trees.append(tree.subgraph(index))
+        packed_graph = data.Graph.pack(trees)
+        repeat_graph = packed_graph.repeat_interleave(repeats)
+        true_graphs = []
+        for i, tree in zip(repeats, trees):
+            true_graphs += [tree] * i
+        true_graph = data.Graph.pack(true_graphs)
+        self.assert_equal(repeat_graph, true_graph, "node reference")
+
+    def test_line_graph(self):
+        graph = data.Graph(self.edge_list, self.edge_weight, self.num_node, edge_feature=self.edge_feature)
+        line_graph = graph.line_graph()
+        adj_result = line_graph.adjacency.to_dense()
+        feat_result = line_graph.node_feature
+        edge_index = torch.arange(graph.num_edge)
+        node_in, node_out = graph.edge_list.t()
+        edge2node_out = torch.zeros(graph.num_edge, graph.num_node)
+        node_in2edge = torch.zeros(graph.num_node, graph.num_edge)
+        edge2node_out[edge_index, node_out] = 1
+        node_in2edge[node_in, edge_index] = 1
+        adj_truth = edge2node_out @ node_in2edge
+        feat_truth = graph.edge_feature
+        self.assertTrue(torch.equal(adj_result, adj_truth), "Incorrect line graph")
+        self.assertTrue(torch.equal(feat_result, feat_truth), "Incorrect line graph")
+
 
 if __name__ == "__main__":
     unittest.main()

test/data/test_molecule.py

@@ -33,7 +33,7 @@ def test_smiles(self):
         self.assertTrue((mols.num_edges == 0).all(), "Incorrect SMILES side case")
 
     def test_feature(self):
-        mol = data.Molecule.from_smiles(self.smiles, graph_feature="ecfp")
+        mol = data.Molecule.from_smiles(self.smiles, mol_feature="ecfp")
         self.assertTrue((mol.graph_feature > 0).any(), "Incorrect ECFP feature")
 
 

test/layers/test_variadic.py

@@ -45,6 +45,21 @@ def test_topk(self):
             self.assertTrue(torch.equal(result_value, truth_value), "Incorrect variadic topk")
             self.assertTrue(torch.equal(result_index, truth_index), "Incorrect variadic topk")
 
+        for _ in range(10):
+            k = torch.randint(self.size.min(), self.size.max(), (self.num_graph,))
+            result_value, result_index = functional.variadic_topk(self.input, self.size, k)
+            _truth_value, _truth_index = self.padded.topk(self.size.max(), dim=1)
+            truth_value, truth_index = [], []
+            for i, size in enumerate(self.size):
+                truth_value.append(_truth_value[i, :k[i]])
+                truth_index.append(_truth_index[i, :k[i]])
+                for j in range(size, k[i].item()):
+                    truth_value[i][j] = truth_value[i][j-1]
+                    truth_index[i][j] = truth_index[i][j-1]
+            truth_value = torch.cat(truth_value, dim=0)
+            truth_index = torch.cat(truth_index, dim=0)
+            self.assertTrue(torch.equal(result_value, truth_value), "Incorrect variadic topk")
+            self.assertTrue(torch.equal(result_index, truth_index), "Incorrect variadic topk")
 
 if __name__ == "__main__":
     unittest.main()