refactor hist_clf_vary.py

Author: janosh

ml_stability/stability_plot_scripts/hist_clf_vary.py

@@ -6,7 +6,7 @@
 import pandas as pd
 from scipy.interpolate import interp1d
 
-from ml_stability import ROOT
+from ml_stability import PKG_DIR, ROOT
 
 
 __author__ = "Rhys Goodall, Janosh Riebesell"
@@ -34,23 +34,23 @@
     ["wren", "voronoi", "cgcnn"],
     ["tab:blue", "tab:orange", "tab:red"],
 ):
-    df = pd.read_csv(
+    data_path = (
         f"{ROOT}/data/2022-06-11-from-rhys/{model_name}-mp-initial-structures.csv"
     )
-    df = df.set_index("material_id")
+    df = pd.read_csv(data_path).set_index("material_id")
 
     df["e_above_hull"] = df_hull.e_above_hull
 
     df = df.dropna(subset=["e_above_hull"])
 
     rare = "all"
 
-    # rare = "nla"
-    # df = df[
-    #     ~df["composition"].apply(
-    #         lambda x: any(el.is_rare_earth_metal for el in Composition(x).elements)
-    #     )
-    # ]
+    # from pymatgen.core import Composition
+    # rare = "no-lanthanides"
+    # df["contains_rare_earths"] = df.composition.map(
+    #     lambda x: any(el.is_rare_earth_metal for el in Composition(x))
+    # )
+    # df = df.query("~contains_rare_earths")
 
     e_above_hull = df.e_above_hull.to_numpy().ravel()
 
@@ -62,7 +62,7 @@
 
     # epistemic_std = np.var(pred, axis=0, ddof=0)
 
-    # aleatoric_std = np.mean(np.square(df.filter(like="ale")), axis=0)
+    aleatoric_std = (df.filter(like="ale") ** 2).mean(axis=0) ** 0.5
 
     # full_std = np.sqrt(epistemic_std + aleatoric_std)
 
@@ -73,7 +73,6 @@
     # test = mean - full_std
 
     crit = "ene"
-    test = mean
 
     bins = 200
     # xlim = (-0.2, 0.2)
@@ -86,20 +85,20 @@
     xticks = (-0.4, -0.2, 0, 0.2, 0.4)
     # yticks = (0, 300, 600, 900, 1200)
 
-    tp = len(e_above_hull[(e_above_hull <= thresh) & (test <= thresh)])
-    fn = len(e_above_hull[(e_above_hull <= thresh) & (test > thresh)])
+    tp = len(e_above_hull[(e_above_hull <= thresh) & (mean <= thresh)])
+    fn = len(e_above_hull[(e_above_hull <= thresh) & (mean > thresh)])
 
     pos = tp + fn
 
-    sort = np.argsort(test)
+    sort = np.argsort(mean)
     e_above_hull = e_above_hull[sort]
-    test = test[sort]
+    mean = mean[sort]
 
     e_type = "pred"
-    tp = np.asarray((e_above_hull <= thresh) & (test <= thresh))
-    fn = np.asarray((e_above_hull <= thresh) & (test > thresh))
-    fp = np.asarray((e_above_hull > thresh) & (test <= thresh))
-    tn = np.asarray((e_above_hull > thresh) & (test > thresh))
+    tp = np.asarray((e_above_hull <= thresh) & (mean <= thresh))
+    fn = np.asarray((e_above_hull <= thresh) & (mean > thresh))
+    fp = np.asarray((e_above_hull > thresh) & (mean <= thresh))
+    tn = np.asarray((e_above_hull > thresh) & (mean > thresh))
     xlabel = r"$\Delta E_{Hull-Pred}$ / eV per atom"
 
     c_tp = np.cumsum(tp)
@@ -114,49 +113,30 @@
 
     x = np.arange(len(ppv))[:end]
 
-    f_ppv = interp1d(x, ppv[:end], kind="cubic")
-    f_tpr = interp1d(x, tpr[:end], kind="cubic")
+    precision_curve = interp1d(x, ppv[:end], kind="cubic")
+    rolling_recall_curve = interp1d(x, tpr[:end], kind="cubic")
 
     line_kwargs = dict(
         linewidth=2, color=color, markevery=[-1], marker="x", markersize=14, mew=2.5
     )
-    ax.plot(x[::100], f_tpr(x[::100]), linestyle=":", **line_kwargs)
-
-    ax.plot(x[::100], f_ppv(x[::100]), linestyle="-", **line_kwargs)
+    ax.plot(x[::100], precision_curve(x[::100]), linestyle="-", **line_kwargs)
+    ax.plot(x[::100], rolling_recall_curve(x[::100]), linestyle=":", **line_kwargs)
 
 
 ax.set(xlabel="Number of Calculations", ylabel="Percentage")
 
-ax.set(xlim=(0, 8e4), ylim=(0, 100), xticks=(0, 2e4, 4e4, 6e4, 8e4))
-
-ax.plot((-1, -1), (-1, -1), color="tab:blue")
-ax.plot((-1, -1), (-1, -1), color="tab:red")
-ax.plot((-1, -1), (-1, -1), color="tab:orange")
-
-# ax.plot((-1, -1), (-1, -1), color="tab:purple")
-
-ax.plot((-1, -1), (-1, -1), "k", linestyle="-")
-ax.plot((-1, -1), (-1, -1), "k", linestyle=":")
-
-lines = ax.get_lines()
-
-legend1 = ax.legend(
-    lines[-2:], ["Precision", "Recall"], frameon=False, loc="upper right"
-)
-legend2 = ax.legend(
-    lines[-5:-2],
-    # ["Wren (This Work)", "CGCNN Pre-relax", "CGCNN-D Pre-relax"],
-    ["Wren (This Work)", "CGCNN Pre-relax", "Voronoi Pre-relax"],
-    frameon=False,
-    loc="lower right",
-)
-
-ax.add_artist(legend1)
-
-ax.set_aspect(1.0 / ax.get_data_ratio())
+xlim = (0, 8e4)
+ax.set(xlim=xlim, ylim=(0, 100), xticks=np.linspace(*xlim, 5))
 
+ax.plot((0, 0), (0, 0), color="tab:blue", label="Wren (This Work)")
+ax.plot((0, 0), (0, 0), color="tab:red", label="CGCNN Pre-relax")
+ax.plot((0, 0), (0, 0), color="tab:orange", label="Voronoi Pre-relax")
+legend_1 = ax.legend(frameon=False, loc="lower right")
+ax.add_artist(legend_1)
 
-# plt.savefig(f"{PKG_DIR}/plots/{today}-vary-{e_type}-{crit}-{rare}.pdf")
-# # plt.savefig(f"{PKG_DIR}/plots/{today}-vary-{e_type}-{crit}-{rare}.png")
+[prec] = ax.plot((0, 0), (0, 0), "black", linestyle="-")
+[recall] = ax.plot((0, 0), (0, 0), "black", linestyle=":")
+ax.legend([prec, recall], ["Precision", "Recall"], frameon=False, loc="upper right")
 
-# plt.show()
+img_path = f"{PKG_DIR}/plots/{today}-vary-{e_type=}-{crit=}-{rare=}.pdf"
+# plt.savefig(img_path)