Repository: pyg-team/pytorch_geometric
Pull request #7827 (2023-08-01 09:16:54)

Added `HalfHop` graph upsampling augmentation
This PR adds the `HalfHop` graph upsampling augmentation from the ICML 2023 paper [Half-Hop: A graph upsampling approach for slowing down message passing](https://openreview.net/forum?id=lXczFIwQkv).

According to the authors, `HalfHop` works by adding artificial slow nodes between neighbors to slow down message propagation, see below.

![image](https://github.com/pyg-team/pytorch_geometric/assets/39986668/62448ab7-3d83-422a-82bc-6397dab55909)

Reference: 
+ https://github.com/nerdslab/halfhop

cc: @mazabou

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7827?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7827](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7827?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (86e2247) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/8c1339c1dcc84c4ee7a6c36d64c93bc085d244e9?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (8c1339c) will **increase** coverage by `0.00%`.
> The diff coverage is `94.87%`.

> :exclamation: Current head 86e2247 differs from pull request most recent head 5084b15. Consider uploading reports for the commit 5084b15 to get more accurate results

```diff
@@           Coverage Diff           @@
##           master    #7827   +/-   ##
=======================================
  Coverage   89.82%   89.82%           
=======================================
  Files         455      456    +1     
  Lines       26327    26366   +39     
=======================================
+ Hits        23647    23683   +36     
- Misses       2680     2683    +3     
```


| [Files Changed](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7827?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/transforms/half\_hop.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7827?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3RyYW5zZm9ybXMvaGFsZl9ob3AucHk=) | `94.73% <94.73%> (ø)` | |
| [torch\_geometric/transforms/\_\_init\_\_.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7827?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3RyYW5zZm9ybXMvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |

... and [1 file with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7827/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 456be23a1..43b13ca3b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added the `HalfHop` graph upsampling augmentation ([#7827](https://github.com/pyg-team/pytorch_geometric/pull/7827))
 - Added the `Wikidata5M` dataset ([#7864](https://github.com/pyg-team/pytorch_geometric/pull/7864))
 - Added TorchScript support inside `BasicGNN` models ([#7865](https://github.com/pyg-team/pytorch_geometric/pull/7865))
 - Added a `batch_size` argument to `unbatch` functionalities ([#7851](https://github.com/pyg-team/pytorch_geometric/pull/7851))
diff --git a/torch_geometric/transforms/__init__.py b/torch_geometric/transforms/__init__.py
index aef8b7913..703219424 100644
--- a/torch_geometric/transforms/__init__.py
+++ b/torch_geometric/transforms/__init__.py
@@ -37,6 +37,7 @@ from .largest_connected_components import LargestConnectedComponents
 from .virtual_node import VirtualNode
 from .add_positional_encoding import AddLaplacianEigenvectorPE, AddRandomWalkPE
 from .feature_propagation import FeaturePropagation
+from .half_hop import HalfHop
 
 from .distance import Distance
 from .cartesian import Cartesian
@@ -106,6 +107,7 @@ graph_transforms = [
     'AddLaplacianEigenvectorPE',
     'AddRandomWalkPE',
     'FeaturePropagation',
+    'HalfHop',
 ]
 
 vision_transforms = [
diff --git a/torch_geometric/transforms/half_hop.py b/torch_geometric/transforms/half_hop.py
new file mode 100644
index 000000000..a03df145e
--- /dev/null
+++ b/torch_geometric/transforms/half_hop.py
@@ -0,0 +1,97 @@
+import torch
+
+from torch_geometric.data import Data
+from torch_geometric.data.datapipes import functional_transform
+from torch_geometric.transforms import BaseTransform
+
+
+@functional_transform('half_hop')
+class HalfHop(BaseTransform):
+    r"""Graph upsampling augmentation from the
+    `"Half-Hop: A graph upsampling approach for slowing down
+    message passing" <https://openreview.net/forum?id=lXczFIwQkv>`_ paper.
+    The graph is augmented by adding artificial slow nodes between neighbors
+    to slow down message propagation. (functional name: :obj:`half_hop`).
+
+    .. note::
+        `HalfHop` augmentation is not supported if `data` has
+        :attr:`edge_weight` or :attr:`edge_attr`.
+
+    Args:
+        alpha (float, optional): The interpolation factor
+            used to compute slow node features
+            :math:`x = \alpha*x_src + (1-\alpha)*x_dst` (default: :obj:`0.5`)
+        p (float, optional): The probability of half-hopping
+            an edge. (default: :obj:`1.0`)
+
+    .. code-block:: python
+
+        import torch_geometric.transforms as T
+
+        data = ...
+        data = T.HalfHop(alpha=0.5)(data) # Apply transformation
+        out = model(data.x, data.edge_index) # Feed-forward
+        out = out[~data.slow_node_mask] # Get rid of slow nodes
+    """
+    def __init__(self, alpha: float = 0.5, p: float = 1.0):
+        if alpha < 0. or alpha > 1.:
+            raise ValueError(f'Interpolation factor has to be between 0 and 1 '
+                             f'(got {alpha}')
+        if p < 0. or p > 1.:
+            raise ValueError(
+                f'Ratio of half-hopped edges has to be between 0 and 1 '
+                f'(got {p}')
+        self.p = p
+        self.alpha = alpha
+
+    def forward(self, data: Data) -> Data:
+
+        if data.edge_weight is not None or data.edge_attr is not None:
+            raise ValueError("'HalfHop' augmentation is not supported if "
+                             "'data' has 'edge_weight' or 'edge_attr'")
+
+        x, edge_index = data.x, data.edge_index
+
+        # isolate self loops which are not half-hopped
+        self_loop_mask = edge_index[0] == edge_index[1]
+        edge_index_self_loop = edge_index[:, self_loop_mask]
+        edge_index = edge_index[:, ~self_loop_mask]
+
+        # randomly sample nodes and half-hop their edges
+        node_mask = torch.rand(data.num_nodes, device=x.device) < self.p
+        edge_mask = node_mask[edge_index[1]]
+        edge_index_to_halfhop = edge_index[:, edge_mask]
+        edge_index_to_keep = edge_index[:, ~edge_mask]
+
+        # add new slow nodes of which features are initialized
+        # by linear interpolation
+        num_halfhop_edges = edge_index_to_halfhop.size(1)
+        slow_node_ids = torch.arange(num_halfhop_edges,
+                                     device=x.device) + data.num_nodes
+        x_src = x[edge_index_to_halfhop[0]]
+        x_dst = x[edge_index_to_halfhop[1]]
+        x_slow_node = self.alpha * x_src + (1 - self.alpha) * x_dst
+        new_x = torch.cat([x, x_slow_node], dim=0)
+
+        # add new edges between slow nodes and the original nodes
+        edge_index_slow = [
+            torch.stack([edge_index_to_halfhop[0], slow_node_ids]),
+            torch.stack([slow_node_ids, edge_index_to_halfhop[1]]),
+            torch.stack([edge_index_to_halfhop[1], slow_node_ids])
+        ]
+        new_edge_index = torch.cat(
+            [edge_index_to_keep, edge_index_self_loop, *edge_index_slow],
+            dim=1)
+
+        # prepare a mask that distinguishes between original nodes & slow nodes
+        slow_node_mask = torch.cat(
+            [x.new_zeros(x.size(0)),
+             x.new_ones(slow_node_ids.size(0))], dim=0).bool()
+
+        data.x, data.edge_index = new_x, new_edge_index
+        data.slow_node_mask = slow_node_mask
+
+        return data
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(alpha={self.alpha}, p={self.p})'


Test patch:
diff --git a/test/transforms/test_half_hop.py b/test/transforms/test_half_hop.py
new file mode 100644
index 000000000..cc0fc6652
--- /dev/null
+++ b/test/transforms/test_half_hop.py
@@ -0,0 +1,46 @@
+import torch
+
+from torch_geometric.data import Data
+from torch_geometric.transforms import HalfHop
+
+
+def test_half_hop():
+    edge_index = torch.tensor([[0, 1, 1, 2, 0, 1, 2], [1, 0, 2, 1, 0, 1, 2]])
+    x = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]],
+                     dtype=torch.float)
+    data = Data(x=x, edge_index=edge_index)
+
+    transform = HalfHop()
+    assert str(transform) == 'HalfHop(alpha=0.5, p=1.0)'
+    data = transform(data)
+
+    expected_edge_index = [[0, 1, 2, 0, 1, 1, 2, 3, 4, 5, 6, 1, 0, 2, 1],
+                           [0, 1, 2, 3, 4, 5, 6, 1, 0, 2, 1, 3, 4, 5, 6]]
+    expected_x = torch.tensor(
+        [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [3, 4, 5, 6],
+         [3, 4, 5, 6], [7, 8, 9, 10], [7, 8, 9, 10]], dtype=torch.float)
+    assert len(data) == 3
+    assert data.num_nodes == 7
+    assert data.edge_index.tolist() == expected_edge_index
+    assert torch.allclose(data.x, expected_x, atol=1e-4)
+    assert data.slow_node_mask.tolist() == [
+        False, False, False, True, True, True, True
+    ]
+
+    torch.manual_seed(1)
+    data = Data(x=x, edge_index=edge_index)
+    transform = HalfHop(p=0.5)
+    assert str(transform) == 'HalfHop(alpha=0.5, p=0.5)'
+    data = transform(data)
+
+    expected_edge_index = [[1, 0, 1, 2, 0, 1, 2, 3, 4, 5, 1, 2, 1],
+                           [0, 0, 1, 2, 3, 4, 5, 1, 2, 1, 3, 4, 5]]
+    expected_x = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12],
+                               [3, 4, 5, 6], [7, 8, 9, 10], [7, 8, 9, 10]],
+                              dtype=torch.float)
+    assert data.num_nodes == 6
+    assert data.edge_index.tolist() == expected_edge_index
+    assert torch.allclose(data.x, expected_x, atol=1e-4)
+    assert data.slow_node_mask.tolist() == [
+        False, False, False, True, True, True
+    ]