Repository: Lightning-AI/torchmetrics

Issue:
Introduce k parameter in RetrievalMAP to compute at cut-off k
## 🚀 Feature

Add a parameter named `k` to RetrievalMAP to compute the mean average precision at cut-off k (e.g. MAP@10), i.e. consider only the top k elements for each query.

### Motivation

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->
Other retrieval metrics such as (Precision and Recall) support limiting results to the top `k` returned results.
This is especially useful when calculating multiple different levels of cutoff for the same metric, e.g. P@1, P@5, P@10... 

### Pitch

Extend the modular metric and the functional interface of RETRIEVAL MEAN AVERAGE PRECISION (MAP) to compute the metric only from the top k elements of each query. By default all elements should be considered. 

### Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->
Not change the API and leave it to the user to limit the results.

### Additional context

<!-- Add any other context or screenshots about the feature request here. -->



Discussion:
Hi! thanks for your contribution!, great first issue!
@lucadiliello opinion here? is it a simple change to the codebase?
Should be a simple change. @Bomme can you open a PR?
Sure! I can prepare a PR. 

Resolving patch:
diff --git a/src/torchmetrics/functional/retrieval/average_precision.py b/src/torchmetrics/functional/retrieval/average_precision.py
--- a/src/torchmetrics/functional/retrieval/average_precision.py
+++ b/src/torchmetrics/functional/retrieval/average_precision.py
@@ -11,13 +11,15 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from typing import Optional
+
 import torch
 from torch import Tensor, tensor
 
 from torchmetrics.utilities.checks import _check_retrieval_functional_inputs
 
 
-def retrieval_average_precision(preds: Tensor, target: Tensor) -> Tensor:
+def retrieval_average_precision(preds: Tensor, target: Tensor, top_k: Optional[int] = None) -> Tensor:
     """Compute average precision (for information retrieval), as explained in `IR Average precision`_.
 
     ``preds`` and ``target`` should be of the same shape and live on the same device. If no ``target`` is ``True``,
@@ -27,10 +29,15 @@ def retrieval_average_precision(preds: Tensor, target: Tensor) -> Tensor:
     Args:
         preds: estimated probabilities of each document to be relevant.
         target: ground truth about each document being relevant or not.
+        top_k: consider only the top k elements (default: ``None``, which considers them all)
 
     Return:
         a single-value tensor with the average precision (AP) of the predictions ``preds`` w.r.t. the labels ``target``.
 
+    Raises:
+        ValueError:
+            If ``top_k`` is not ``None`` or an integer larger than 0.
+
     Example:
         >>> from torchmetrics.functional import retrieval_average_precision
         >>> preds = tensor([0.2, 0.3, 0.5])
@@ -40,10 +47,14 @@ def retrieval_average_precision(preds: Tensor, target: Tensor) -> Tensor:
     """
     preds, target = _check_retrieval_functional_inputs(preds, target)
 
+    top_k = top_k or preds.shape[-1]
+    if not isinstance(top_k, int) and top_k <= 0:
+        raise ValueError(f"Argument ``top_k`` has to be a positive integer or None, but got {top_k}.")
+
+    target = target[preds.topk(min(top_k, preds.shape[-1]), sorted=True, dim=-1)[1]]
     if not target.sum():
         return tensor(0.0, device=preds.device)
 
-    target = target[torch.argsort(preds, dim=-1, descending=True)]
     positions = torch.arange(1, len(target) + 1, device=target.device, dtype=torch.float32)[target > 0]
     res = torch.div((torch.arange(len(positions), device=positions.device, dtype=torch.float32) + 1), positions).mean()
     return res
diff --git a/src/torchmetrics/retrieval/average_precision.py b/src/torchmetrics/retrieval/average_precision.py
--- a/src/torchmetrics/retrieval/average_precision.py
+++ b/src/torchmetrics/retrieval/average_precision.py
@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from typing import Any, Optional
+
 from torch import Tensor
 
 from torchmetrics.functional.retrieval.average_precision import retrieval_average_precision
@@ -49,6 +51,7 @@ class RetrievalMAP(RetrievalMetric):
 
         ignore_index:
             Ignore predictions where the target is equal to this number.
+        top_k: consider only the top k elements for each query (default: ``None``, which considers them all)
         kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.
 
     Raises:
@@ -56,6 +59,8 @@ class RetrievalMAP(RetrievalMetric):
             If ``empty_target_action`` is not one of ``error``, ``skip``, ``neg`` or ``pos``.
         ValueError:
             If ``ignore_index`` is not `None` or an integer.
+        ValueError:
+            If ``top_k`` is not ``None`` or an integer larger than 0.
 
     Example:
         >>> from torch import tensor
@@ -72,5 +77,22 @@ class RetrievalMAP(RetrievalMetric):
     higher_is_better: bool = True
     full_state_update: bool = False
 
+    def __init__(
+        self,
+        empty_target_action: str = "neg",
+        ignore_index: Optional[int] = None,
+        top_k: Optional[int] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(
+            empty_target_action=empty_target_action,
+            ignore_index=ignore_index,
+            **kwargs,
+        )
+
+        if top_k is not None and not isinstance(top_k, int) and top_k <= 0:
+            raise ValueError(f"Argument ``top_k`` has to be a positive integer or None, but got {top_k}")
+        self.k = top_k
+
     def _metric(self, preds: Tensor, target: Tensor) -> Tensor:
-        return retrieval_average_precision(preds, target)
+        return retrieval_average_precision(preds, target, top_k=self.k)


Test patch:
diff --git a/tests/unittests/retrieval/helpers.py b/tests/unittests/retrieval/helpers.py
--- a/tests/unittests/retrieval/helpers.py
+++ b/tests/unittests/retrieval/helpers.py
@@ -118,9 +118,10 @@ def _compute_sklearn_metric(
             res = metric(trg, pds, **kwargs)
             sk_results.append(res)
 
-    if len(sk_results) > 0:
-        return np.mean(sk_results)
-    return np.array(0.0)
+    sk_results = np.array(sk_results)
+    sk_results[np.isnan(sk_results)] = 0.0  # this is needed with old versions of sklearn
+
+    return sk_results.mean() if len(sk_results) > 0 else np.array(0.0)
 
 
 def _concat_tests(*tests: Tuple[Dict]) -> Dict:
diff --git a/tests/unittests/retrieval/test_map.py b/tests/unittests/retrieval/test_map.py
--- a/tests/unittests/retrieval/test_map.py
+++ b/tests/unittests/retrieval/test_map.py
@@ -11,6 +11,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from typing import Optional
+
+import numpy as np
 import pytest
 from sklearn.metrics import average_precision_score as sk_average_precision_score
 from torch import Tensor
@@ -32,10 +35,21 @@
 seed_all(42)
 
 
+def _average_precision_at_k(target: np.ndarray, preds: np.ndarray, top_k: Optional[int] = None):
+    """Wrapper around reference metric to account for top_k argument."""
+    assert target.shape == preds.shape
+    assert len(target.shape) == 1
+    top_k = top_k or len(preds)
+    idx = np.argsort(preds, axis=0)[::-1][:top_k]
+    target, preds = target[idx], preds[idx]
+    return sk_average_precision_score(target, preds)
+
+
 class TestMAP(RetrievalMetricTester):
     @pytest.mark.parametrize("ddp", [True, False])
     @pytest.mark.parametrize("empty_target_action", ["skip", "neg", "pos"])
     @pytest.mark.parametrize("ignore_index", [None, 1])  # avoid setting 0, otherwise test with all 0 targets will fail
+    @pytest.mark.parametrize("top_k", [None, 1, 4, 10])
     @pytest.mark.parametrize(**_default_metric_class_input_arguments)
     def test_class_metric(
         self,
@@ -45,8 +59,9 @@ def test_class_metric(
         target: Tensor,
         empty_target_action: str,
         ignore_index: int,
+        top_k: int,
     ):
-        metric_args = {"empty_target_action": empty_target_action, "ignore_index": ignore_index}
+        metric_args = {"empty_target_action": empty_target_action, "ignore_index": ignore_index, "top_k": top_k}
 
         self.run_class_metric_test(
             ddp=ddp,
@@ -54,12 +69,13 @@ def test_class_metric(
             preds=preds,
             target=target,
             metric_class=RetrievalMAP,
-            reference_metric=sk_average_precision_score,
+            reference_metric=_average_precision_at_k,
             metric_args=metric_args,
         )
 
     @pytest.mark.parametrize("ddp", [True, False])
     @pytest.mark.parametrize("empty_target_action", ["skip", "neg", "pos"])
+    @pytest.mark.parametrize("top_k", [None, 1, 4, 10])
     @pytest.mark.parametrize(**_default_metric_class_input_arguments_ignore_index)
     def test_class_metric_ignore_index(
         self,
@@ -68,8 +84,9 @@ def test_class_metric_ignore_index(
         preds: Tensor,
         target: Tensor,
         empty_target_action: str,
+        top_k: int,
     ):
-        metric_args = {"empty_target_action": empty_target_action, "ignore_index": -100}
+        metric_args = {"empty_target_action": empty_target_action, "ignore_index": -100, "top_k": top_k}
 
         self.run_class_metric_test(
             ddp=ddp,
@@ -77,18 +94,20 @@ def test_class_metric_ignore_index(
             preds=preds,
             target=target,
             metric_class=RetrievalMAP,
-            reference_metric=sk_average_precision_score,
+            reference_metric=_average_precision_at_k,
             metric_args=metric_args,
         )
 
     @pytest.mark.parametrize(**_default_metric_functional_input_arguments)
-    def test_functional_metric(self, preds: Tensor, target: Tensor):
+    @pytest.mark.parametrize("top_k", [None, 1, 4, 10])
+    def test_functional_metric(self, preds: Tensor, target: Tensor, top_k: int):
         self.run_functional_metric_test(
             preds=preds,
             target=target,
             metric_functional=retrieval_average_precision,
-            reference_metric=sk_average_precision_score,
+            reference_metric=_average_precision_at_k,
             metric_args={},
+            top_k=top_k,
         )
 
     @pytest.mark.parametrize(**_default_metric_class_input_arguments)