Repository: pyg-team/pytorch_geometric
Pull request #7192 (2023-04-17 13:36:09)

Add `AddRemainingSelfLoops` transform
This PR adds the `AddRemainingSelfLoops` transform to generate only remaining self-loops instead of generating them on all the vertices like in the `AddSelfLoops` transform.

codecov[bot]: ## [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7192?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7192](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7192?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (5383ae3) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/8999e507602d45ba9b53cc94f501e320cb579a92?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (8999e50) will **decrease** coverage by `0.45%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master    #7192      +/-   ##
==========================================
- Coverage   91.72%   91.28%   -0.45%     
==========================================
  Files         436      436              
  Lines       23922    23924       +2     
==========================================
- Hits        21942    21838     -104     
- Misses       1980     2086     +106     
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7192?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/add\_self\_loops.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7192?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3RyYW5zZm9ybXMvYWRkX3NlbGZfbG9vcHMucHk=) | `100.00% <100.00%> (ø)` | |

... and [23 files with indirect coverage changes](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7192/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)

wsad1: @bwroblew if its ok, could you please give edit access to `maintainers`, for this and future PRs. 
bwroblew: > @bwroblew if its ok, could you please give edit access to `maintainers`, for this and future PRs.

I'm afraid it's not possible since our fork lives under the `graphcore` organization... I didn't know that limitation but don't see an option here in the PR and just discovered there is an issue for that - https://github.com/orgs/community/discussions/5634. It has not been resolved yet.

I believe @rusty1s have access since was added as an outside collaborator with write access to our fork. Unfortunately, I don't have access to add outside collaborators to the fork, please let me discuss that inside Graphcore and see what we can do about that.
wsad1: @bwroblew ah I see. 
In that case can you resolve conflicts and merge this PR.
rusty1s: Thanks all. Let me take care of merging :)

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b35f93d61..560e87b60 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 a `AddRemainingSelfLoops` transform ([#7192](https://github.com/pyg-team/pytorch_geometric/pull/7192))
 - Added `optimizer_resolver` ([#7209](https://github.com/pyg-team/pytorch_geometric/pull/7209))
 - Added `type_ptr` argument to `HeteroLayerNorm` ([#7208](https://github.com/pyg-team/pytorch_geometric/pull/7208))
 - Added an option to benchmark scripts to write PyTorch profiler results to CSV ([#7114](https://github.com/pyg-team/pytorch_geometric/pull/7114))
diff --git a/torch_geometric/transforms/__init__.py b/torch_geometric/transforms/__init__.py
index f56db80f8..96612f7cc 100644
--- a/torch_geometric/transforms/__init__.py
+++ b/torch_geometric/transforms/__init__.py
@@ -18,6 +18,7 @@ from .one_hot_degree import OneHotDegree
 from .target_indegree import TargetIndegree
 from .local_degree_profile import LocalDegreeProfile
 from .add_self_loops import AddSelfLoops
+from .add_remaining_self_loops import AddRemainingSelfLoops
 from .remove_isolated_nodes import RemoveIsolatedNodes
 from .remove_duplicated_edges import RemoveDuplicatedEdges
 from .knn_graph import KNNGraph
@@ -82,6 +83,7 @@ graph_transforms = [
     'TargetIndegree',
     'LocalDegreeProfile',
     'AddSelfLoops',
+    'AddRemainingSelfLoops',
     'RemoveIsolatedNodes',
     'RemoveDuplicatedEdges',
     'KNNGraph',
diff --git a/torch_geometric/transforms/add_remaining_self_loops.py b/torch_geometric/transforms/add_remaining_self_loops.py
new file mode 100644
index 000000000..74a1470a0
--- /dev/null
+++ b/torch_geometric/transforms/add_remaining_self_loops.py
@@ -0,0 +1,49 @@
+from typing import Optional, Union
+
+from torch import Tensor
+
+from torch_geometric.data import Data, HeteroData
+from torch_geometric.data.datapipes import functional_transform
+from torch_geometric.transforms import BaseTransform
+from torch_geometric.utils import add_remaining_self_loops
+
+
+@functional_transform('add_remaining_self_loops')
+class AddRemainingSelfLoops(BaseTransform):
+    r"""Adds remaining self-loops to the given homogeneous or heterogeneous
+    graph (functional name: :obj:`add_remaining_self_loops`).
+
+    Args:
+        attr (str, optional): The name of the attribute of edge weights
+            or multi-dimensional edge features to pass to
+            :meth:`torch_geometric.utils.add_remaining_self_loops`.
+            (default: :obj:`"edge_weight"`)
+        fill_value (float or Tensor or str, optional): The way to generate
+            edge features of self-loops (in case :obj:`attr != None`).
+            If given as :obj:`float` or :class:`torch.Tensor`, edge features of
+            self-loops will be directly given by :obj:`fill_value`.
+            If given as :obj:`str`, edge features of self-loops are computed by
+            aggregating all features of edges that point to the specific node,
+            according to a reduce operation. (:obj:`"add"`, :obj:`"mean"`,
+            :obj:`"min"`, :obj:`"max"`, :obj:`"mul"`). (default: :obj:`1.`)
+    """
+    def __init__(self, attr: Optional[str] = 'edge_weight',
+                 fill_value: Union[float, Tensor, str] = 1.0):
+        self.attr = attr
+        self.fill_value = fill_value
+
+    def __call__(
+        self,
+        data: Union[Data, HeteroData],
+    ) -> Union[Data, HeteroData]:
+        for store in data.edge_stores:
+            if store.is_bipartite() or 'edge_index' not in store:
+                continue
+
+            store.edge_index, edge_weight = add_remaining_self_loops(
+                store.edge_index, getattr(store, self.attr, None),
+                fill_value=self.fill_value, num_nodes=store.size(0))
+
+            setattr(store, self.attr, edge_weight)
+
+        return data


Test patch:
diff --git a/test/transforms/test_add_remaining_self_loops.py b/test/transforms/test_add_remaining_self_loops.py
new file mode 100644
index 000000000..1de56f5a3
--- /dev/null
+++ b/test/transforms/test_add_remaining_self_loops.py
@@ -0,0 +1,72 @@
+import torch
+
+from torch_geometric.data import Data, HeteroData
+from torch_geometric.transforms import AddRemainingSelfLoops
+
+
+def test_add_remaining_self_loops():
+    assert str(AddRemainingSelfLoops()) == 'AddRemainingSelfLoops()'
+
+    assert len(AddRemainingSelfLoops()(Data())) == 0
+
+    # No self-loops in `edge_index`.
+    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+    edge_weight = torch.tensor([1, 2, 3, 4])
+    edge_attr = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
+
+    data = Data(edge_index=edge_index, num_nodes=3)
+    data = AddRemainingSelfLoops()(data)
+    assert len(data) == 2
+    assert data.edge_index.tolist() == [[0, 1, 1, 2, 0, 1, 2],
+                                        [1, 0, 2, 1, 0, 1, 2]]
+    assert data.num_nodes == 3
+
+    # Single self-loop in `edge_index`.
+    edge_index = torch.tensor([[0, 0, 1, 2], [1, 0, 2, 1]])
+    data = Data(edge_index=edge_index, num_nodes=3)
+    data = AddRemainingSelfLoops()(data)
+    assert len(data) == 2
+    assert data.edge_index.tolist() == [[0, 1, 2, 0, 1, 2], [1, 2, 1, 0, 1, 2]]
+    assert data.num_nodes == 3
+
+    data = Data(edge_index=edge_index, edge_weight=edge_weight, num_nodes=3)
+    data = AddRemainingSelfLoops(attr='edge_weight', fill_value=5)(data)
+    assert data.edge_index.tolist() == [[0, 1, 2, 0, 1, 2], [1, 2, 1, 0, 1, 2]]
+    assert data.num_nodes == 3
+    assert data.edge_weight.tolist() == [1, 3, 4, 2, 5, 5]
+
+    data = Data(edge_index=edge_index, edge_attr=edge_attr, num_nodes=3)
+    data = AddRemainingSelfLoops(attr='edge_attr', fill_value='add')(data)
+    assert data.edge_index.tolist() == [[0, 1, 2, 0, 1, 2], [1, 2, 1, 0, 1, 2]]
+    assert data.num_nodes == 3
+    assert data.edge_attr.tolist() == [[1, 2], [5, 6], [7, 8], [3, 4], [8, 10],
+                                       [5, 6]]
+
+
+def test_add_remaining_self_loops_all_loops_exist():
+    # All self-loops already exist in the data object.
+    edge_index = torch.tensor([[0, 1, 2], [0, 1, 2]])
+    data = Data(edge_index=edge_index, num_nodes=3)
+    data = AddRemainingSelfLoops()(data)
+    assert data.edge_index.tolist() == edge_index.tolist()
+
+    # All self-loops already exist in the data object, some of them appear
+    # multiple times.
+    edge_index = torch.tensor([[0, 0, 1, 1, 2], [0, 0, 1, 1, 2]])
+    data = Data(edge_index=edge_index, num_nodes=3)
+    data = AddRemainingSelfLoops()(data)
+    assert data.edge_index.tolist() == [[0, 1, 2], [0, 1, 2]]
+
+
+def test_hetero_add_remaining_self_loops():
+    edge_index = torch.tensor([[0, 0, 1, 2], [1, 0, 2, 1]])
+
+    data = HeteroData()
+    data['v'].num_nodes = 3
+    data['w'].num_nodes = 3
+    data['v', 'v'].edge_index = edge_index
+    data['v', 'w'].edge_index = edge_index
+    data = AddRemainingSelfLoops()(data)
+    assert data['v', 'v'].edge_index.tolist() == [[0, 1, 2, 0, 1, 2],
+                                                  [1, 2, 1, 0, 1, 2]]
+    assert data['v', 'w'].edge_index.tolist() == edge_index.tolist()
diff --git a/test/transforms/test_add_self_loops.py b/test/transforms/test_add_self_loops.py
index 45455bcda..38a4ed0ff 100644
--- a/test/transforms/test_add_self_loops.py
+++ b/test/transforms/test_add_self_loops.py
@@ -7,12 +7,12 @@ from torch_geometric.transforms import AddSelfLoops
 def test_add_self_loops():
     assert str(AddSelfLoops()) == 'AddSelfLoops()'
 
+    assert len(AddSelfLoops()(Data())) == 0
+
     edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
     edge_weight = torch.tensor([1, 2, 3, 4])
     edge_attr = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
 
-    assert len(AddSelfLoops()(Data())) == 0
-
     data = Data(edge_index=edge_index, num_nodes=3)
     data = AddSelfLoops()(data)
     assert len(data) == 2
@@ -36,6 +36,14 @@ def test_add_self_loops():
                                        [8, 10], [5, 6]]
 
 
+def test_add_self_loops_with_existing_self_loops():
+    edge_index = torch.tensor([[0, 1, 2], [0, 1, 2]])
+    data = Data(edge_index=edge_index, num_nodes=3)
+    data = AddSelfLoops()(data)
+    assert data.edge_index.tolist() == [[0, 1, 2, 0, 1, 2], [0, 1, 2, 0, 1, 2]]
+    assert data.num_nodes == 3
+
+
 def test_hetero_add_self_loops():
     edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])