Repository: pyg-team/pytorch_geometric

Issue:
[Roadmap] PyG for Recommendation 🚀
### 🚀 The feature, motivation and pitch

**This roadmap aims to bring better support for recommendation tasks to PyG.**

Currently, all/most of our link prediction models are trained and evaluated using binary classification metrics. However, this usually requires that we have a set of candidates in advance, from which we can then infer the existence of links. This is not necessarily practical, since in most cases, we want to find the top-k most likely links from the full set of `O(N^2)` pairs.

While training can still be done via negative sampling and binary classification, this roadmap resolves around bringing better support for link prediction evaluation into PyG, with the following end-to-end pipeline:
1. Embed all source and destination nodes
1. Use "Maximum Inner Product Search" (MIPS) to find the top-k most likely links (via [`MIPSKNNIndex`](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.pool.MIPSKNNIndex.html#torch_geometric.nn.pool.MIPSKNNIndex))
1. Evaluate using common metrics for recommendation, e.g., `map@k`, `precision@k`, `recall@k`, `f1@k`, `ndcg@k`.

### Metrics

We need to support recommendation metrics, which can be updated and computed in a mini-batch fashion. A related issue can be found [here](https://github.com/pyg-team/pytorch_geometric/issues/8271). Its interface can/should follow the `torchmetrics.Metric` interface, *e.g.*:
```python
class LinkPredMetric(torchmetrics.Metric):
    def __init__(self, k: int):
        pass

    def update(self, top_k_pred_mat: Tensor, edge_label_index: Tensor):
        pass

    def compute(self):
        pass
```
where `top_k_pred_mat` holds the top-k indices for each left-hand-side (LHS) entity, and `edge_label_index` holds the ground-truth information as a `[2, num_targets]` matrix.

* [x] Implement `LinkPredMetric` interface
* [x] Implement `map@k`
* [x] Implement `precision@k`
* [x] Implement `recall@k`
* [x] Implement `f1@k`
* [x] Implement `ndcg@k` (#8326)

### Examples

With this, we can build one or more clear and descriptive examples of how to leverage PyG for recommendation.

* [x] Select and implement one or two datasets commonly used for recommendation
* [x] Add exclusion logic to `MIPSKNNIndex`
* [x] Build an example that implements this pipeline
* [ ] Write a tutorial about recommendation in PyG
* [ ] Advanced: Combine PyG's recommendation capabilities with its temporal GNN support (see #3230)


Discussion:
@rusty1s I have done some similar work in the past extending PyG for better recommendation support (at a company where I couldn't contribute it to open source). Would love to contribute to this!

Resolving patch:
diff --git a/torch_geometric/nn/metrics.py b/torch_geometric/nn/metrics.py
--- a/torch_geometric/nn/metrics.py
+++ b/torch_geometric/nn/metrics.py
@@ -19,12 +19,11 @@ class LinkPredMetric(BaseMetric, ABC):
     r"""An abstract class for computing link prediction retrieval metrics.
 
     Args:
-        k (int): The number of top-:math:`k` predictions to evaluate
-            against.
+        k (int): The number of top-:math:`k` predictions to evaluate against.
     """
-    is_differentiable: Optional[bool] = None
+    is_differentiable: bool = False
+    full_state_update: bool = False
     higher_is_better: Optional[bool] = None
-    full_state_update: Optional[bool] = None
 
     def __init__(self, k: int):
         super().__init__()
@@ -132,35 +131,44 @@ def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
         raise NotImplementedError
 
     def __repr__(self) -> str:
-        return f'{self.__class__.__name__}({self.k})'
+        return f'{self.__class__.__name__}(k={self.k})'
 
 
 class LinkPredPrecision(LinkPredMetric):
     r"""A link prediction metric to compute Precision@:math`k`.
 
     Args:
-        k (int): The number of top-:math:`k` predictions to evaluate
-            against.
+        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 LinkPredRecall(LinkPredMetric):
+    r"""A link prediction metric to compute Recall@:math:`k`.
+
+    Args:
+        k (int): The number of top-:math:`k` predictions to evaluate against.
+    """
+    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 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.
+        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)


Test patch:
diff --git a/test/nn/test_metrics.py b/test/nn/test_metrics.py
--- a/test/nn/test_metrics.py
+++ b/test/nn/test_metrics.py
@@ -3,7 +3,11 @@
 import pytest
 import torch
 
-from torch_geometric.nn.metrics import LinkPredNDCG, LinkPredPrecision
+from torch_geometric.nn.metrics import (
+    LinkPredNDCG,
+    LinkPredPrecision,
+    LinkPredRecall,
+)
 
 
 @pytest.mark.parametrize('num_src_nodes', [100])
@@ -21,7 +25,7 @@ def test_precision(num_src_nodes, num_dst_nodes, num_edges, batch_size, k):
     top_k_pred_mat = pred.topk(k, dim=1)[1]
 
     metric = LinkPredPrecision(k)
-    assert str(metric) == f'LinkPredPrecision({k})'
+    assert str(metric) == f'LinkPredPrecision(k={k})'
 
     for node_id in torch.split(torch.randperm(num_src_nodes), batch_size):
         mask = torch.isin(edge_label_index[0], node_id)
@@ -48,11 +52,24 @@ def test_precision(num_src_nodes, num_dst_nodes, num_edges, batch_size, k):
     assert torch.allclose(out, expected)
 
 
+def test_recall():
+    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 = LinkPredRecall(k=2)
+    assert str(metric) == f'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_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})'
     metric.update(pred_mat, edge_label_index)
     result = metric.compute()