diff --git a/great_ai/utilities/evaluate_ranking/evaluate_ranking.py b/great_ai/utilities/evaluate_ranking/evaluate_ranking.py index bb6ce34..4f297b0 100644 --- a/great_ai/utilities/evaluate_ranking/evaluate_ranking.py +++ b/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 diff --git a/tests/utilities/test_evaluate_ranking.py b/tests/utilities/test_evaluate_ranking.py index 23cf14d..c1ababe 100644 --- a/tests/utilities/test_evaluate_ranking.py +++ b/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()