Merge pull request #62 from hotchpotch/report_to_df

Add `report.to_dataframe()` method

ranx/data_structures/report.py

@@ -3,6 +3,7 @@
 import json
 from typing import Dict, List, Tuple
 
+import pandas as pd
 from tabulate import tabulate
 
 from .frozenset_dict import FrozensetDict
@@ -319,6 +320,23 @@ def to_dict(self) -> Dict:
 
         return d
 
+    def to_dataframe(self) -> pd.DataFrame:
+        """Returns the Report data as a Pandas DataFrame.
+
+        Returns:
+            pd.DataFrame: Report data as a Pandas DataFrame
+        """
+        report_dict = self.to_dict()
+        report_scores = {
+            name: report_dict[name]["scores"] for name in report_dict["model_names"]
+        }
+        df = pd.DataFrame.from_dict(report_scores, orient="index")
+        df = df.reset_index().rename(
+            columns={"index": "model_names"}
+        )  # index to model_names column
+
+        return df
+
     def save(self, path: str):
         """Save the Report data as JSON file.
         See [**Report.to_dict**][ranx.report.to_dict] for more details.

tests/unit/ranx/data_structures/report_test.py

@@ -144,3 +144,19 @@ def test_stat_test(qrels, runs, metrics):
     assert report.stat_test == "tukey"
     assert report.get_stat_test_label(report.stat_test) == "Tukey's HSD test"
     assert report.get_stat_test_label(report.stat_test) in report.to_latex()
+
+
+def test_to_dataframe(qrels, runs, metrics):
+    report = compare(qrels, runs, metrics)
+    report_df = report.to_dataframe()
+
+    assert report_df["model_names"].tolist() == report.model_names
+    assert report_df.columns.tolist() == ["model_names"] + metrics
+
+    assert all(
+        all(
+            report_df[report_df["model_names"] == model][metric].notnull().all()
+            for metric in metrics
+        )
+        for model in report_df["model_names"]
+    )