Repository: pyg-team/pytorch_geometric
Pull request #7927 (2023-08-24 12:48:20)

Added a test with `to_hetero` for `Sequential` models


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4946196c0..fdf2ce5f8 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.ppr` for personalized PageRank computation ([#7917](https://github.com/pyg-team/pytorch_geometric/pull/7917))
 - Added support for XPU device in `PrefetchLoader` ([#7918](https://github.com/pyg-team/pytorch_geometric/pull/7918))
 - Added support for floating-point slicing in `Dataset`, *e.g.*, `dataset[:0.9]` ([#7915](https://github.com/pyg-team/pytorch_geometric/pull/7915))
 - Added nightly GPU tests ([#7895](https://github.com/pyg-team/pytorch_geometric/pull/7895))
diff --git a/torch_geometric/transforms/gdc.py b/torch_geometric/transforms/gdc.py
index 348be9fe6..e5c600dcc 100644
--- a/torch_geometric/transforms/gdc.py
+++ b/torch_geometric/transforms/gdc.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, Tuple
 
 import numpy as np
 import torch
@@ -11,11 +11,12 @@ from torch_geometric.transforms import BaseTransform
 from torch_geometric.utils import (
     add_self_loops,
     coalesce,
+    get_ppr,
     is_undirected,
     scatter,
+    sort_edge_index,
     to_dense_adj,
 )
-from torch_geometric.utils.sparse import index2ptr
 
 
 @functional_transform('gdc')
@@ -83,9 +84,6 @@ class GDC(BaseTransform):
                                                      avg_degree=64),
         exact: bool = True,
     ):
-
-        self.__calc_ppr__ = get_calc_ppr()
-
         self.self_loop_weight = self_loop_weight
         self.normalization_in = normalization_in
         self.normalization_out = normalization_out
@@ -303,20 +301,16 @@ class GDC(BaseTransform):
                 _, col = edge_index
                 deg = scatter(edge_weight, col, 0, num_nodes, reduce='sum')
 
-            edge_index_np = edge_index.cpu().numpy()
-
-            # Assumes sorted and coalesced edge indices:
-            indptr = index2ptr(edge_index[0], num_nodes).cpu().numpy()
-            out_degree = indptr[1:] - indptr[:-1]
+            edge_index, edge_weight = get_ppr(
+                edge_index,
+                alpha=kwargs['alpha'],
+                eps=kwargs['eps'],
+                num_nodes=num_nodes,
+            )
 
-            neighbors, neighbor_weights = self.__calc_ppr__(
-                indptr, edge_index_np[1], out_degree, kwargs['alpha'],
-                kwargs['eps'])
-            ppr_normalization = 'col' if normalization == 'col' else 'row'
-            edge_index, edge_weight = self.__neighbors_to_graph__(
-                neighbors, neighbor_weights, ppr_normalization,
-                device=edge_index.device)
-            edge_index = edge_index.to(torch.long)
+            if normalization == 'col':
+                edge_index, edge_weight = sort_edge_index(
+                    edge_index.flip([0]), edge_weight, num_nodes)
 
             if normalization == 'sym':
                 # We can change the normalization from row-normalized to
@@ -501,104 +495,3 @@ class GDC(BaseTransform):
         left = sorted_edges[avg_degree * num_nodes - 1]
         right = sorted_edges[avg_degree * num_nodes]
         return (left + right) / 2.0
-
-    def __neighbors_to_graph__(
-        self,
-        neighbors: List[List[int]],
-        neighbor_weights: List[List[float]],
-        normalization: str = 'row',
-        device: torch.device = 'cpu',
-    ) -> Tuple[Tensor, Tensor]:
-        r"""Combine a list of neighbors and neighbor weights to create a sparse
-        graph.
-
-        Args:
-            neighbors (List[List[int]]): List of neighbors for each node.
-            neighbor_weights (List[List[float]]): List of weights for the
-                neighbors of each node.
-            normalization (str): Normalization of resulting matrix
-                (options: :obj:`"row"`, :obj:`"col"`). (default: :obj:`"row"`)
-            device (torch.device): Device to create output tensors on.
-                (default: :obj:`"cpu"`)
-
-        :rtype: (:class:`LongTensor`, :class:`Tensor`)
-        """
-        edge_weight = torch.from_numpy(np.concatenate(neighbor_weights))
-        edge_weight = edge_weight.to(device, torch.get_default_dtype())
-        i = np.repeat(np.arange(len(neighbors)),
-                      np.fromiter(map(len, neighbors), dtype=int))
-        j = np.concatenate(neighbors)
-        if normalization == 'col':
-            edge_index = torch.from_numpy(np.vstack([j, i])).to(device)
-            N = len(neighbors)
-            edge_index, edge_weight = coalesce(edge_index, edge_weight, N, N)
-        elif normalization == 'row':
-            edge_index = torch.from_numpy(np.vstack([i, j])).to(device)
-        else:
-            raise ValueError(
-                f"PPR matrix normalization {normalization} unknown.")
-        return edge_index, edge_weight
-
-
-def get_calc_ppr():
-    import numba
-
-    @numba.jit(nopython=True, parallel=True)
-    def calc_ppr(
-        indptr: np.ndarray,
-        indices: np.ndarray,
-        out_degree: np.ndarray,
-        alpha: float,
-        eps: float,
-    ) -> Tuple[List[List[int]], List[List[float]]]:
-        r"""Calculate the personalized PageRank vector for all nodes
-        using a variant of the Andersen algorithm
-        (see Andersen et al. :Local Graph Partitioning using PageRank Vectors.)
-
-        Args:
-            indptr (np.ndarray): Index pointer for the sparse matrix
-                (CSR-format).
-            indices (np.ndarray): Indices of the sparse matrix entries
-                (CSR-format).
-            out_degree (np.ndarray): Out-degree of each node.
-            alpha (float): Alpha of the PageRank to calculate.
-            eps (float): Threshold for PPR calculation stopping criterion
-                (:obj:`edge_weight >= eps * out_degree`).
-
-        :rtype: (:class:`List[List[int]]`, :class:`List[List[float]]`)
-        """
-
-        alpha_eps = alpha * eps
-        js = [[0]] * len(out_degree)
-        vals = [[0.]] * len(out_degree)
-        for inode_uint in numba.prange(len(out_degree)):
-            inode = numba.int64(inode_uint)
-            p = {inode: 0.0}
-            r = {}
-            r[inode] = alpha
-            q = [inode]
-            while len(q) > 0:
-                unode = q.pop()
-
-                res = r[unode] if unode in r else 0
-                if unode in p:
-                    p[unode] += res
-                else:
-                    p[unode] = res
-                r[unode] = 0
-                for vnode in indices[indptr[unode]:indptr[unode + 1]]:
-                    _val = (1 - alpha) * res / out_degree[unode]
-                    if vnode in r:
-                        r[vnode] += _val
-                    else:
-                        r[vnode] = _val
-
-                    res_vnode = r[vnode] if vnode in r else 0
-                    if res_vnode >= alpha_eps * out_degree[vnode]:
-                        if vnode not in q:
-                            q.append(vnode)
-            js[inode] = list(p.keys())
-            vals[inode] = list(p.values())
-        return js, vals
-
-    return calc_ppr
diff --git a/torch_geometric/utils/__init__.py b/torch_geometric/utils/__init__.py
index 9609d13c5..46cc541ae 100644
--- a/torch_geometric/utils/__init__.py
+++ b/torch_geometric/utils/__init__.py
@@ -51,6 +51,7 @@ 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
+from .ppr import get_ppr
 from .train_test_split_edges import train_test_split_edges
 
 __all__ = [
@@ -135,6 +136,7 @@ __all__ = [
     'tree_decomposition',
     'get_embeddings',
     'trim_to_layer',
+    'get_ppr',
     'train_test_split_edges',
 ]
 
diff --git a/torch_geometric/utils/ppr.py b/torch_geometric/utils/ppr.py
new file mode 100644
index 000000000..c3e4b34a2
--- /dev/null
+++ b/torch_geometric/utils/ppr.py
@@ -0,0 +1,138 @@
+from itertools import chain
+from typing import Callable, List, Optional, Tuple
+
+import numpy as np
+import torch
+from torch import Tensor
+
+from torch_geometric.utils import is_torch_sparse_tensor, to_torch_csr_tensor
+from torch_geometric.utils.num_nodes import maybe_num_nodes
+
+try:
+    import numba
+    WITH_NUMBA = True
+except ImportError:  # pragma: no cover
+    WITH_NUMBA = False
+
+
+def _get_ppr(  # pragma: no cover
+    rowptr: np.ndarray,
+    col: np.ndarray,
+    alpha: float,
+    eps: float,
+    target: Optional[np.ndarray] = None,
+) -> Tuple[List[List[int]], List[List[float]]]:
+
+    num_nodes = len(rowptr) - 1 if target is None else len(target)
+    alpha_eps = alpha * eps
+    js = [[0]] * num_nodes
+    vals = [[0.]] * num_nodes
+
+    for inode_uint in numba.prange(num_nodes):
+        if target is None:
+            inode = numba.int64(inode_uint)
+        else:
+            inode = target[inode_uint]
+
+        p = {inode: 0.0}
+        r = {}
+        r[inode] = alpha
+        q = [inode]
+
+        while len(q) > 0:
+            unode = q.pop()
+
+            res = r[unode] if unode in r else 0
+            if unode in p:
+                p[unode] += res
+            else:
+                p[unode] = res
+
+            r[unode] = 0
+            start, end = rowptr[unode], rowptr[unode + 1]
+            ucount = end - start
+
+            for vnode in col[start:end]:
+                _val = (1 - alpha) * res / ucount
+                if vnode in r:
+                    r[vnode] += _val
+                else:
+                    r[vnode] = _val
+
+                res_vnode = r[vnode] if vnode in r else 0
+                vcount = rowptr[vnode + 1] - rowptr[vnode]
+                if res_vnode >= alpha_eps * vcount:
+                    if vnode not in q:
+                        q.append(vnode)
+
+        js[inode_uint] = list(p.keys())
+        vals[inode_uint] = list(p.values())
+
+    return js, vals
+
+
+_get_ppr_numba: Optional[Callable] = None
+
+
+def get_ppr(
+    edge_index: Tensor,
+    alpha: float = 0.2,
+    eps: float = 1e-5,
+    target: Optional[Tensor] = None,
+    num_nodes: Optional[int] = None,
+) -> Tuple[Tensor, Tensor]:
+    r"""Calculates the personalized PageRank (PPR) vector for all or a subset
+    of nodes using a variant of the `Andersen algorithm
+    <https://mathweb.ucsd.edu/~fan/wp/localpartition.pdf>`_.
+
+    Args:
+        edge_index (torch.Tensor): The indices of the graph.
+        alpha (float, optional): The alpha value of the PageRank algorithm.
+            (default: :obj:`0.2`)
+        eps (float, optional): The threshold for stopping the PPR calculation
+            (:obj:`edge_weight >= eps * out_degree`). (default: :obj:`1e-5`)
+        target (torch.Tensor, optional): The target nodes to compute PPR for.
+            If not given, calculates PPR vectors for all nodes.
+            (default: :obj:`None`)
+        num_nodes (int, optional): The number of nodes. (default: :obj:`None`)
+
+    :rtype: (:class:`torch.Tensor`, :class:`torch.Tensor`)
+    """
+    if not WITH_NUMBA:  # pragma: no cover
+        raise ImportError("'get_ppr' requires the 'numba' package")
+
+    global _get_ppr_numba
+    if _get_ppr_numba is None:
+        _get_ppr_numba = numba.jit(nopython=True, parallel=True)(_get_ppr)
+
+    num_nodes = maybe_num_nodes(edge_index, num_nodes)
+
+    if not is_torch_sparse_tensor(edge_index):
+        size = (num_nodes, num_nodes)
+        edge_index = to_torch_csr_tensor(edge_index, size=size)
+    else:
+        edge_index = edge_index.to_sparse_csr()
+
+    assert edge_index.layout == torch.sparse_csr
+
+    rowptr, col = edge_index.crow_indices(), edge_index.col_indices()
+
+    cols, weights = _get_ppr_numba(
+        rowptr.cpu().numpy(),
+        col.cpu().numpy(),
+        alpha,
+        eps,
+        None if target is None else target.cpu().numpy(),
+    )
+
+    device = edge_index.device
+    col = torch.tensor(list(chain.from_iterable(cols)), device=device)
+    weight = torch.tensor(list(chain.from_iterable(weights)), device=device)
+    deg = torch.tensor([len(value) for value in cols], device=device)
+
+    row = torch.arange(num_nodes) if target is None else target
+    row = row.repeat_interleave(deg, output_size=col.numel())
+
+    edge_index = torch.stack([row, col], dim=0)
+
+    return edge_index, weight


Test patch:
diff --git a/test/nn/test_sequential.py b/test/nn/test_sequential.py
index 086e194c3..43adf1b27 100644
--- a/test/nn/test_sequential.py
+++ b/test/nn/test_sequential.py
@@ -9,8 +9,10 @@ from torch_geometric.nn import (
     GCNConv,
     JumpingKnowledge,
     MessagePassing,
+    SAGEConv,
     Sequential,
     global_mean_pool,
+    to_hetero,
 )
 from torch_geometric.typing import SparseTensor
 
@@ -142,3 +144,33 @@ def test_sequential_with_ordered_dict():
 
     x = model(x, edge_index)
     assert x.size() == (4, 64)
+
+
+def test_sequential_to_hetero():
+    model = Sequential('x, edge_index', [
+        (SAGEConv((-1, -1), 32), 'x, edge_index -> x1'),
+        ReLU(),
+        (SAGEConv((-1, -1), 64), 'x1, edge_index -> x2'),
+        ReLU(),
+    ])
+
+    x_dict = {
+        'paper': torch.randn(100, 16),
+        'author': torch.randn(100, 16),
+    }
+    edge_index_dict = {
+        ('paper', 'cites', 'paper'):
+        torch.randint(100, (2, 200), dtype=torch.long),
+        ('paper', 'written_by', 'author'):
+        torch.randint(100, (2, 200), dtype=torch.long),
+        ('author', 'writes', 'paper'):
+        torch.randint(100, (2, 200), dtype=torch.long),
+    }
+    metadata = list(x_dict.keys()), list(edge_index_dict.keys())
+
+    model = to_hetero(model, metadata, debug=False)
+
+    out_dict = model(x_dict, edge_index_dict)
+    assert isinstance(out_dict, dict) and len(out_dict) == 2
+    assert out_dict['paper'].size() == (100, 64)
+    assert out_dict['author'].size() == (100, 64)
diff --git a/test/utils/test_ppr.py b/test/utils/test_ppr.py
new file mode 100644
index 000000000..ca1a268d7
--- /dev/null
+++ b/test/utils/test_ppr.py
@@ -0,0 +1,28 @@
+import pytest
+import torch
+
+from torch_geometric.datasets import KarateClub
+from torch_geometric.testing import withPackage
+from torch_geometric.utils import get_ppr
+
+
+@withPackage('numba')
+@pytest.mark.parametrize('target', [None, torch.tensor([0, 4, 5, 6])])
+def test_get_ppr(target):
+    data = KarateClub()[0]
+
+    edge_index, edge_weight = get_ppr(
+        data.edge_index,
+        alpha=0.1,
+        eps=1e-5,
+        target=target,
+    )
+
+    assert edge_index.size(0) == 2
+    assert edge_index.size(1) == edge_weight.numel()
+
+    min_row = 0 if target is None else target.min()
+    max_row = data.num_nodes - 1 if target is None else target.max()
+    assert edge_index[0].min() == min_row and edge_index[0].max() == max_row
+    assert edge_index[1].min() >= 0 and edge_index[1].max() < data.num_nodes
+    assert edge_weight.min() >= 0.0 and edge_weight.max() <= 1.0