Repository: pyg-team/pytorch_geometric
Pull request #10022 (2025-02-12 11:44:22)

Add `LinkPredAveragePopularity`


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index eb97043e5..4e7038eae 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 `AveragePopularity` metric for link prediction ([#10022](https://github.com/pyg-team/pytorch_geometric/pull/10022))
 - Added `Personalization` metric for link prediction ([#10015](https://github.com/pyg-team/pytorch_geometric/pull/10015))
 - Added `HitRatio` metric for link prediction ([#10013](https://github.com/pyg-team/pytorch_geometric/pull/10013))
 - Added Data Splitting Tutorial ([#8366](https://github.com/pyg-team/pytorch_geometric/pull/8366))
diff --git a/torch_geometric/metrics/__init__.py b/torch_geometric/metrics/__init__.py
index cc096f9e0..a0130d2bc 100644
--- a/torch_geometric/metrics/__init__.py
+++ b/torch_geometric/metrics/__init__.py
@@ -13,6 +13,7 @@ from .link_pred import (
     LinkPredCoverage,
     LinkPredDiversity,
     LinkPredPersonalization,
+    LinkPredAveragePopularity,
 )
 
 link_pred_metrics = [
@@ -28,6 +29,7 @@ link_pred_metrics = [
     'LinkPredCoverage',
     'LinkPredDiversity',
     'LinkPredPersonalization',
+    'LinkPredAveragePopularity',
 ]
 
 __all__ = link_pred_metrics
diff --git a/torch_geometric/metrics/link_pred.py b/torch_geometric/metrics/link_pred.py
index 4995be2f8..6ec323fe0 100644
--- a/torch_geometric/metrics/link_pred.py
+++ b/torch_geometric/metrics/link_pred.py
@@ -722,12 +722,14 @@ class LinkPredPersonalization(_LinkPredMetric):
             of user recommendations should be processed at once.
             (default: :obj:`2**16`)
     """
+    higher_is_better: bool = True
+
     def __init__(self, k: int, batch_size: int = 2**16) -> None:
         super().__init__(k)
         self.batch_size = batch_size
 
         if WITH_TORCHMETRICS:
-            self.add_state("preds", default=[], dist_reduce_fx="cat")
+            self.add_state('preds', default=[], dist_reduce_fx='cat')
             self.add_state('dev_tensor', torch.empty(0), dist_reduce_fx='sum')
         else:
             self.preds: List[Tensor] = []
@@ -787,3 +789,51 @@ class LinkPredPersonalization(_LinkPredMetric):
 
     def _reset(self) -> None:
         self.preds = []
+
+
+class LinkPredAveragePopularity(_LinkPredMetric):
+    r"""A link prediction metric to compute the Average Recommendation
+    Popularity (ARP) @ :math:`k`, which provides insights into the model's
+    tendency to recommend popular items by averaging the popularity scores of
+    items within the top-:math:`k` recommendations.
+
+    Args:
+        k (int): The number of top-:math:`k` predictions to evaluate against.
+        popularity (torch.Tensor): The popularity of every item in the training
+            set, *e.g.*, the number of times an item has been rated.
+    """
+    higher_is_better: bool = False
+
+    def __init__(self, k: int, popularity: 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.popularity: Tensor
+        self.register_buffer('popularity', popularity)
+
+    def update(
+        self,
+        pred_index_mat: Tensor,
+        edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
+        edge_label_weight: Optional[Tensor] = None,
+    ) -> None:
+        pred_index_mat = pred_index_mat[:, :self.k]
+        popularity = self.popularity[pred_index_mat]
+        popularity = popularity.to(self.accum.dtype).mean(dim=-1)
+        self.accum += popularity.sum()
+        self.total += popularity.numel()
+
+    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 d7df7fb12..7f52c3172 100644
--- a/test/metrics/test_link_pred_metric.py
+++ b/test/metrics/test_link_pred_metric.py
@@ -4,6 +4,7 @@ import pytest
 import torch
 
 from torch_geometric.metrics import (
+    LinkPredAveragePopularity,
     LinkPredCoverage,
     LinkPredDiversity,
     LinkPredF1,
@@ -259,6 +260,20 @@ def test_personalization(device):
     metric.reset()
 
 
+def test_average_popularity():
+    pred_index_mat = torch.tensor([[0, 1, 2], [3, 1, 0]])
+    popularity = torch.tensor([10, 5, 2, 1])
+    edge_label_index = torch.empty(2, 0, dtype=torch.long)
+
+    metric = LinkPredAveragePopularity(k=3, popularity=popularity)
+    assert str(metric) == 'LinkPredAveragePopularity(k=3)'
+    metric.update(pred_index_mat, edge_label_index)
+    result = metric.compute()
+    metric.reset()
+
+    assert pytest.approx(float(result)) == (10 + 5 + 2 + 1 + 5 + 10) / 6
+
+
 @pytest.mark.parametrize('num_src_nodes', [10])
 @pytest.mark.parametrize('num_dst_nodes', [50])
 @pytest.mark.parametrize('num_edges', [200])