Render ranking exports and save the figure owning the supplied axes

great_ai/utilities/evaluate_ranking/evaluate_ranking.py

@@ -1,9 +1,10 @@
 from pathlib import Path
-from typing import Dict, List, Optional, TypeVar
+from typing import Dict, List, Optional, TypeVar, cast
 
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
+from matplotlib.figure import Figure
 from sklearn.metrics import average_precision_score, precision_recall_curve
 
 from ..unique import unique
@@ -45,7 +46,8 @@ def evaluate_ranking(
 
     assert 0 <= target_recall <= 1
 
-    if plot and axes is None:
+    render = plot or output_svg is not None
+    if render and axes is None:
         fig = plt.figure(figsize=(10, 10))
         fig.patch.set_facecolor("white")
         ax = plt.axes()
@@ -56,7 +58,7 @@ def evaluate_ranking(
     str_classes = [str(c) for c in classes]
 
     with matplotlib.rc_context({"font.size": 20}):
-        if plot:
+        if render:
             ax.set_xmargin(0)
 
             draw_f1_iso_lines(axes=ax)
@@ -87,14 +89,14 @@ def evaluate_ranking(
             )
             results[classes[i]] = precision_at_closest_recall
 
-            if plot:
+            if render:
                 ax.plot(
                     recall,
                     precision,
                     label=f"{'|'.join(str_classes[:i + 1])} ↔ {'|'.join(str_classes[i+1:])} (P@{target_recall:.2f}={precision_at_closest_recall:.2f}, AP={average_precision:.2f})",
                 )
 
-        if plot:
+        if render:
             ax.legend(loc="upper right")
             ax.axvline(x=target_recall, linestyle="--", color="#55c6bb", linewidth=2.0)
 
@@ -112,6 +114,8 @@ def evaluate_ranking(
             if axes is None:
                 plt.show()
         elif output_svg:
-            plt.savefig(output_svg, format="svg")
+            cast(Figure, ax.get_figure(root=True)).savefig(output_svg, format="svg")
+            if not plot and axes is None:
+                plt.close(ax.get_figure(root=True))
 
     return results

tests/utilities/test_evaluate_ranking.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 
 import matplotlib
+import matplotlib.pyplot as plt
 import pytest
 
 matplotlib.use("Agg")  # don't show a window for each test
@@ -108,3 +109,32 @@ def test_save() -> None:
 
     assert path.exists()
     path.unlink()
+
+
+def test_export_saves_the_supplied_axes_figure(tmp_path: Path) -> None:
+    intended, axes = plt.subplots()
+    unrelated, other = plt.subplots()
+    other.set_title("unrelated figure")
+    output = tmp_path / "ranking.svg"
+    try:
+        evaluate_ranking(
+            [0, 1],
+            [0.1, 0.9],
+            target_recall=1,
+            axes=axes,
+            output_svg=output,
+            title="intended ranking",
+        )
+        content = output.read_text()
+        assert "intended ranking" in content
+        assert "unrelated figure" not in content
+    finally:
+        plt.close(intended)
+        plt.close(unrelated)
+
+
+def test_export_without_display_still_renders(tmp_path: Path) -> None:
+    output = tmp_path / "ranking.svg"
+    evaluate_ranking([0, 1], [0.1, 0.9], target_recall=1, plot=False, output_svg=output)
+    assert "Precision" in output.read_text()
+    assert "Recall" in output.read_text()