Repository: pyg-team/pytorch_geometric
Pull request #8326 (2023-11-05 17:34:52)

`NDCG@k` metric for link-prediction
Hello!

This branch implements an example of link prediction using the [ndcg@k](https://en.wikipedia.org/wiki/Discounted_cumulative_gain) metric.

This has been discussed in the #8271 issue. 

I have tried to first implement it in a new example as a starting point. It probably should be moved somewhere else, or have a more distinctive example. This example is a simple copy from the `link_prediction.py` example, replacing AUC with NDCG. 

There remain some questions about the signature of the `update` method, as well as how to call the method in the example.

For example, we might want to demonstrate that we can call it during training, adding mini-batches one at a time. 

We might also want to avoid predicting links from another set, that way, we are not considering valid links as undesired.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8326?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#8326](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8326?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (83eedc3) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/9f7e82442efb41687e60469f1e63b0efaa793165?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (9f7e824) will **decrease** coverage by `0.62%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master    #8326      +/-   ##
==========================================
- Coverage   89.00%   88.39%   -0.62%     
==========================================
  Files         475      477       +2     
  Lines       28841    28866      +25     
==========================================
- Hits        25670    25516     -154     
- Misses       3171     3350     +179     
```


| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8326?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/\_\_init\_\_.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8326?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/metrics/\_\_init\_\_.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8326?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL21ldHJpY3MvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/metrics/ndcg.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8326?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL21ldHJpY3MvbmRjZy5weQ==) | `100.00% <100.00%> (ø)` | |

... and [36 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8326/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: Codecov offers a browser extension for seamless coverage viewing on GitHub. Try it in [Chrome](https://chrome.google.com/webstore/detail/codecov/gedikamndpbemklijjkncpnolildpbgo) or [Firefox](https://addons.mozilla.org/en-US/firefox/addon/codecov/) today!


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index ff6f11081..1d8ca38ad 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 the `NDCG@k` metric for link-prediction ([#8326](https://github.com/pyg-team/pytorch_geometric/pull/8326))
 - 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))
diff --git a/torch_geometric/nn/metrics.py b/torch_geometric/nn/metrics.py
index a52bc7985..25badcde1 100644
--- a/torch_geometric/nn/metrics.py
+++ b/torch_geometric/nn/metrics.py
@@ -4,7 +4,7 @@ from typing import Optional, Tuple, Union
 import torch
 from torch import Tensor
 
-from torch_geometric.utils import scatter
+from torch_geometric.utils import cumsum, scatter
 
 try:
     import torchmetrics  # noqa
@@ -19,21 +19,21 @@ class LinkPredMetric(BaseMetric, ABC):
     r"""An abstract class for computing link prediction retrieval metrics.
 
     Args:
-        top_k (int): The number of top-:math:`k` predictions to evaluate
+        k (int): The number of top-:math:`k` predictions to evaluate
             against.
     """
     is_differentiable: Optional[bool] = None
     higher_is_better: Optional[bool] = None
     full_state_update: Optional[bool] = None
 
-    def __init__(self, top_k: int):
+    def __init__(self, k: int):
         super().__init__()
 
-        if top_k <= 0:
-            raise ValueError(f"'top_k' needs to be a positive integer in "
-                             f"'{self.__class__.__name__}' (got {top_k})")
+        if k <= 0:
+            raise ValueError(f"'k' needs to be a positive integer in "
+                             f"'{self.__class__.__name__}' (got {k})")
 
-        self.top_k = top_k
+        self.k = k
 
         if WITH_TORCHMETRICS:
             self.add_state('accum', torch.tensor(0.), dist_reduce_fx='sum')
@@ -50,16 +50,20 @@ class LinkPredMetric(BaseMetric, ABC):
         r"""Updates the state variables based on the current mini-batch
         prediction.
 
+        :meth:`update` can be repeated multiple times to accumulate the results
+        of successive predictions, *e.g.*, inside a mini-batch training or
+        evaluation loop.
+
         Args:
             pred_index_mat (torch.Tensor): The top-:math:`k` predictions of
                 every example in the mini-batch with shape
-                :obj:`[batch_size, top_k]`.
+                :obj:`[batch_size, k]`.
             edge_label_index (torch.Tensor): The ground-truth indices for every
                 example in the mini-batch, given in COO format of shape
                 :obj:`[2, num_ground_truth_indices]`.
         """
-        if pred_index_mat.size(1) != self.top_k:
-            raise ValueError(f"Expected 'pred_index_mat' to hold {self.top_k} "
+        if pred_index_mat.size(1) != self.k:
+            raise ValueError(f"Expected 'pred_index_mat' to hold {self.k} "
                              f"many indices for every entry "
                              f"(got {pred_index_mat.size(1)})")
 
@@ -128,15 +132,49 @@ class LinkPredMetric(BaseMetric, ABC):
         raise NotImplementedError
 
     def __repr__(self) -> str:
-        return f'{self.__class__.__name__}({self.top_k})'
+        return f'{self.__class__.__name__}({self.k})'
 
 
 class LinkPredPrecision(LinkPredMetric):
     r"""A link prediction metric to compute Precision@:math`k`.
 
     Args:
-        top_k (int): The number of top-:math:`k` predictions to evaluate
+        k (int): The number of top-:math:`k` predictions to evaluate
+            against.
+    """
+    is_differentiable: bool = False
+    higher_is_better: bool = True
+    full_state_update: bool = False
+
+    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
+        return pred_isin_mat.sum(dim=-1) / self.k
+
+
+class LinkPredNDCG(LinkPredMetric):
+    r"""A link prediction metric to compute the Normalized Discounted
+    Cumulative Gain (NDCG).
+
+    Args:
+        k (int): The number of top-:math:`k` predictions to evaluate
             against.
     """
+    is_differentiable: bool = False
+    higher_is_better: bool = True
+    full_state_update: bool = False
+
+    def __init__(self, k: int):
+        super().__init__(k=k)
+
+        dtype = torch.get_default_dtype()
+        multiplier = 1.0 / torch.arange(2, k + 2, dtype=dtype).log2()
+
+        self.register_buffer('multiplier', multiplier)
+        self.register_buffer('idcg', cumsum(multiplier))
+
     def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
-        return pred_isin_mat.sum(dim=-1) / self.top_k
+        dcg = (pred_isin_mat * self.multiplier.view(1, -1)).sum(dim=-1)
+        idcg = self.idcg[y_count.clamp(max=self.k)]
+
+        out = dcg / idcg
+        out[out.isnan() | out.isinf()] = 0.0
+        return out


Test patch:
diff --git a/test/nn/test_metrics.py b/test/nn/test_metrics.py
index a40d5f43e..04f268bd8 100644
--- a/test/nn/test_metrics.py
+++ b/test/nn/test_metrics.py
@@ -3,25 +3,25 @@ from typing import List
 import pytest
 import torch
 
-from torch_geometric.nn.metrics import LinkPredPrecision
+from torch_geometric.nn.metrics import LinkPredNDCG, LinkPredPrecision
 
 
 @pytest.mark.parametrize('num_src_nodes', [100])
 @pytest.mark.parametrize('num_dst_nodes', [1000])
 @pytest.mark.parametrize('num_edges', [3000])
 @pytest.mark.parametrize('batch_size', [32])
-@pytest.mark.parametrize('top_k', [1, 10, 100])
-def test_metrics(num_src_nodes, num_dst_nodes, num_edges, batch_size, top_k):
+@pytest.mark.parametrize('k', [1, 10, 100])
+def test_precision(num_src_nodes, num_dst_nodes, num_edges, batch_size, k):
     row = torch.randint(0, num_src_nodes, (num_edges, ))
     col = torch.randint(0, num_dst_nodes, (num_edges, ))
     edge_label_index = torch.stack([row, col], dim=0)
 
     pred = torch.rand(num_src_nodes, num_dst_nodes)
     pred[row, col] += 0.3  # Offset positive links by a little.
-    top_k_pred_mat = pred.topk(top_k, dim=1)[1]
+    top_k_pred_mat = pred.topk(k, dim=1)[1]
 
-    metric = LinkPredPrecision(top_k)
-    assert str(metric) == f'LinkPredPrecision({top_k})'
+    metric = LinkPredPrecision(k)
+    assert str(metric) == f'LinkPredPrecision({k})'
 
     for node_id in torch.split(torch.randperm(num_src_nodes), batch_size):
         mask = torch.isin(edge_label_index[0], node_id)
@@ -42,7 +42,18 @@ def test_metrics(num_src_nodes, num_dst_nodes, num_edges, batch_size, top_k):
         y_index = col[row == i]
         if y_index.numel() > 0:
             mask = torch.isin(top_k_pred_mat[i], y_index)
-            precision = float(mask.sum() / top_k)
+            precision = float(mask.sum() / k)
             values.append(precision)
     expected = torch.tensor(values).mean()
     assert torch.allclose(out, expected)
+
+
+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)
+    metric.update(pred_mat, edge_label_index)
+    result = metric.compute()
+
+    assert float(result) == pytest.approx(0.6934264)