add test_rolling_mae_vs_hull_dist() to test_plot_funcs.py

Author: janosh

mb_discovery/plot_scripts/plot_funcs.py

@@ -90,15 +90,15 @@ def hist_classified_stable_as_func_of_hull_dist(
         false_neg = e_above_hull_vals[actual_pos & model_neg]
         false_pos = e_above_hull_vals[actual_neg & model_pos]
         true_neg = e_above_hull_vals[actual_neg & model_neg]
-        xlabel = r"$\Delta E_{Hull-MP}$ / eV per atom"
+        xlabel = r"$\Delta E_{Hull-MP}$ (eV / atom)"
 
     # --- histogram by model-predicted distance to convex hull
     if energy_type == "pred":
         true_pos = residuals[actual_pos & model_pos]
         false_neg = residuals[actual_pos & model_neg]
         false_pos = residuals[actual_neg & model_pos]
         true_neg = residuals[actual_neg & model_neg]
-        xlabel = r"$\Delta E_{Hull-Pred}$ / eV per atom"
+        xlabel = r"$\Delta E_{Hull-Pred}$ (eV / atom)"
 
     ax.hist(
         [true_pos, false_neg, false_pos, true_neg],
@@ -153,7 +153,7 @@ def rolling_mae_vs_hull_dist(
     e_above_hull_col: str,
     residual_col: str = "residual",
     half_window: float = 0.02,
-    increment: float = 0.002,
+    bin_width: float = 0.002,
     x_lim: tuple[float, float] = (-0.2, 0.3),
     ax: plt.Axes = None,
     **kwargs: Any,
@@ -174,7 +174,7 @@ def rolling_mae_vs_hull_dist(
 
     is_fresh_ax = len(ax.lines) == 0
 
-    bins = np.arange(*x_lim, increment)
+    bins = np.arange(*x_lim, bin_width)
 
     rolling_maes = np.zeros_like(bins)
     rolling_stds = np.zeros_like(bins)
@@ -254,7 +254,7 @@ def rolling_mae_vs_hull_dist(
 
     ax.text(0, 0.13, r"$|\Delta E_{Hull-MP}| > $MAE", horizontalalignment="center")
 
-    ax.set(xlabel=r"$\Delta E_{Hull-MP}$ / eV per atom", ylabel="MAE / eV per atom")
+    ax.set(xlabel=r"$\Delta E_{Hull-MP}$ (eV / atom)", ylabel="MAE (eV / atom)")
 
     ax.set(xlim=x_lim, ylim=(0.0, 0.14))
 
@@ -389,7 +389,7 @@ def precision_recall_vs_calc_count(
         # previous call
         return ax
 
-    ax.set(xlabel="Number of Calculations", ylabel="Percentage")
+    ax.set(xlabel="Number of Calculations", ylabel="Precision and Recall (%)")
 
     ax.set(ylim=(0, 100))
 

tests/test_plot_funcs.py

@@ -2,11 +2,15 @@
 
 from typing import Any, Sequence
 
+import matplotlib.pyplot as plt
 import pandas as pd
 import pytest
 
 from mb_discovery import ROOT
-from mb_discovery.plot_scripts.plot_funcs import precision_recall_vs_calc_count
+from mb_discovery.plot_scripts.plot_funcs import (
+    precision_recall_vs_calc_count,
+    rolling_mae_vs_hull_dist,
+)
 
 
 DATA_DIR = f"{ROOT}/data/2022-06-11-from-rhys"
@@ -28,9 +32,9 @@
     "intersect_lines, stability_crit, stability_threshold, expected_line_count",
     [
         ((), "energy", 0, 11),
-        ("precision_x", "energy+std", 0, 23),
-        (["recall_y"], "energy", -0.1, 35),
-        ("all", "energy-std", 0.1, 56),
+        ("precision_x", "energy+std", 0, 14),
+        (["recall_y"], "energy", -0.1, 14),
+        ("all", "energy-std", 0.1, 23),
     ],
 )
 def test_precision_recall_vs_calc_count(
@@ -39,7 +43,7 @@ def test_precision_recall_vs_calc_count(
     stability_threshold: float,
     expected_line_count: int,
 ) -> None:
-    ax = None
+    ax = plt.figure().gca()  # ensure test functions use different axes
 
     for (model_name, df), color in zip(
         test_dfs.items(), ("tab:blue", "tab:orange", "tab:pink")
@@ -66,6 +70,9 @@ def test_precision_recall_vs_calc_count(
     assert ax.get_ylim() == (0, 100)
     assert ax.get_xlim() == pytest.approx((-1.4, 29.4))
 
+    assert ax.get_xlabel() == "Number of Calculations"
+    assert ax.get_ylabel() == "Precision and Recall (%)"
+
 
 @pytest.mark.parametrize(
     "kwargs, expected_exc, match_pat",
@@ -84,3 +91,37 @@ def test_precision_recall_vs_calc_count_raises(
             e_above_hull_col="e_above_mp_hull",
             **kwargs,
         )
+
+
+@pytest.mark.parametrize("half_window", (0.02, 0.002))
+@pytest.mark.parametrize("bin_width", (0.1, 0.001))
+@pytest.mark.parametrize("x_lim", ((0, 0.6), (-0.2, 0.8)))
+def test_rolling_mae_vs_hull_dist(
+    half_window: float, bin_width: float, x_lim: tuple[float, float]
+) -> None:
+    ax = plt.figure().gca()  # ensure test functions use different axes
+
+    for (model_name, df), color in zip(
+        test_dfs.items(), ("tab:blue", "tab:orange", "tab:pink")
+    ):
+        model_preds = df.filter(like=r"_pred").mean(axis=1)
+        targets = df.e_form_target
+
+        df["residual"] = model_preds - targets + df.e_above_mp_hull
+
+        ax = rolling_mae_vs_hull_dist(
+            df,
+            residual_col="residual",
+            e_above_hull_col="e_above_mp_hull",
+            color=color,
+            label=model_name,
+            ax=ax,
+            x_lim=x_lim,
+            half_window=half_window,
+            bin_width=bin_width,
+        )
+
+    assert ax is not None
+    assert ax.get_ylim() == pytest.approx((0, 0.14))
+    assert ax.get_ylabel() == "MAE (eV / atom)"
+    assert ax.get_xlabel() == r"$\Delta E_{Hull-MP}$ (eV / atom)"