Repository: terrierteam/ir_measures (language: python)

Issue:
rank_eval
Another evaluation library appearing at ECIR'22, would be nice to add as another provider: https://github.com/AmenRa/rank_eval

It doesn't include any measures not yet supported by other libraries (except `Hits`, which is similar to `NumRelRet` but averaging instead of summing), but it reports being faster than others.

Available measures:
 - `P@k`
 - `R@k`
 - `RPrec`
 - `RR`
 - `AP`
 - `nDCG` (both `dcg='log2'` and `dcg='exp-log2'`)
 - `Hits@k` (new)


Discussion / hints:

Fix patch:
diff --git a/docs/measures.rst b/docs/measures.rst
index 5f38f33..3f0836d 100644
--- a/docs/measures.rst
+++ b/docs/measures.rst
@@ -112,6 +112,7 @@ by taking the mean of AP over the query set.
 - ``cwl_eval``: ``AP(rel=ANY)@NOT_PROVIDED``
 - ``pytrec_eval``: ``AP(rel=ANY)@ANY``
 - ``trectools``: ``AP(rel=1)@ANY``
+- ``ranx``: ``AP(rel=ANY)@ANY``
 
 
 
@@ -421,6 +422,7 @@ It is normalized wrt. the Ideal NDCG, i.e. documents ranked in descending order
 - ``pytrec_eval``: ``nDCG(dcg='log2')@ANY``
 - ``gdeval``: ``nDCG(dcg='exp-log2')@REQUIRED``
 - ``trectools``: ``nDCG(dcg=ANY)@ANY``
+- ``ranx``: ``nDCG(dcg=('log2', 'exp-log2'))@ANY``
 
 
 
@@ -668,6 +670,7 @@ returned with at least that relevance score (inclusive).
 **Provided by:**
 
 - ``pytrec_eval``: ``NumRet(rel=ANY)``
+- ``ranx``: ``NumRet(rel=REQUIRED)``
 
 
 
@@ -699,6 +702,7 @@ P@cutoff.
 - ``cwl_eval``: ``P(rel=ANY)@ANY``
 - ``pytrec_eval``: ``P(rel=ANY)@ANY``
 - ``trectools``: ``P(rel=1)@ANY``
+- ``ranx``: ``P(rel=ANY)@ANY``
 
 
 
@@ -739,6 +743,7 @@ This software follows the TREC convention and refers to that measure as Success@
 **Provided by:**
 
 - ``pytrec_eval``: ``R@ANY``
+- ``ranx``: ``R@ANY``
 
 
 
@@ -798,6 +803,7 @@ it is also the recall at R.
 
 - ``pytrec_eval``: ``Rprec(rel=ANY)``
 - ``trectools``: ``Rprec(rel=1)``
+- ``ranx``: ``Rprec(rel=ANY)``
 
 
 
@@ -833,6 +839,7 @@ depth explored. rel (default 1) controls which relevance level is considered rel
 - ``pytrec_eval``: ``RR(rel=ANY)@NOT_PROVIDED``
 - ``trectools``: ``RR(rel=1)@NOT_PROVIDED``
 - ``msmarco``: ``RR(rel=ANY)@ANY``
+- ``ranx``: ``RR(rel=ANY)@NOT_PROVIDED``
 
 
 
@@ -907,6 +914,7 @@ The Set Precision (SetP); i.e., the number of relevant docs divided by the total
 **Provided by:**
 
 - ``pytrec_eval``: ``SetP(rel=ANY, relative=ANY)``
+- ``ranx``: ``SetP(rel=ANY)``
 
 
 
@@ -924,6 +932,7 @@ The Set Recall (SetR); i.e., the number of relevant docs divided by the total nu
 **Provided by:**
 
 - ``pytrec_eval``: ``SetR(rel=ANY)``
+- ``ranx``: ``SetR(rel=ANY)``
 
 
 
@@ -963,6 +972,7 @@ Recall@k is defined as the proportion of known relevant documents retrieved in t
 **Provided by:**
 
 - ``pytrec_eval``: ``Success(rel=ANY)@ANY``
+- ``ranx``: ``Success(rel=ANY)@REQUIRED``
 
 
 
diff --git a/docs/providers.rst b/docs/providers.rst
index 24be84f..584f7a2 100644
--- a/docs/providers.rst
+++ b/docs/providers.rst
@@ -165,6 +165,40 @@ Provider for the accuracy metric
 
 
 
+``ranx``
+-------------------------
+
+
+ ranx
+
+ https://amenra.github.io/ranx/
+
+::
+
+ @misc{ranx2021,
+   title = {ranx: A Blazing-Fast Python Library for Ranking Evaluation and Comparison},
+   author = {Bassani, Elias},
+   year = {2021},
+   publisher = {GitHub},
+   howpublished = {\url{https://github.com/AmenRa/ranx}},
+ }
+
+ 
+**Supported Measures:**
+
+ - ``P(rel=ANY)@ANY``
+ - ``SetP(rel=ANY)``
+ - ``RR(rel=ANY)@NOT_PROVIDED``
+ - ``Rprec(rel=ANY)``
+ - ``AP(rel=ANY)@ANY``
+ - ``nDCG(dcg=('log2', 'exp-log2'))@ANY``
+ - ``R@ANY``
+ - ``SetR(rel=ANY)``
+ - ``NumRet(rel=REQUIRED)``
+ - ``Success(rel=ANY)@REQUIRED``
+
+
+
 ``trectools``
 -------------------------
 
diff --git a/ir_measures/__init__.py b/ir_measures/__init__.py
index c33cfd9..20e0b0c 100644
--- a/ir_measures/__init__.py
+++ b/ir_measures/__init__.py
@@ -33,6 +33,7 @@
 judged = providers.registry['judged']
 msmarco = providers.registry['msmarco']
 pyndeval = providers.registry['pyndeval']
+ranx = providers.registry['ranx']
 
 CwlMetric = providers.CwlMetric
 
@@ -46,6 +47,7 @@
     msmarco,
     gdeval,  # doesn't work when installed from package #9
     accuracy,
+    ranx,
 ])
 evaluator = DefaultPipeline.evaluator
 calc_ctxt = DefaultPipeline.calc_ctxt # deprecated; replaced with evaluator
diff --git a/ir_measures/providers/__init__.py b/ir_measures/providers/__init__.py
index 09f8263..0a36e28 100644
--- a/ir_measures/providers/__init__.py
+++ b/ir_measures/providers/__init__.py
@@ -14,3 +14,4 @@ def register(provider):
 from .gdeval_provider import GdevalProvider
 from .trectools_provider import TrectoolsProvider
 from .msmarco_provider import MsMarcoProvider
+from .ranx_provider import RanxProvider
diff --git a/ir_measures/providers/ranx_provider.py b/ir_measures/providers/ranx_provider.py
new file mode 100644
index 0000000..0cae3cf
--- /dev/null
+++ b/ir_measures/providers/ranx_provider.py
@@ -0,0 +1,153 @@
+import contextlib
+import ir_measures
+from ir_measures import providers, measures, Metric
+from ir_measures.providers.base import Any, Choices, NOT_PROVIDED
+
+
+class RanxProvider(providers.Provider):
+    """
+    ranx
+
+    https://amenra.github.io/ranx/
+
+::
+
+    @misc{ranx2021,
+      title = {ranx: A Blazing-Fast Python Library for Ranking Evaluation and Comparison},
+      author = {Bassani, Elias},
+      year = {2021},
+      publisher = {GitHub},
+      howpublished = {\\url{https://github.com/AmenRa/ranx}},
+    }
+
+    """
+    NAME = 'ranx'
+    SUPPORTED_MEASURES = [
+        measures._P(cutoff=Any(), rel=Any()),
+        measures._SetP(rel=Any()),
+        measures._RR(cutoff=Choices(NOT_PROVIDED), rel=Any()),
+        measures._Rprec(rel=Any()),
+        measures._AP(cutoff=Any(), rel=Any()),
+        measures._nDCG(cutoff=Any(), dcg=Choices('log2', 'exp-log2')),
+        measures._R(cutoff=Any()),
+        measures._SetR(rel=Any()),
+        measures._NumRet(rel=Any(required=True)),
+        measures._Success(cutoff=Any(required=True), rel=Any()),
+    ]
+
+    def __init__(self):
+        super().__init__()
+        self.ranx = None
+
+    def _evaluator(self, measures, qrels):
+        measures = ir_measures.util.flatten_measures(measures)
+        # Convert qrels to dict_of_dict (input format used by pytrec_eval)
+        qrels = ir_measures.util.QrelsConverter(qrels).as_pd_dataframe()
+        qids = set(qrels['query_id'].unique())
+
+        # Depending on the measure params, we may need multiple invocations of ranx
+        # (e.g., with different rel_level, since it only supports running with 1 rel_level at a time)
+        invokers = self._build_invokers(measures, qrels)
+        return RanxEvaluator(self.ranx, measures, invokers, qrels, qids=qids)
+
+    def _build_invokers(self, measures, qrels):
+        invocations = {}
+        setf_count = 0
+        for measure in measures:
+            match_str = None
+            if measure.NAME == 'P':
+                invocation_key = (measure['rel'], 0)
+                measure_str = f'precision@{measure["cutoff"]}'
+            elif measure.NAME == 'SetP':
+                invocation_key = (measure['rel'], 0)
+                measure_str = f'precision'
+            elif measure.NAME == 'R':
+                invocation_key = (measure['rel'], 0)
+                measure_str = f'recall@{measure["cutoff"]}'
+            elif measure.NAME == 'SetR':
+                invocation_key = (measure['rel'], 0)
+                measure_str = f'recall'
+            elif measure.NAME == 'RR':
+                invocation_key = (measure['rel'], 0)
+                if 'cutoff' in measure.params:
+                    measure_str = f'mrr@{measure["cutoff"]}'
+                else:
+                    measure_str = f'mrr'
+            elif measure.NAME == 'AP':
+                invocation_key = (measure['rel'], 0)
+                if 'cutoff' in measure.params:
+                    measure_str = f'map@{measure["cutoff"]}'
+                else:
+                    measure_str = f'map'
+            elif measure.NAME == 'Success':
+                invocation_key = (measure['rel'], 0)
+                measure_str = f'hit_rate@{measure["cutoff"]}'
+            elif measure.NAME == 'NumRet':
+                invocation_key = (measure['rel'], 0)
+                measure_str = f'hits'
+            elif measure.NAME == 'nDCG':
+                invocation_key = (None, 0)
+                name = 'ndcg_burges' if measure.params.get('dcg', measure.SUPPORTED_PARAMS['dcg'].default) == 'exp-log2' else 'ndcg'
+                if 'cutoff' in measure.params:
+                    measure_str = f'{name}@{measure["cutoff"]}'
+                else:
+                    measure_str = name
+            elif measure.NAME == 'Rprec':
+                invocation_key = (measure['rel'], 0)
+                measure_str = f'r-precision'
+            else:
+                raise ValueError(f'unsupported measure {measure}')
+
+            if match_str is None:
+                match_str = measure_str
+
+            if invocation_key not in invocations:
+                invocations[invocation_key] = {}
+            invocations[invocation_key][match_str] = (measure, measure_str)
+
+        invokers = []
+        for (rel_level, it), measure_map in invocations.items():
+            if rel_level is not None:
+                these_qrels = qrels.assign(relevance=(qrels['relevance']>=rel_level).astype(int))
+            else:
+                these_qrels = qrels
+            these_qrels = self.ranx.Qrels.from_df(these_qrels, q_id_col='query_id', doc_id_col='doc_id', score_col='relevance')
+            invokers.append(RanxInvoker(self.ranx, these_qrels, measure_map))
+
+        return invokers
+
+    def initialize(self):
+        try:
+            import ranx
+            self.ranx = ranx
+        except ImportError as ex:
+            raise RuntimeError('ranx not available (do you need to `pip install ranx`?)', ex)
+
+
+class RanxEvaluator(providers.Evaluator):
+    def __init__(self, ranx, measures, invokers, qrels, qids):
+        super().__init__(measures, qids)
+        self.ranx = ranx
+        self.invokers = invokers
+
+    def _iter_calc(self, run):
+        run = self.ranx.Run.from_df(ir_measures.util.RunConverter(run).as_pd_dataframe(), q_id_col='query_id', doc_id_col='doc_id', score_col='score')
+        for invoker in self.invokers:
+            yield from invoker.iter_calc(run)
+
+
+class RanxInvoker:
+    def __init__(self, ranx, qrels, measure_map):
+        self.ranx = ranx
+        self.qrels = qrels
+        self.measure_map = measure_map
+
+    def iter_calc(self, run):
+        self.ranx.evaluate(self.qrels, run, list(self.measure_map))
+        for measure, qid_value_map in run.scores.items():
+            for query_id, value in qid_value_map.items():
+                yield Metric(query_id=query_id, measure=self.measure_map[measure][0], value=value)
+        run.scores.clear()
+
+
+providers.register(RanxProvider())
diff --git a/setup.py b/setup.py
index c7d4e5d..3f2a4e8 100644
--- a/setup.py
+++ b/setup.py
@@ -35,6 +35,7 @@ def get_version(rel_path):
     install_requires=list(open('requirements.txt')),
     extras_require={
         "trectools": ["trectools>=0.0.44"],
+        "ranx": ["ranx>=0.1.11"],
     },
     classifiers=[],
     python_requires='>=3.6',


Test patch:
diff --git a/requirements-test.txt b/requirements-test.txt
index 0ef2283..8d6863a 100644
--- a/requirements-test.txt
+++ b/requirements-test.txt
@@ -1,1 +1,2 @@
-trectools>=0.0.44
\ No newline at end of file
+trectools>=0.0.44
+ranx>=0.1.11
diff --git a/test/test_ranx.py b/test/test_ranx.py
new file mode 100644
index 0000000..e388c55
--- /dev/null
+++ b/test/test_ranx.py
@@ -0,0 +1,144 @@
+import unittest
+import itertools
+import ir_measures
+from ir_measures import Metric
+
+
+class TestRanx(unittest.TestCase):
+
+    def test_P(self):
+        qrels = list(ir_measures.read_trec_qrels('''
+0 0 D0 0
+0 0 D1 1
+0 0 D2 1
+0 0 D3 2
+0 0 D4 0
+'''))
+        run = list(ir_measures.read_trec_run('''
+0 0 D0 1 0.8 run
+0 0 D2 2 0.7 run
+0 0 D1 3 0.3 run
+0 0 D3 4 0.4 run
+0 0 D4 5 0.1 run
+'''))
+        measure = ir_measures.P@5
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.6))
+        measure = ir_measures.P(rel=2)@5
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.2))
+
+        measure = ir_measures.SetP
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.6))
+        measure = ir_measures.SetP(rel=2)@5
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.2))
+
+        measure = ir_measures.R@5
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 1.0))
+        measure = ir_measures.R(rel=2)@5
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 1.0))
+        measure = ir_measures.R@2
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.3333333333333333))
+        measure = ir_measures.R(rel=2)@2
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.0))
+
+        measure = ir_measures.SetR
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 1.0))
+        measure = ir_measures.SetR(rel=2)@5
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 1.0))
+
+        measure = ir_measures.RR
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.5))
+        measure = ir_measures.RR(rel=2)
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.3333333333333333))
+        measure = ir_measures.RR@10
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.5))
+        measure = ir_measures.RR(rel=2)@10
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.3333333333333333))
+        measure = ir_measures.RR@2
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.5))
+        measure = ir_measures.RR(rel=2)@2
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.))
+
+        measure = ir_measures.AP
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.6388888888888888))
+        measure = ir_measures.AP(rel=2)
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.3333333333333333))
+        measure = ir_measures.AP@10
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.6388888888888888))
+        measure = ir_measures.AP(rel=2)@10
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.3333333333333333))
+        measure = ir_measures.AP@2
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.16666666666666666))
+        measure = ir_measures.AP(rel=2)@2
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.0))
+
+        measure = ir_measures.Success@10
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 1.))
+        measure = ir_measures.Success(rel=2)@10
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 1.))
+        measure = ir_measures.Success@2
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 1.))
+        measure = ir_measures.Success(rel=2)@2
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.))
+
+        measure = ir_measures.NumRet(rel=1)
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 3.))
+        measure = ir_measures.NumRet(rel=2)
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+
+        measure = ir_measures.nDCG
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.6584645692843067))
+        measure = ir_measures.nDCG@5
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.6584645692843067))
+        measure = ir_measures.nDCG@2
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.23981246656813146))
+
+        measure = ir_measures.nDCG(dcg='exp-log2')
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.6201040599710453))
+        measure = ir_measures.nDCG(dcg='exp-log2')@5
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.6201040599710453))
+        measure = ir_measures.nDCG(dcg='exp-log2')@2
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.17376534287144002))
+
+        measure = ir_measures.Rprec
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.6666666666666666))
+        measure = ir_measures.Rprec(rel=2)
+        result = list(ir_measures.ranx.iter_calc([measure], qrels, run))
+        self.assertEqual(result[0], Metric('0', measure, 0.))
+
+
+if __name__ == '__main__':
+    unittest.main()