Mean reciprocal rank metric (#9632)

Addresses #9631. Implements `LinkPredMRR` as a `LinkPredMetric`, with
test included.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: rusty1s <[email protected]>

Author: 3 people

CHANGELOG.md

@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added the `LinkPredMRR` metric ([#9632](https://github.com/pyg-team/pytorch_geometric/pull/9632))
 - Added PyTorch 2.4 support ([#9594](https://github.com/pyg-team/pytorch_geometric/pull/9594))
 - Added `utils.normalize_edge_index` for symmetric/asymmetric normalization of graph edges ([#9554](https://github.com/pyg-team/pytorch_geometric/pull/9554))
 - Added the `RemoveSelfLoops` transformation ([#9562](https://github.com/pyg-team/pytorch_geometric/pull/9562))

test/metrics/test_link_pred_metric.py

@@ -6,6 +6,7 @@
 from torch_geometric.metrics import (
     LinkPredF1,
     LinkPredMAP,
+    LinkPredMRR,
     LinkPredNDCG,
     LinkPredPrecision,
     LinkPredRecall,
@@ -98,3 +99,15 @@ def test_ndcg():
     result = metric.compute()
 
     assert float(result) == pytest.approx(0.6934264)
+
+
+def test_mrr():
+    pred_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]])
+
+    metric = LinkPredMRR(k=2)
+    assert str(metric) == 'LinkPredMRR(k=2)'
+    metric.update(pred_mat, edge_label_index)
+    result = metric.compute()
+
+    assert float(result) == pytest.approx((1 + 0.5 + 0) / 3)

torch_geometric/metrics/__init__.py

@@ -1,14 +1,21 @@
 # flake8: noqa
 
-from .link_pred import (LinkPredPrecision, LinkPredRecall, LinkPredF1,
-                        LinkPredMAP, LinkPredNDCG)
+from .link_pred import (
+    LinkPredPrecision,
+    LinkPredRecall,
+    LinkPredF1,
+    LinkPredMAP,
+    LinkPredNDCG,
+    LinkPredMRR,
+)
 
 link_pred_metrics = [
     'LinkPredPrecision',
     'LinkPredRecall',
     'LinkPredF1',
     'LinkPredMAP',
     'LinkPredNDCG',
+    'LinkPredMRR',
 ]
 
 __all__ = link_pred_metrics

torch_geometric/metrics/link_pred.py

@@ -216,3 +216,20 @@ def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
         out = dcg / idcg
         out[out.isnan() | out.isinf()] = 0.0
         return out
+
+
+class LinkPredMRR(LinkPredMetric):
+    r"""A link prediction metric to compute the MRR @ :math:`k` (Mean
+    Reciprocal Rank).
+
+    Args:
+        k (int): The number of top-:math:`k` predictions to evaluate against.
+    """
+    higher_is_better: bool = True
+
+    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
+        rank = pred_isin_mat.type(torch.uint8).argmax(dim=-1)
+        is_correct = pred_isin_mat.gather(1, rank.view(-1, 1)).view(-1)
+        reciprocals = 1.0 / (rank + 1)
+        reciprocals[~is_correct] = 0.0
+        return reciprocals