VirtualNode transform (#4163)

* add virtualnode impl

* add test

* typo

test/transforms/test_virtual_node.py

@@ -0,0 +1,38 @@
+import torch
+
+from torch_geometric.data import Data
+from torch_geometric.transforms import VirtualNode
+
+
+def test_virtual_node():
+    assert str(VirtualNode()) == 'VirtualNode()'
+
+    x = torch.randn(4, 16)
+    edge_index = torch.tensor([[2, 0, 2], [3, 1, 0]])
+    edge_weight = torch.rand(edge_index.size(1))
+    edge_attr = torch.randn(edge_index.size(1), 8)
+
+    data = Data(x=x, edge_index=edge_index, edge_weight=edge_weight,
+                edge_attr=edge_attr, num_nodes=x.size(0))
+
+    data = VirtualNode()(data)
+    assert len(data) == 6
+
+    assert data.x.size() == (5, 16)
+    assert torch.allclose(data.x[:4], x)
+    assert data.x[4:].abs().sum() == 0
+
+    assert data.edge_index.tolist() == [[2, 0, 2, 0, 1, 2, 3, 4, 4, 4, 4],
+                                        [3, 1, 0, 4, 4, 4, 4, 0, 1, 2, 3]]
+
+    assert data.edge_weight.size() == (11, )
+    assert torch.allclose(data.edge_weight[:3], edge_weight)
+    assert data.edge_weight[3:].abs().sum() == 8
+
+    assert data.edge_attr.size() == (11, 8)
+    assert torch.allclose(data.edge_attr[:3], edge_attr)
+    assert data.edge_attr[3:].abs().sum() == 0
+
+    assert data.num_nodes == 5
+
+    assert data.edge_type.tolist() == [0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]

torch_geometric/transforms/__init__.py

@@ -47,6 +47,7 @@
 from .random_link_split import RandomLinkSplit
 from .add_metapaths import AddMetaPaths
 from .largest_connected_components import LargestConnectedComponents
+from .virtual_node import VirtualNode
 
 __all__ = [
     'BaseTransform',
@@ -98,6 +99,7 @@
     'RandomLinkSplit',
     'AddMetaPaths',
     'LargestConnectedComponents',
+    'VirtualNode',
 ]
 
 classes = __all__

torch_geometric/transforms/add_self_loops.py

@@ -41,6 +41,3 @@ def __call__(self, data: Union[Data, HeteroData]):
             setattr(store, self.attr, edge_weight)
 
         return data
-
-    def __repr__(self) -> str:
-        return f'{self.__class__.__name__}()'

torch_geometric/transforms/virtual_node.py

@@ -0,0 +1,64 @@
+import torch
+from torch import Tensor
+
+from torch_geometric.data import Data
+from torch_geometric.transforms import BaseTransform
+
+
+class VirtualNode(BaseTransform):
+    r"""Appends a virtual node to the given homogeneous graph that is connected
+    to all other nodes, as described in the `"Neural Message Passing for
+    Quantum Chemistry" <https://arxiv.org/abs/1704.01212>`_ paper.
+    The virtual node serves as a global scratch space that each node both reads
+    from and writes to in every step of message passing.
+    This allows information to travel long distances during the propagation
+    phase.
+
+    Node and edge features of the virtual node are added as zero-filled input
+    features.
+    Furthermore, special edge types will be added both for in-coming and
+    out-going information to and from the virtual node.
+    """
+    def __call__(self, data: Data) -> Data:
+        num_nodes, (row, col) = data.num_nodes, data.edge_index
+        edge_type = data.get('edge_type', torch.zeros_like(row))
+
+        arange = torch.arange(num_nodes, device=row.device)
+        full = row.new_full((num_nodes, ), num_nodes)
+        row = torch.cat([row, arange, full], dim=0)
+        col = torch.cat([col, full, arange], dim=0)
+        edge_index = torch.stack([row, col], dim=0)
+
+        new_type = edge_type.new_full((num_nodes, ), int(edge_type.max()) + 1)
+        edge_type = torch.cat([edge_type, new_type, new_type + 1], dim=0)
+
+        for key, value in data.items():
+            if key == 'edge_index' or key == 'edge_type':
+                continue
+
+            if isinstance(value, Tensor):
+                dim = data.__cat_dim__(key, value)
+                size = list(value.size())
+
+                fill_value = None
+                if key == 'edge_weight':
+                    size[dim] = 2 * num_nodes
+                    fill_value = 1.
+                elif data.is_edge_attr(key):
+                    size[dim] = 2 * num_nodes
+                    fill_value = 0.
+                elif data.is_node_attr(key):
+                    size[dim] = 1
+                    fill_value = 0.
+
+                if fill_value is not None:
+                    new_value = value.new_full(size, fill_value)
+                    data[key] = torch.cat([value, new_value], dim=dim)
+
+        data.edge_index = edge_index
+        data.edge_type = edge_type
+
+        if 'num_nodes' in data:
+            data.num_nodes = data.num_nodes + 1
+
+        return data