Repository: pyg-team/pytorch_geometric
Pull request #10009 (2025-02-07 09:59:05)

Add `LinkPredDiversity`


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3d02d0e91..c066a8ea2 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 `Diversity` metric for link prediction ([#10009](https://github.com/pyg-team/pytorch_geometric/pull/10009))
 - Added `Coverage` metric for link prediction ([#10006](https://github.com/pyg-team/pytorch_geometric/pull/10006))
 - Added Graph Transformer Tutorial ([#8144](https://github.com/pyg-team/pytorch_geometric/pull/8144))
 - Consolidate Cugraph examples into ogbn_train_cugraph.py and ogbn_train_cugraph_multigpu.py for ogbn-arxiv, ogbn-products and ogbn-papers100M ([#9953](https://github.com/pyg-team/pytorch_geometric/pull/9953))
diff --git a/torch_geometric/metrics/__init__.py b/torch_geometric/metrics/__init__.py
index 1b4462ac5..4859acd06 100644
--- a/torch_geometric/metrics/__init__.py
+++ b/torch_geometric/metrics/__init__.py
@@ -10,6 +10,7 @@ from .link_pred import (
     LinkPredNDCG,
     LinkPredMRR,
     LinkPredCoverage,
+    LinkPredDiversity,
 )
 
 link_pred_metrics = [
@@ -22,6 +23,7 @@ link_pred_metrics = [
     'LinkPredNDCG',
     'LinkPredMRR',
     'LinkPredCoverage',
+    'LinkPredDiversity',
 ]
 
 __all__ = link_pred_metrics
diff --git a/torch_geometric/metrics/link_pred.py b/torch_geometric/metrics/link_pred.py
index f4204494b..5be78188f 100644
--- a/torch_geometric/metrics/link_pred.py
+++ b/torch_geometric/metrics/link_pred.py
@@ -565,7 +565,8 @@ class LinkPredMRR(LinkPredMetric):
 
 
 class LinkPredCoverage(_LinkPredMetric):
-    r"""A link prediction metric to compute the Coverage @ :math:`k`.
+    r"""A link prediction metric to compute the Coverage @ :math:`k` of
+    predictions.
 
     Args:
         k (int): The number of top-:math:`k` predictions to evaluate against.
@@ -601,3 +602,68 @@ class LinkPredCoverage(_LinkPredMetric):
     def __repr__(self) -> str:
         return (f'{self.__class__.__name__}(k={self.k}, '
                 f'num_dst_nodes={self.num_dst_nodes})')
+
+
+class LinkPredDiversity(_LinkPredMetric):
+    r"""A link prediction metric to compute the Diversity @ :math:`k` of
+    predictions according to item categories.
+
+    Diversity is computed as
+
+    .. math::
+        div_{u@k} = 1 - \left( \frac{1}{k \cdot (k-1)} \right) \sum_{i \neq j}
+        sim(i, j)
+
+    where
+
+    .. math::
+        sim(i,j) = \begin{cases}
+            1 & \quad \text{if } i,j \text{ share category,}\\
+            0 & \quad \text{otherwise.}
+        \end{cases}
+
+    which measures the pair-wise inequality of recommendations according to
+    item categories.
+
+    Args:
+        k (int): The number of top-:math:`k` predictions to evaluate against.
+        category (torch.Tensor): A vector that assigns each destination node to
+            a specific category.
+    """
+    higher_is_better: bool = True
+
+    def __init__(self, k: int, category: Tensor) -> None:
+        super().__init__(k)
+
+        if WITH_TORCHMETRICS:
+            self.add_state('accum', torch.tensor(0.), dist_reduce_fx='sum')
+            self.add_state('total', torch.tensor(0), dist_reduce_fx='sum')
+        else:
+            self.register_buffer('accum', torch.tensor(0.))
+            self.register_buffer('total', torch.tensor(0))
+
+        self.category: Tensor
+        self.register_buffer('category', category)
+
+    def update(
+        self,
+        pred_index_mat: Tensor,
+        edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
+        edge_label_weight: Optional[Tensor] = None,
+    ) -> None:
+        category = self.category[pred_index_mat[:, :self.k]]
+
+        sim = (category.unsqueeze(-2) == category.unsqueeze(-1)).sum(dim=-1)
+        div = 1 - 1 / (self.k * (self.k - 1)) * (sim - 1).sum(dim=-1)
+
+        self.accum += div.sum()
+        self.total += pred_index_mat.size(0)
+
+    def compute(self) -> Tensor:
+        if self.total == 0:
+            return torch.zeros_like(self.accum)
+        return self.accum / self.total
+
+    def _reset(self) -> None:
+        self.accum.zero_()
+        self.total.zero_()


Test patch:
diff --git a/test/metrics/test_link_pred_metric.py b/test/metrics/test_link_pred_metric.py
index 8d76c6d25..b50bfb7e7 100644
--- a/test/metrics/test_link_pred_metric.py
+++ b/test/metrics/test_link_pred_metric.py
@@ -5,6 +5,7 @@ import torch
 
 from torch_geometric.metrics import (
     LinkPredCoverage,
+    LinkPredDiversity,
     LinkPredF1,
     LinkPredMAP,
     LinkPredMetricCollection,
@@ -181,7 +182,7 @@ def test_mrr():
 
 def test_coverage():
     pred_index_mat = torch.tensor([[1, 0], [1, 2], [0, 2], [0, 1]])
-    edge_label_index = torch.tensor([[0, 0, 2, 2, 3], [0, 1, 2, 1, 2]])
+    edge_label_index = torch.empty(2, 0, dtype=torch.long)
 
     metric = LinkPredCoverage(k=2, num_dst_nodes=3)
     assert str(metric) == 'LinkPredCoverage(k=2, num_dst_nodes=3)'
@@ -202,6 +203,20 @@ def test_coverage():
     assert float(result) == 2 / 4
 
 
+def test_diversity():
+    pred_index_mat = torch.tensor([[0, 1, 2], [3, 1, 0]])
+    category = torch.tensor([0, 1, 2, 0])
+    edge_label_index = torch.empty(2, 0, dtype=torch.long)
+
+    metric = LinkPredDiversity(k=3, category=category)
+    assert str(metric) == 'LinkPredDiversity(k=3)'
+    metric.update(pred_index_mat, edge_label_index)
+    result = metric.compute()
+    metric.reset()
+
+    assert pytest.approx(float(result)) == (1 + 2 / 3) / 2
+
+
 @pytest.mark.parametrize('num_src_nodes', [10])
 @pytest.mark.parametrize('num_dst_nodes', [50])
 @pytest.mark.parametrize('num_edges', [200])