From ed5081f417529735c6c6001e6f0a4d0ac8daaeeb Mon Sep 17 00:00:00 2001 From: Jintang Li Date: Wed, 5 Oct 2022 14:26:06 +0800 Subject: [PATCH] Add `assortativity` to `torch_geometric.utils` (#5587) * add assortativity * test * doc-string * doc-string * Update torch_geometric/utils/assortativity.py Co-authored-by: Matthias Fey * Update torch_geometric/utils/assortativity.py Co-authored-by: Matthias Fey * Update torch_geometric/utils/assortativity.py Co-authored-by: Padarn Wilson * update test * doc-string * fix test * changelog Co-authored-by: Matthias Fey Co-authored-by: Padarn Wilson --- CHANGELOG.md | 1 + test/utils/test_assortativity.py | 18 +++++++ torch_geometric/utils/__init__.py | 4 +- torch_geometric/utils/assortativity.py | 66 ++++++++++++++++++++++++++ 4 files changed, 88 insertions(+), 1 deletion(-) create mode 100644 test/utils/test_assortativity.py create mode 100644 torch_geometric/utils/assortativity.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 1cb3b145da50..2f67cf300cd3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). ## [2.2.0] - 2022-MM-DD ### Added +- Added `assortativity` that computes degree assortativity coefficient ([#5587](https://github.com/pyg-team/pytorch_geometric/pull/5587)) - Added `SSGConv` layer ([#5599](https://github.com/pyg-team/pytorch_geometric/pull/5599)) - Added `shuffle_node`, `mask_feature` and `add_random_edge` augmentation methdos ([#5548](https://github.com/pyg-team/pytorch_geometric/pull/5548)) - Added `dropout_path` augmentation that drops edges from a graph based on random walks ([#5531](https://github.com/pyg-team/pytorch_geometric/pull/5531)) diff --git a/test/utils/test_assortativity.py b/test/utils/test_assortativity.py new file mode 100644 index 000000000000..09ed3b958c11 --- /dev/null +++ b/test/utils/test_assortativity.py @@ -0,0 +1,18 @@ +import pytest +import torch + +from torch_geometric.utils import assortativity + + +def test_assortativity(): + # completely assortative graph + edge_index = torch.tensor([[0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 5], + [1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2, 5, 4]]) + out = assortativity(edge_index) + assert pytest.approx(out, abs=1e-5) == 1.0 + + # completely disassortative graph + edge_index = torch.tensor([[0, 1, 2, 3, 4, 5, 5, 5, 5, 5], + [5, 5, 5, 5, 5, 0, 1, 2, 3, 4]]) + out = assortativity(edge_index) + assert pytest.approx(out, abs=1e-5) == -1.0 diff --git a/torch_geometric/utils/__init__.py b/torch_geometric/utils/__init__.py index 770464784a46..cb2cdef01bc0 100644 --- a/torch_geometric/utils/__init__.py +++ b/torch_geometric/utils/__init__.py @@ -1,6 +1,7 @@ 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 @@ -11,6 +12,7 @@ from .subgraph import (get_num_hops, subgraph, k_hop_subgraph, bipartite_subgraph) from .homophily import homophily +from .assortativity import assortativity from .get_laplacian import get_laplacian from .get_mesh_laplacian import get_mesh_laplacian from .mask import index_to_mask, mask_to_index @@ -34,7 +36,6 @@ structured_negative_sampling_feasible) from .train_test_split_edges import train_test_split_edges from .scatter import scatter -from .augmentation import shuffle_node, mask_feature, add_random_edge __all__ = [ 'degree', @@ -63,6 +64,7 @@ 'bipartite_subgraph', 'k_hop_subgraph', 'homophily', + 'assortativity', 'get_laplacian', 'get_mesh_laplacian', 'index_to_mask', diff --git a/torch_geometric/utils/assortativity.py b/torch_geometric/utils/assortativity.py new file mode 100644 index 000000000000..8fb0ea622dd9 --- /dev/null +++ b/torch_geometric/utils/assortativity.py @@ -0,0 +1,66 @@ +import torch +from torch_sparse import SparseTensor + +from torch_geometric.typing import Adj +from torch_geometric.utils import coalesce, degree + +from .to_dense_adj import to_dense_adj + + +def assortativity(edge_index: Adj) -> float: + r"""The degree assortativity coefficient from the + `"Mixing patterns in networks" + `_ paper. + Assortativity in a network refers to the tendency of nodes to + connect with other similar nodes over dissimilar nodes. + It is computed from Pearson correlation coefficient of the node degrees. + + Args: + edge_index (Tensor or SparseTensor): The graph connectivity. + + Returns: + The value of the degree assortativity coefficient for the input + graph :math:`\in [-1, 1]` + + Example: + + >>> edge_index = torch.tensor([[0, 1, 2, 3, 2], + ... [1, 2, 0, 1, 3]]) + >>> assortativity(edge_index) + -0.666667640209198 + """ + if isinstance(edge_index, SparseTensor): + row, col, _ = edge_index.coo() + else: + row, col = edge_index + + device = row.device + out_deg = degree(row, dtype=torch.long) + in_deg = degree(col, dtype=torch.long) + degrees = torch.unique(torch.cat([out_deg, in_deg])) + mapping = row.new_zeros(degrees.max().item() + 1) + mapping[degrees] = torch.arange(degrees.size(0), device=device) + + # Compute degree mixing matrix (joint probability distribution) `M` + num_degrees = degrees.size(0) + src_deg = mapping[out_deg[row]] + dst_deg = mapping[in_deg[col]] + + pairs = torch.stack([src_deg, dst_deg], dim=0) + occurrence = torch.ones(pairs.size(1), device=device) + pairs, occurrence = coalesce(pairs, occurrence) + M = to_dense_adj(pairs, edge_attr=occurrence, max_num_nodes=num_degrees)[0] + # normalization + M /= M.sum() + + # numeric assortativity coefficient, computed by + # Pearson correlation coefficient of the node degrees + x = y = degrees.float() + a, b = M.sum(0), M.sum(1) + + vara = (a * x**2).sum() - ((a * x).sum())**2 + varb = (b * x**2).sum() - ((b * x).sum())**2 + xy = torch.outer(x, y) + ab = torch.outer(a, b) + out = (xy * (M - ab)).sum() / (vara * varb).sqrt() + return out.item()