Repository: pyg-team/pytorch_geometric
Pull request #6894 (2023-03-10 18:38:21)

[Transforms] Add `NodePropertySplit` transform
New functional for creating node-level data splits, which allows to induce distributional shifts in graph and conduct experiments with graph models in more challenging setups. It is implemented via the `NodePropertySplit` transform and allows to use the class itself or its static method `mask_nodes_by_property` (the latter is for users who already knows their structural node property and just need to perform data splitting).

Test coverage is provided and passed successfully. I decided to edit `transforms.rst` in `docs/source/modules` to make the static method `mask_nodes_by_property` visible in documentation, but not sure if this is the right way.

codecov[bot]: ## [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6894?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#6894](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6894?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (346247f) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/ad2fe06b5f5a69ddcf201e3c21257c9abc92c2e1?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (ad2fe06) will **increase** coverage by `0.04%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master    #6894      +/-   ##
==========================================
+ Coverage   91.53%   91.58%   +0.04%     
==========================================
  Files         427      428       +1     
  Lines       23166    23247      +81     
==========================================
+ Hits        21205    21290      +85     
+ Misses       1961     1957       -4     
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6894?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/\_\_init\_\_.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6894?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%> (ø)` | |
| [torch\_geometric/transforms/node\_property\_split.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6894?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3RyYW5zZm9ybXMvbm9kZV9wcm9wZXJ0eV9zcGxpdC5weQ==) | `100.00% <100.00%> (ø)` | |

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

gvbazhenov: I see that `networkx` is not a necessary dependency for `torch-geometric`. How can I fix the implementation so that the library and documentation build successfully?
rusty1s: You can move the imports of `networkx` to the corresponding functions that make use of it. In this case, you probably need to also replace the `nx.Graph` type hints with something more general such as `Any`.
gvbazhenov: As I can see, test fails on `import networkx as nx` when `NodePropertySplit` instance calls `to_networkx`. This is rather strange, since I just rely on the existing utilities in `utils/convert`. Could you please explain me what I am doing wrong?
gvbazhenov: @EdisonLeeeee @rusty1s thanks for your PR review and help!
gvbazhenov: @rusty1s @EdisonLeeeee excuse me, could you please share when it will be possible to merge this pull request? Is there anything else that requires my participation?
rusty1s: Will take a final look soon, but overall this looks good. Sorry for the delay, I was quite busy with PyG 2.3 release :(
gvbazhenov: > Will take a final look soon, but overall this looks good. Sorry for the delay, I was quite busy with PyG 2.3 release :(

Roger that, thank you!

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f7634c095..e5f57ed66 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 `NodePropertySplit` transform for creating node-level splits using structural node properties ([#6894](https://github.com/pyg-team/pytorch_geometric/pull/6894))
 - Added an option to preserve directed graphs in `CitationFull` datasets ([#7275](https://github.com/pyg-team/pytorch_geometric/pull/7275))
 - Added support for `torch.sparse.Tensor` in `DataLoader` ([#7252](https://github.com/pyg-team/pytorch_geometric/pull/7252))
 - Added `save` and `load` methods to `InMemoryDataset` ([#7250](https://github.com/pyg-team/pytorch_geometric/pull/7250))
diff --git a/torch_geometric/graphgym/model_builder.py b/torch_geometric/graphgym/model_builder.py
index 90e7237fa..246ccfde5 100644
--- a/torch_geometric/graphgym/model_builder.py
+++ b/torch_geometric/graphgym/model_builder.py
@@ -74,7 +74,7 @@ def create_model(to_device=True, dim_in=None, dim_out=None) -> GraphGymModule:
     dim_in = cfg.share.dim_in if dim_in is None else dim_in
     dim_out = cfg.share.dim_out if dim_out is None else dim_out
     # binary classification, output dim = 1
-    if 'classification' in cfg.dataset.task_type and dim_out == 2:
+    if 'classification' == cfg.dataset.task_type and dim_out == 2:
         dim_out = 1
 
     model = GraphGymModule(dim_in, dim_out, cfg)
diff --git a/torch_geometric/transforms/__init__.py b/torch_geometric/transforms/__init__.py
index 96612f7cc..aef8b7913 100644
--- a/torch_geometric/transforms/__init__.py
+++ b/torch_geometric/transforms/__init__.py
@@ -10,6 +10,7 @@ from .svd_feature_reduction import SVDFeatureReduction
 from .remove_training_classes import RemoveTrainingClasses
 from .random_node_split import RandomNodeSplit
 from .random_link_split import RandomLinkSplit
+from .node_property_split import NodePropertySplit
 from .mask import IndexToMask, MaskToIndex
 from .pad import Pad
 
@@ -72,6 +73,7 @@ general_transforms = [
     'RemoveTrainingClasses',
     'RandomNodeSplit',
     'RandomLinkSplit',
+    'NodePropertySplit',
     'IndexToMask',
     'MaskToIndex',
     'Pad',
diff --git a/torch_geometric/transforms/node_property_split.py b/torch_geometric/transforms/node_property_split.py
new file mode 100644
index 000000000..df4a313a5
--- /dev/null
+++ b/torch_geometric/transforms/node_property_split.py
@@ -0,0 +1,165 @@
+from typing import Any, Dict, List
+
+import torch
+from torch import Tensor
+
+from torch_geometric.data import Data
+from torch_geometric.data.datapipes import functional_transform
+from torch_geometric.transforms import BaseTransform
+from torch_geometric.utils import to_networkx
+
+
+@functional_transform('node_property_split')
+class NodePropertySplit(BaseTransform):
+    r"""Creates a node-level split with distributional shift based on a given
+    node property, as proposed in the `"Evaluating Robustness and Uncertainty
+    of Graph Models Under Structural Distributional Shifts"
+    <https://arxiv.org/abs/2302.13875v1>`__ paper
+    (functional name: :obj:`node_property_split`).
+
+    It splits the nodes in a given graph into five non-intersecting parts
+    based on their structural properties.
+    This can be used for transductive node prediction tasks with distributional
+    shifts.
+    It considers the in-distribution (ID) and out-of-distribution (OOD) subsets
+    of nodes.
+    The ID subset includes training, validation and testing parts, while
+    the OOD subset includes validation and testing parts.
+    As a result, it creates five associated node mask vectors for each graph,
+    three which are for the ID nodes (:obj:`id_train_mask`,
+    :obj:`id_val_mask`, :obj:`id_test_mask`), and two which are for the OOD
+    nodes (:obj:`ood_val_mask`, :obj:`ood_test_mask`).
+
+    This class implements three particular strategies for inducing
+    distributional shifts in a graph — based on **popularity**, **locality**
+    or **density**.
+
+    Args:
+        property_name (str): The name of the node property to be used
+            (:obj:`"popularity"`, :obj:`"locality"`, :obj:`"density"`).
+        ratios ([float]): A list of five ratio values for ID training,
+            ID validation, ID test, OOD validation and OOD test parts.
+            The values must sum to :obj:`1.0`.
+        ascending (bool, optional): Whether to sort nodes in ascending order
+            of the node property, so that nodes with greater values of the
+            property are considered to be OOD (default: :obj:`True`)
+
+    Example:
+
+    .. code-block:: python
+
+        from torch_geometric.transforms import NodePropertySplit
+        from torch_geometric.datasets.graph_generator import ERGraph
+
+        data = ERGraph(num_nodes=1000, edge_prob=0.4)()
+
+        property_name = 'popularity'
+        ratios = [0.3, 0.1, 0.1, 0.3, 0.2]
+        tranaform = NodePropertySplit(property_name, ratios)
+
+        data = transform(data)
+    """
+    def __init__(
+        self,
+        property_name: str,
+        ratios: List[float],
+        ascending: bool = True,
+    ):
+        if property_name not in {'popularity', 'locality', 'density'}:
+            raise ValueError(f"Unexpected 'property_name' "
+                             f"(got '{property_name}')")
+
+        if len(ratios) != 5:
+            raise ValueError(f"'ratios' must contain 5 values "
+                             f"(got {len(ratios)})")
+
+        if sum(ratios) != 1.0:
+            raise ValueError(f"'ratios' must sum to 1.0 (got {sum(ratios)})")
+
+        self.property_name = property_name
+        self.compute_fn = _property_name_to_compute_fn[property_name]
+        self.ratios = ratios
+        self.ascending = ascending
+
+    def __call__(self, data: Data) -> Data:
+
+        G = to_networkx(data, to_undirected=True, remove_self_loops=True)
+        property_values = self.compute_fn(G, self.ascending)
+        mask_dict = self._mask_nodes_by_property(property_values, self.ratios)
+
+        for key, mask in mask_dict.items():
+            data[key] = mask
+
+        return data
+
+    @staticmethod
+    def _compute_popularity_property(G: Any, ascending: bool = True) -> Tensor:
+        import networkx.algorithms as A
+
+        property_values = torch.tensor(list(A.pagerank(G).values()))
+        property_values *= -1 if ascending else 1
+        return property_values
+
+    @staticmethod
+    def _compute_locality_property(G: Any, ascending: bool = True) -> Tensor:
+        import networkx.algorithms as A
+
+        pagerank_values = torch.tensor(list(A.pagerank(G).values()))
+
+        num_nodes = G.number_of_nodes()
+        personalization = dict(zip(range(num_nodes), [0.0] * num_nodes))
+        personalization[int(pagerank_values.argmax())] = 1.0
+
+        property_values = torch.tensor(
+            list(A.pagerank(G, personalization=personalization).values()))
+        property_values *= -1 if ascending else 1
+        return property_values
+
+    @staticmethod
+    def _compute_density_property(G: Any, ascending: bool = True) -> Tensor:
+        import networkx.algorithms as A
+
+        property_values = torch.tensor(list(A.clustering(G).values()))
+        property_values *= -1 if ascending else 1
+        return property_values
+
+    @staticmethod
+    def _mask_nodes_by_property(
+        property_values: Tensor,
+        ratios: List[float],
+    ) -> Dict[str, Tensor]:
+
+        num_nodes = property_values.size(0)
+        sizes = (num_nodes * torch.tensor(ratios)).round().long()
+        sizes[-1] -= sizes.sum() - num_nodes
+
+        perm = torch.randperm(num_nodes)
+        id_size = int(sizes[:3].sum())
+        perm = perm[property_values[perm].argsort()]
+        perm[:id_size] = perm[:id_size][torch.randperm(id_size)]
+
+        node_splits = perm.split(sizes.tolist())
+        names = [
+            'id_train_mask',
+            'id_val_mask',
+            'id_test_mask',
+            'ood_val_mask',
+            'ood_test_mask',
+        ]
+
+        split_masks = {}
+        for name, node_split in zip(names, node_splits):
+            split_mask = torch.zeros(num_nodes, dtype=torch.bool)
+            split_mask[node_split] = True
+            split_masks[name] = split_mask
+        return split_masks
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.property_name})'
+
+
+_property_name_to_compute_fn = {
+    'popularity': NodePropertySplit._compute_popularity_property,
+    'locality': NodePropertySplit._compute_locality_property,
+    'density': NodePropertySplit._compute_density_property,
+}


Test patch:
diff --git a/test/transforms/test_node_property_split.py b/test/transforms/test_node_property_split.py
new file mode 100644
index 000000000..27686fbc3
--- /dev/null
+++ b/test/transforms/test_node_property_split.py
@@ -0,0 +1,39 @@
+import pytest
+import torch
+
+from torch_geometric.datasets import graph_generator
+from torch_geometric.testing import withPackage
+from torch_geometric.transforms import NodePropertySplit
+
+
+@withPackage('networkx')
+@pytest.mark.parametrize('property_name', [
+    'popularity',
+    'locality',
+    'density',
+])
+def test_node_property_split(property_name):
+    ratios = [0.3, 0.1, 0.1, 0.2, 0.3]
+
+    transform = NodePropertySplit(property_name, ratios)
+    assert str(transform) == f'NodePropertySplit({property_name})'
+
+    data = graph_generator.ERGraph(num_nodes=100, edge_prob=0.4)()
+    data = transform(data)
+
+    node_ids = []
+    for name, ratio in zip([
+            'id_train_mask',
+            'id_val_mask',
+            'id_test_mask',
+            'ood_val_mask',
+            'ood_test_mask',
+    ], ratios):
+        assert data[name].dtype == torch.bool
+        assert data[name].size() == (100, )
+        assert int(data[name].sum()) == 100 * ratio
+        node_ids.extend(data[name].nonzero().view(-1).tolist())
+
+    # Check that masks are non-intersecting and cover all nodes:
+    node_ids = torch.tensor(node_ids)
+    assert node_ids.numel() == torch.unique(node_ids).numel() == 100