Only add true negatives in add_random_edge augmentation (#7654)

Fixes #7653

Author: rusty1s

CHANGELOG.md

@@ -70,6 +70,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Changed `add_random_edge` to only add true negative edges ([#7654](https://github.com/pyg-team/pytorch_geometric/pull/7654))
 - Allowed the usage of `BasicGNN` models in `DeepGraphInfomax` ([#7648](https://github.com/pyg-team/pytorch_geometric/pull/7648))
 - Breaking Change: Made `Data.keys` a method rather than a property ([#7629](https://github.com/pyg-team/pytorch_geometric/pull/7629))
 - Added a `num_edges` parameter to the forward method of `HypergraphConv` ([#7560](https://github.com/pyg-team/pytorch_geometric/pull/7560))

test/utils/test_augmentation.py

@@ -1,6 +1,7 @@
 import pytest
 import torch
 
+from torch_geometric import seed_everything
 from torch_geometric.utils import (
     add_random_edge,
     is_undirected,
@@ -77,28 +78,26 @@ def test_add_random_edge():
     assert out[0].tolist() == edge_index.tolist()
     assert out[1].tolist() == [[], []]
 
-    torch.manual_seed(5)
+    seed_everything(5)
     out = add_random_edge(edge_index, p=0.5)
-    assert out[0].tolist() == [[0, 1, 1, 2, 2, 3, 3, 2, 3],
-                               [1, 0, 2, 1, 3, 2, 1, 2, 2]]
-
-    assert out[1].tolist() == [[3, 2, 3], [1, 2, 2]]
+    assert out[0].tolist() == [[0, 1, 1, 2, 2, 3, 3, 1, 2],
+                               [1, 0, 2, 1, 3, 2, 0, 3, 0]]
+    assert out[1].tolist() == [[3, 1, 2], [0, 3, 0]]
 
-    torch.manual_seed(6)
+    seed_everything(6)
     out = add_random_edge(edge_index, p=0.5, force_undirected=True)
-    assert out[0].tolist() == [[0, 1, 1, 2, 2, 3, 1, 2],
-                               [1, 0, 2, 1, 3, 2, 2, 1]]
-    assert out[1].tolist() == [[1, 2], [2, 1]]
+    assert out[0].tolist() == [[0, 1, 1, 2, 2, 3, 1, 3],
+                               [1, 0, 2, 1, 3, 2, 3, 1]]
+    assert out[1].tolist() == [[1, 3], [3, 1]]
     assert is_undirected(out[0])
     assert is_undirected(out[1])
 
-    # test with bipartite graph
-    torch.manual_seed(7)
+    # Test for bipartite graph:
+    seed_everything(7)
     edge_index = torch.tensor([[0, 1, 2, 3, 4, 5], [2, 3, 1, 4, 2, 1]])
-    with pytest.raises(RuntimeError,
-                       match="not supported for heterogeneous graphs"):
-        out = add_random_edge(edge_index, p=0.5, force_undirected=True,
-                              num_nodes=(6, 5))
+    with pytest.raises(RuntimeError, match="not supported for bipartite"):
+        add_random_edge(edge_index, force_undirected=True, num_nodes=(6, 5))
     out = add_random_edge(edge_index, p=0.5, num_nodes=(6, 5))
-    out[0].tolist() == [[0, 1, 2, 3, 4, 5, 3, 4, 1],
-                        [2, 3, 1, 4, 2, 1, 1, 3, 2]]
+    assert out[0].tolist() == [[0, 1, 2, 3, 4, 5, 2, 0, 2],
+                               [2, 3, 1, 4, 2, 1, 0, 4, 2]]
+    assert out[1].tolist() == [[2, 0, 2], [0, 4, 2]]

torch_geometric/utils/__init__.py

@@ -6,7 +6,6 @@
 from .degree import degree
 from .softmax import softmax
 from .dropout import dropout_adj, dropout_node, dropout_edge, dropout_path
-from .augmentation import shuffle_node, mask_feature, add_random_edge
 from .sort_edge_index import sort_edge_index
 from .coalesce import coalesce
 from .undirected import is_undirected, to_undirected
@@ -47,6 +46,7 @@
 from .negative_sampling import (negative_sampling, batched_negative_sampling,
                                 structured_negative_sampling,
                                 structured_negative_sampling_feasible)
+from .augmentation import shuffle_node, mask_feature, add_random_edge
 from .tree_decomposition import tree_decomposition
 from .embedding import get_embeddings
 from .trim_to_layer import trim_to_layer
@@ -62,9 +62,6 @@
     'dropout_edge',
     'dropout_path',
     'dropout_adj',
-    'shuffle_node',
-    'mask_feature',
-    'add_random_edge',
     'sort_edge_index',
     'coalesce',
     'is_undirected',
@@ -130,6 +127,9 @@
     'batched_negative_sampling',
     'structured_negative_sampling',
     'structured_negative_sampling_feasible',
+    'shuffle_node',
+    'mask_feature',
+    'add_random_edge',
     'tree_decomposition',
     'get_embeddings',
     'trim_to_layer',

torch_geometric/utils/augmentation.py

@@ -3,12 +3,14 @@
 import torch
 from torch import Tensor
 
-from torch_geometric.utils import scatter
-from torch_geometric.utils.num_nodes import maybe_num_nodes
+from torch_geometric.utils import negative_sampling, scatter
 
 
-def shuffle_node(x: Tensor, batch: Optional[Tensor] = None,
-                 training: bool = True) -> Tuple[Tensor, Tensor]:
+def shuffle_node(
+    x: Tensor,
+    batch: Optional[Tensor] = None,
+    training: bool = True,
+) -> Tuple[Tensor, Tensor]:
     r"""Randomly shuffle the feature matrix :obj:`x` along the
     first dimmension.
 
@@ -67,9 +69,13 @@ def shuffle_node(x: Tensor, batch: Optional[Tensor] = None,
     return x[perm], perm
 
 
-def mask_feature(x: Tensor, p: float = 0.5, mode: str = 'col',
-                 fill_value: float = 0.,
-                 training: bool = True) -> Tuple[Tensor, Tensor]:
+def mask_feature(
+    x: Tensor,
+    p: float = 0.5,
+    mode: str = 'col',
+    fill_value: float = 0.,
+    training: bool = True,
+) -> Tuple[Tensor, Tensor]:
     r"""Randomly masks feature from the feature matrix
     :obj:`x` with probability :obj:`p` using samples from
     a Bernoulli distribution.
@@ -149,9 +155,13 @@ def mask_feature(x: Tensor, p: float = 0.5, mode: str = 'col',
     return x, mask
 
 
-def add_random_edge(edge_index, p: float, force_undirected: bool = False,
-                    num_nodes: Optional[Union[Tuple[int], int]] = None,
-                    training: bool = True) -> Tuple[Tensor, Tensor]:
+def add_random_edge(
+    edge_index,
+    p: float = 0.5,
+    force_undirected: bool = False,
+    num_nodes: Optional[Union[int, Tuple[int, int]]] = None,
+    training: bool = True,
+) -> Tuple[Tensor, Tensor]:
     r"""Randomly adds edges to :obj:`edge_index`.
 
     The method returns (1) the retained :obj:`edge_index`, (2) the added
@@ -160,6 +170,7 @@ def add_random_edge(edge_index, p: float, force_undirected: bool = False,
     Args:
         edge_index (LongTensor): The edge indices.
         p (float): Ratio of added edges to the existing edges.
+            (default: :obj:`0.5`)
         force_undirected (bool, optional): If set to :obj:`True`,
             added edges will be undirected.
             (default: :obj:`False`)
@@ -208,30 +219,24 @@ def add_random_edge(edge_index, p: float, force_undirected: bool = False,
                 [1, 3, 2]])
     """
     if p < 0. or p > 1.:
-        raise ValueError(f'Ratio of added edges has to be between 0 and 1 '
-                         f'(got {p}')
+        raise ValueError(f"Ratio of added edges has to be between 0 and 1 "
+                         f"(got '{p}')")
     if force_undirected and isinstance(num_nodes, (tuple, list)):
-        raise RuntimeError('`force_undirected` is not supported for'
-                           ' heterogeneous graphs')
+        raise RuntimeError("'force_undirected' is not supported for "
+                           "bipartite graphs")
 
     device = edge_index.device
     if not training or p == 0.0:
         edge_index_to_add = torch.tensor([[], []], device=device)
         return edge_index, edge_index_to_add
 
-    if not isinstance(num_nodes, (tuple, list)):
-        num_nodes = (num_nodes, num_nodes)
-    num_src_nodes = maybe_num_nodes(edge_index, num_nodes[0])
-    num_dst_nodes = maybe_num_nodes(edge_index, num_nodes[1])
-
-    num_edges_to_add = round(edge_index.size(1) * p)
-    row = torch.randint(0, num_src_nodes, size=(num_edges_to_add, ))
-    col = torch.randint(0, num_dst_nodes, size=(num_edges_to_add, ))
+    edge_index_to_add = negative_sampling(
+        edge_index=edge_index,
+        num_nodes=num_nodes,
+        num_neg_samples=round(edge_index.size(1) * p),
+        force_undirected=force_undirected,
+    )
 
-    if force_undirected:
-        mask = row < col
-        row, col = row[mask], col[mask]
-        row, col = torch.cat([row, col]), torch.cat([col, row])
-    edge_index_to_add = torch.stack([row, col], dim=0).to(device)
     edge_index = torch.cat([edge_index, edge_index_to_add], dim=1)
+
     return edge_index, edge_index_to_add

torch_geometric/utils/negative_sampling.py

@@ -9,11 +9,13 @@
 from torch_geometric.utils.num_nodes import maybe_num_nodes
 
 
-def negative_sampling(edge_index: Tensor,
-                      num_nodes: Optional[Union[int, Tuple[int, int]]] = None,
-                      num_neg_samples: Optional[int] = None,
-                      method: str = "sparse",
-                      force_undirected: bool = False) -> Tensor:
+def negative_sampling(
+    edge_index: Tensor,
+    num_nodes: Optional[Union[int, Tuple[int, int]]] = None,
+    num_neg_samples: Optional[int] = None,
+    method: str = "sparse",
+    force_undirected: bool = False,
+) -> Tensor:
     r"""Samples random negative edges of a graph given by :attr:`edge_index`.
 
     Args: