Repository: pyg-team/pytorch_geometric
Pull request #8647 (2023-12-20 16:12:19)

Add `f1@k` and `map@k` metrics
From https://github.com/pyg-team/pytorch_geometric/issues/8452

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8647?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`11a29b0`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/11a29b097a7448b76375045e753154f79f4e5bd2?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.12% compared to head [(`8609857`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8647?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.12%.

> :exclamation: Current head 8609857 differs from pull request most recent head 44122d3. Consider uploading reports for the commit 44122d3 to get more accurate results


<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master    #8647   +/-   ##
=======================================
  Coverage   89.12%   89.12%           
=======================================
  Files         481      481           
  Lines       30830    30839    +9     
=======================================
+ Hits        27477    27486    +9     
  Misses       3353     3353           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8647?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 022560cc5..e0ac84f5b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,7 +17,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added warning when calling `dataset.num_classes` on regression problems ([#8550](https://github.com/pyg-team/pytorch_geometric/pull/8550))
 - Added relabel node functionality to `dropout_node` ([#8524](https://github.com/pyg-team/pytorch_geometric/pull/8524))
 - Added support for type checking via `mypy` ([#8254](https://github.com/pyg-team/pytorch_geometric/pull/8254))
-- Added support for link-prediction retrieval metrics ([#8499](https://github.com/pyg-team/pytorch_geometric/pull/8499), [#8326](https://github.com/pyg-team/pytorch_geometric/pull/8326), [#8566](https://github.com/pyg-team/pytorch_geometric/pull/8566))
+- Added support for link-prediction retrieval metrics ([#8499](https://github.com/pyg-team/pytorch_geometric/pull/8499), [#8326](https://github.com/pyg-team/pytorch_geometric/pull/8326), [#8566](https://github.com/pyg-team/pytorch_geometric/pull/8566), [#8647](https://github.com/pyg-team/pytorch_geometric/pull/8647))
 - Added METIS partitioning with CSC/CSR format selection in `ClusterData` ([#8438](https://github.com/pyg-team/pytorch_geometric/pull/8438))
 - Added `is_torch_instance` to check against the original class of compiled models ([#8461](https://github.com/pyg-team/pytorch_geometric/pull/8461))
 - Added dense computation for `AddRandomWalkPE` ([#8431](https://github.com/pyg-team/pytorch_geometric/pull/8431))
diff --git a/torch_geometric/nn/metrics.py b/torch_geometric/nn/metrics.py
index 092f02b7f..eb93d8922 100644
--- a/torch_geometric/nn/metrics.py
+++ b/torch_geometric/nn/metrics.py
@@ -154,11 +154,38 @@ class LinkPredRecall(LinkPredMetric):
     """
     higher_is_better: bool = True
 
-    def __init__(self, k: int):
-        super().__init__(k)
+    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
+        return pred_isin_mat.sum(dim=-1) / y_count.clamp(min=1e-7)
+
+
+class LinkPredF1(LinkPredMetric):
+    r"""A link prediction metric to compute F1@:math`k`.
+
+    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:
+        isin_count = pred_isin_mat.sum(dim=-1)
+        precision = isin_count / self.k
+        recall = isin_count = isin_count / y_count.clamp(min=1e-7)
+        return 2 * precision * recall / (precision + recall).clamp(min=1e-7)
+
+
+class LinkPredMAP(LinkPredMetric):
+    r"""A link prediction metric to compute MAP@:math`k`.
+
+    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:
-        return pred_isin_mat.sum(dim=1) / y_count.clamp(min=1e-7)
+        cum_precision = (torch.cumsum(pred_isin_mat, dim=1) /
+                         torch.arange(1, self.k + 1, device=y_count.device))
+        return ((cum_precision * pred_isin_mat).sum(dim=-1) /
+                y_count.clamp(min=1e-7, max=self.k))
 
 
 class LinkPredNDCG(LinkPredMetric):


Test patch:
diff --git a/test/nn/test_metrics.py b/test/nn/test_metrics.py
index 8f06905ad..6bc942165 100644
--- a/test/nn/test_metrics.py
+++ b/test/nn/test_metrics.py
@@ -4,6 +4,8 @@ import pytest
 import torch
 
 from torch_geometric.nn.metrics import (
+    LinkPredF1,
+    LinkPredMAP,
     LinkPredNDCG,
     LinkPredPrecision,
     LinkPredRecall,
@@ -57,19 +59,41 @@ def test_recall():
     edge_label_index = torch.tensor([[0, 0, 0, 2, 2], [0, 1, 2, 2, 1]])
 
     metric = LinkPredRecall(k=2)
-    assert str(metric) == f'LinkPredRecall(k={2})'
+    assert str(metric) == 'LinkPredRecall(k=2)'
     metric.update(pred_mat, edge_label_index)
     result = metric.compute()
 
     assert float(result) == pytest.approx(0.5 * (2 / 3 + 0.5))
 
 
+def test_f1():
+    pred_mat = torch.tensor([[1, 0], [1, 2], [0, 2]])
+    edge_label_index = torch.tensor([[0, 0, 0, 2, 2], [0, 1, 2, 2, 1]])
+
+    metric = LinkPredF1(k=2)
+    assert str(metric) == 'LinkPredF1(k=2)'
+    metric.update(pred_mat, edge_label_index)
+    result = metric.compute()
+    assert float(result) == pytest.approx(0.6500)
+
+
+def test_map():
+    pred_mat = torch.tensor([[1, 0], [1, 2], [0, 2]])
+    edge_label_index = torch.tensor([[0, 0, 0, 2, 2], [0, 1, 2, 2, 1]])
+
+    metric = LinkPredMAP(k=2)
+    assert str(metric) == 'LinkPredMAP(k=2)'
+    metric.update(pred_mat, edge_label_index)
+    result = metric.compute()
+    assert float(result) == pytest.approx(0.6250)
+
+
 def test_ndcg():
     pred_mat = torch.tensor([[1, 0], [1, 2], [0, 2]])
     edge_label_index = torch.tensor([[0, 0, 2, 2], [0, 1, 2, 1]])
 
     metric = LinkPredNDCG(k=2)
-    assert str(metric) == f'LinkPredNDCG(k={2})'
+    assert str(metric) == 'LinkPredNDCG(k=2)'
     metric.update(pred_mat, edge_label_index)
     result = metric.compute()