Repository: Lightning-AI/torchmetrics

Issue:
Allow automatically adjusting `k` for `RetrievalPrecision` when `k > num(retrieved documents)`
## 🚀 Feature

Allow `RetrievalPrecision` to use `k=min(k, preds.shape[-1])` for individual queries.

### Motivation

When one calculates `precision at k` for a query over the list of documents `D`, where `len(D) < k`, there's an ambiguity whether `k` should be fixed or adjusted w.r.t. `len(D)` when `D < k`.

### Pitch

Add an input argument for an option of how to handle these edge cases. Set a default to match the current implementation.

### Alternatives

Leave it as it is.

### Additional context

The implementation described above is used by the [`catboost` package](https://github.com/catboost/catboost).


Discussion:
@lucadiliello any opinion on this?
It makes sense for me. Actually someone may want to use it because, at the moment, if you have `k` greater than the number of documents you cannot get `1` on that query even if your system is perfect.

I can implement it if you want. I would add a `adaptive_k` boolean parameter. Do you suggest other names?
`adaptive_k` is fine by me :]
I will add this argument to all the retrieval metrics which have `k` as parameter ok?
@lucadiliello sounds good to me :]

Resolving patch:
diff --git a/torchmetrics/functional/retrieval/fall_out.py b/torchmetrics/functional/retrieval/fall_out.py
--- a/torchmetrics/functional/retrieval/fall_out.py
+++ b/torchmetrics/functional/retrieval/fall_out.py
@@ -53,7 +53,7 @@ def retrieval_fall_out(preds: Tensor, target: Tensor, k: Optional[int] = None) -
     if not (isinstance(k, int) and k > 0):
         raise ValueError("`k` has to be a positive integer or None")
 
-    target = 1 - target
+    target = 1 - target  # we want to compute the probability of getting a non-relevant doc among all non-relevant docs
 
     if not target.sum():
         return tensor(0.0, device=preds.device)
diff --git a/torchmetrics/functional/retrieval/precision.py b/torchmetrics/functional/retrieval/precision.py
--- a/torchmetrics/functional/retrieval/precision.py
+++ b/torchmetrics/functional/retrieval/precision.py
@@ -18,7 +18,7 @@
 from torchmetrics.utilities.checks import _check_retrieval_functional_inputs
 
 
-def retrieval_precision(preds: Tensor, target: Tensor, k: Optional[int] = None) -> Tensor:
+def retrieval_precision(preds: Tensor, target: Tensor, k: Optional[int] = None, adaptive_k: bool = False) -> Tensor:
     """Computes the precision metric (for information retrieval). Precision is the fraction of relevant documents
     among all the retrieved documents.
 
@@ -30,13 +30,16 @@ def retrieval_precision(preds: Tensor, target: Tensor, k: Optional[int] = None)
         preds: estimated probabilities of each document to be relevant.
         target: ground truth about each document being relevant or not.
         k: consider only the top k elements (default: `None`, which considers them all)
+        adaptive_k: adjust `k` to `min(k, number of documents)` for each query
 
     Returns:
         a single-value tensor with the precision (at ``k``) of the predictions ``preds`` w.r.t. the labels ``target``.
 
     Raises:
         ValueError:
-            If ``k`` parameter is not `None` or an integer larger than 0
+            If ``k`` is not `None` or an integer larger than 0.
+        ValueError:
+            If ``adaptive_k`` is not boolean.
 
     Example:
         >>> preds = tensor([0.2, 0.3, 0.5])
@@ -46,7 +49,10 @@ def retrieval_precision(preds: Tensor, target: Tensor, k: Optional[int] = None)
     """
     preds, target = _check_retrieval_functional_inputs(preds, target)
 
-    if k is None:
+    if not isinstance(adaptive_k, bool):
+        raise ValueError("`adaptive_k` has to be a boolean")
+
+    if k is None or (adaptive_k and k > preds.shape[-1]):
         k = preds.shape[-1]
 
     if not (isinstance(k, int) and k > 0):
diff --git a/torchmetrics/retrieval/precision.py b/torchmetrics/retrieval/precision.py
--- a/torchmetrics/retrieval/precision.py
+++ b/torchmetrics/retrieval/precision.py
@@ -47,6 +47,7 @@ class RetrievalPrecision(RetrievalMetric):
         ignore_index:
             Ignore predictions where the target is equal to this number.
         k: consider only the top k elements for each query (default: `None`, which considers them all)
+        adaptive_k: adjust `k` to `min(k, number of documents)` for each query
         compute_on_step:
             Forward only calls ``update()`` and returns None if this is set to False.
 
@@ -62,7 +63,9 @@ class RetrievalPrecision(RetrievalMetric):
         ValueError:
             If ``ignore_index`` is not `None` or an integer.
         ValueError:
-            If ``k`` parameter is not `None` or an integer larger than 0.
+            If ``k`` is not `None` or an integer larger than 0.
+        ValueError:
+            If ``adaptive_k`` is not boolean.
 
     Example:
         >>> from torchmetrics import RetrievalPrecision
@@ -81,6 +84,7 @@ def __init__(
         empty_target_action: str = "neg",
         ignore_index: Optional[int] = None,
         k: Optional[int] = None,
+        adaptive_k: bool = False,
         compute_on_step: Optional[bool] = None,
         **kwargs: Dict[str, Any],
     ) -> None:
@@ -93,7 +97,10 @@ def __init__(
 
         if (k is not None) and not (isinstance(k, int) and k > 0):
             raise ValueError("`k` has to be a positive integer or None")
+        if not isinstance(adaptive_k, bool):
+            raise ValueError("`adaptive_k` has to be a boolean")
         self.k = k
+        self.adaptive_k = adaptive_k
 
     def _metric(self, preds: Tensor, target: Tensor) -> Tensor:
-        return retrieval_precision(preds, target, k=self.k)
+        return retrieval_precision(preds, target, k=self.k, adaptive_k=self.adaptive_k)


Test patch:
diff --git a/tests/retrieval/helpers.py b/tests/retrieval/helpers.py
--- a/tests/retrieval/helpers.py
+++ b/tests/retrieval/helpers.py
@@ -27,6 +27,7 @@
 from tests.retrieval.inputs import _input_retrieval_scores_empty as _irs_empty
 from tests.retrieval.inputs import _input_retrieval_scores_extra as _irs_extra
 from tests.retrieval.inputs import _input_retrieval_scores_float_target as _irs_float_tgt
+from tests.retrieval.inputs import _input_retrieval_scores_for_adaptive_k as _irs_adpt_k
 from tests.retrieval.inputs import _input_retrieval_scores_int_target as _irs_int_tgt
 from tests.retrieval.inputs import _input_retrieval_scores_mismatching_sizes as _irs_mis_sz
 from tests.retrieval.inputs import _input_retrieval_scores_mismatching_sizes_func as _irs_mis_sz_fn
@@ -162,6 +163,14 @@ def _concat_tests(*tests: Tuple[Dict]) -> Dict:
     ],
 )
 
+_errors_test_functional_metric_parameters_adaptive_k = dict(
+    argnames="preds,target,message,metric_args",
+    argvalues=[
+        (_irs.preds, _irs.target, "`adaptive_k` has to be a boolean", dict(adaptive_k=10)),
+        (_irs.preds, _irs.target, "`adaptive_k` has to be a boolean", dict(adaptive_k=None)),
+    ],
+)
+
 _errors_test_class_metric_parameters_no_pos_target = dict(
     argnames="indexes,preds,target,message,metric_args",
     argvalues=[
@@ -302,6 +311,15 @@ def _concat_tests(*tests: Tuple[Dict]) -> Dict:
     argnames="indexes,preds,target,message,metric_args",
     argvalues=[
         (_irs.index, _irs.preds, _irs.target, "`k` has to be a positive integer or None", dict(k=-10)),
+        (_irs.index, _irs.preds, _irs.target, "`k` has to be a positive integer or None", dict(k=4.0)),
+    ],
+)
+
+_errors_test_class_metric_parameters_adaptive_k = dict(
+    argnames="indexes,preds,target,message,metric_args",
+    argvalues=[
+        (_irs.index, _irs.preds, _irs.target, "`adaptive_k` has to be a boolean", dict(adaptive_k=10)),
+        (_irs.index, _irs.preds, _irs.target, "`adaptive_k` has to be a boolean", dict(adaptive_k=None)),
     ],
 )
 
@@ -311,6 +329,7 @@ def _concat_tests(*tests: Tuple[Dict]) -> Dict:
         (_irs.indexes, _irs.preds, _irs.target),
         (_irs_extra.indexes, _irs_extra.preds, _irs_extra.target),
         (_irs_no_tgt.indexes, _irs_no_tgt.preds, _irs_no_tgt.target),
+        (_irs_adpt_k.indexes, _irs_adpt_k.preds, _irs_adpt_k.target),
     ],
 )
 
diff --git a/tests/retrieval/inputs.py b/tests/retrieval/inputs.py
--- a/tests/retrieval/inputs.py
+++ b/tests/retrieval/inputs.py
@@ -26,6 +26,12 @@
     target=torch.randint(high=2, size=(NUM_BATCHES, BATCH_SIZE)),
 )
 
+_input_retrieval_scores_for_adaptive_k = Input(
+    indexes=torch.randint(high=NUM_BATCHES * BATCH_SIZE // 2, size=(NUM_BATCHES, BATCH_SIZE)),
+    preds=torch.rand(NUM_BATCHES, BATCH_SIZE),
+    target=torch.randint(high=2, size=(NUM_BATCHES, BATCH_SIZE)),
+)
+
 _input_retrieval_scores_extra = Input(
     indexes=torch.randint(high=10, size=(NUM_BATCHES, BATCH_SIZE, EXTRA_DIM)),
     preds=torch.rand(NUM_BATCHES, BATCH_SIZE, EXTRA_DIM),
diff --git a/tests/retrieval/test_precision.py b/tests/retrieval/test_precision.py
--- a/tests/retrieval/test_precision.py
+++ b/tests/retrieval/test_precision.py
@@ -22,9 +22,11 @@
     _default_metric_class_input_arguments,
     _default_metric_class_input_arguments_ignore_index,
     _default_metric_functional_input_arguments,
+    _errors_test_class_metric_parameters_adaptive_k,
     _errors_test_class_metric_parameters_default,
     _errors_test_class_metric_parameters_k,
     _errors_test_class_metric_parameters_no_pos_target,
+    _errors_test_functional_metric_parameters_adaptive_k,
     _errors_test_functional_metric_parameters_default,
     _errors_test_functional_metric_parameters_k,
 )
@@ -34,7 +36,7 @@
 seed_all(42)
 
 
-def _precision_at_k(target: np.ndarray, preds: np.ndarray, k: int = None):
+def _precision_at_k(target: np.ndarray, preds: np.ndarray, k: int = None, adaptive_k: bool = False):
     """Didn't find a reliable implementation of Precision in Information Retrieval, so, reimplementing here.
 
     A good explanation can be found
@@ -43,7 +45,7 @@ def _precision_at_k(target: np.ndarray, preds: np.ndarray, k: int = None):
     assert target.shape == preds.shape
     assert len(target.shape) == 1  # works only with single dimension inputs
 
-    if k is None:
+    if k is None or adaptive_k and k > len(preds):
         k = len(preds)
 
     if target.sum() > 0:
@@ -59,6 +61,7 @@ class TestPrecision(RetrievalMetricTester):
     @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("k", [None, 1, 4, 10])
+    @pytest.mark.parametrize("adaptive_k", [False, True])
     @pytest.mark.parametrize(**_default_metric_class_input_arguments)
     def test_class_metric(
         self,
@@ -70,8 +73,11 @@ def test_class_metric(
         empty_target_action: str,
         ignore_index: int,
         k: int,
+        adaptive_k: bool,
     ):
-        metric_args = dict(empty_target_action=empty_target_action, k=k, ignore_index=ignore_index)
+        metric_args = dict(
+            empty_target_action=empty_target_action, k=k, ignore_index=ignore_index, adaptive_k=adaptive_k
+        )
 
         self.run_class_metric_test(
             ddp=ddp,
@@ -88,6 +94,7 @@ def test_class_metric(
     @pytest.mark.parametrize("dist_sync_on_step", [True, False])
     @pytest.mark.parametrize("empty_target_action", ["skip", "neg", "pos"])
     @pytest.mark.parametrize("k", [None, 1, 4, 10])
+    @pytest.mark.parametrize("adaptive_k", [False, True])
     @pytest.mark.parametrize(**_default_metric_class_input_arguments_ignore_index)
     def test_class_metric_ignore_index(
         self,
@@ -98,8 +105,9 @@ def test_class_metric_ignore_index(
         dist_sync_on_step: bool,
         empty_target_action: str,
         k: int,
+        adaptive_k: bool,
     ):
-        metric_args = dict(empty_target_action=empty_target_action, k=k, ignore_index=-100)
+        metric_args = dict(empty_target_action=empty_target_action, k=k, ignore_index=-100, adaptive_k=adaptive_k)
 
         self.run_class_metric_test(
             ddp=ddp,
@@ -114,7 +122,8 @@ def test_class_metric_ignore_index(
 
     @pytest.mark.parametrize(**_default_metric_functional_input_arguments)
     @pytest.mark.parametrize("k", [None, 1, 4, 10])
-    def test_functional_metric(self, preds: Tensor, target: Tensor, k: int):
+    @pytest.mark.parametrize("adaptive_k", [False, True])
+    def test_functional_metric(self, preds: Tensor, target: Tensor, k: int, adaptive_k: bool):
         self.run_functional_metric_test(
             preds=preds,
             target=target,
@@ -122,6 +131,7 @@ def test_functional_metric(self, preds: Tensor, target: Tensor, k: int):
             sk_metric=_precision_at_k,
             metric_args={},
             k=k,
+            adaptive_k=adaptive_k,
         )
 
     @pytest.mark.parametrize(**_default_metric_class_input_arguments)
@@ -149,6 +159,7 @@ def test_precision_gpu(self, indexes: Tensor, preds: Tensor, target: Tensor):
             _errors_test_class_metric_parameters_default,
             _errors_test_class_metric_parameters_no_pos_target,
             _errors_test_class_metric_parameters_k,
+            _errors_test_class_metric_parameters_adaptive_k,
         )
     )
     def test_arguments_class_metric(
@@ -169,6 +180,7 @@ def test_arguments_class_metric(
         **_concat_tests(
             _errors_test_functional_metric_parameters_default,
             _errors_test_functional_metric_parameters_k,
+            _errors_test_functional_metric_parameters_adaptive_k,
         )
     )
     def test_arguments_functional_metric(self, preds: Tensor, target: Tensor, message: str, metric_args: dict):