Repository: pyg-team/pytorch_geometric
Pull request #9554 (2024-07-31 18:33:29)

Add `normalize_edge_index` function to utils
Normalizes the edge index of a graph for use in GNNs.

This function can add self-loops to the graph and apply symmetric or asymmetric normalization based on the node degrees. The function returns the new edge index and the edge weights.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7478077ff..6f7c41ffe 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 `utils.normalize_edge_index` for symmetric/asymmetric normalization of graph edges ([#9554](https://github.com/pyg-team/pytorch_geometric/pull/9554))
 - Added the `RemoveSelfLoops` transformation ([#9562](https://github.com/pyg-team/pytorch_geometric/pull/9562))
 - Added ONNX export for `scatter` with min/max reductions ([#9587](https://github.com/pyg-team/pytorch_geometric/pull/9587))
 - Added a `residual` option in `GATConv` and `GATv2Conv` ([#9515](https://github.com/pyg-team/pytorch_geometric/pull/9515))
diff --git a/torch_geometric/utils/__init__.py b/torch_geometric/utils/__init__.py
index f57908e4d..ecae6cc1c 100644
--- a/torch_geometric/utils/__init__.py
+++ b/torch_geometric/utils/__init__.py
@@ -21,6 +21,7 @@ from ._subgraph import (get_num_hops, subgraph, k_hop_subgraph,
 from .dropout import dropout_adj, dropout_node, dropout_edge, dropout_path
 from ._homophily import homophily
 from ._assortativity import assortativity
+from ._normalize_edge_index import normalize_edge_index
 from .laplacian import get_laplacian
 from .mesh_laplacian import get_mesh_laplacian
 from .mask import mask_select, index_to_mask, mask_to_index
@@ -89,6 +90,7 @@ __all__ = [
     'dropout_adj',
     'homophily',
     'assortativity',
+    'normalize_edge_index',
     'get_laplacian',
     'get_mesh_laplacian',
     'mask_select',
diff --git a/torch_geometric/utils/_normalize_edge_index.py b/torch_geometric/utils/_normalize_edge_index.py
new file mode 100644
index 000000000..add31b4d5
--- /dev/null
+++ b/torch_geometric/utils/_normalize_edge_index.py
@@ -0,0 +1,46 @@
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor
+
+from torch_geometric.utils import add_self_loops as add_self_loops_fn
+from torch_geometric.utils import degree
+
+
+def normalize_edge_index(
+    edge_index: Tensor,
+    num_nodes: Optional[int] = None,
+    add_self_loops: bool = True,
+    symmetric: bool = True,
+) -> Tuple[Tensor, Tensor]:
+    """Applies normalization to the edges of a graph.
+
+    This function can add self-loops to the graph and apply either symmetric or
+    asymmetric normalization based on the node degrees.
+
+    Args:
+        edge_index (LongTensor): The edge indices.
+        num_nodes (int, int], optional): The number of nodes, *i.e.*
+            :obj:`max_val + 1` of :attr:`edge_index`. (default: :obj:`None`)
+        add_self_loops (bool, optional): If set to :obj:`False`, will not add
+            self-loops to the input graph. (default: :obj:`True`)
+        symmetric (bool, optional):  If set to :obj:`True`, symmetric
+            normalization (:math:`D^{-1/2} A D^{-1/2}`) is used, otherwise
+            asymmetric normalization (:math:`D^{-1} A`).
+    """
+    if add_self_loops:
+        edge_index, _ = add_self_loops_fn(edge_index, num_nodes=num_nodes)
+
+    row, col = edge_index[0], edge_index[1]
+    deg = degree(row, num_nodes, dtype=torch.get_default_dtype())
+
+    if symmetric:  # D^-1/2 * A * D^-1/2
+        deg_inv_sqrt = deg.pow(-0.5)
+        deg_inv_sqrt[torch.isinf(deg_inv_sqrt)] = 0
+        edge_weight = deg_inv_sqrt[row] * deg_inv_sqrt[col]
+    else:  # D^-1 * A
+        deg_inv = deg.pow(-1)
+        deg_inv[torch.isinf(deg_inv)] = 0
+        edge_weight = deg_inv[row]
+
+    return edge_index, edge_weight


Test patch:
diff --git a/test/utils/test_normalize_edge_index.py b/test/utils/test_normalize_edge_index.py
new file mode 100644
index 000000000..9ddd1df0b
--- /dev/null
+++ b/test/utils/test_normalize_edge_index.py
@@ -0,0 +1,27 @@
+import pytest
+import torch
+
+from torch_geometric.utils import normalize_edge_index
+
+
+@pytest.mark.parametrize('add_self_loops', [False, True])
+@pytest.mark.parametrize('symmetric', [False, True])
+def test_normalize_edge_index(add_self_loops: bool, symmetric: bool):
+    edge_index = torch.tensor([[0, 2, 2, 3], [2, 0, 3, 0]])
+
+    out = normalize_edge_index(
+        edge_index,
+        add_self_loops=add_self_loops,
+        symmetric=symmetric,
+    )
+    assert isinstance(out, tuple) and len(out) == 2
+    if not add_self_loops:
+        assert out[0].equal(edge_index)
+    else:
+        assert out[0].tolist() == [
+            [0, 2, 2, 3, 0, 1, 2, 3],
+            [2, 0, 3, 0, 0, 1, 2, 3],
+        ]
+
+    assert out[1].min() >= 0.0
+    assert out[1].min() <= 1.0