add rolling_accuracy kwarg and doc str to hist_classified_stable_as_func_of_hull_dist()

extend global plot settings:
plt.rc("legend", title_fontsize=16)
plt.rc("axes", titlesize=16, labelsize=16)

Author: janosh

matbench_discovery/plot_scripts/hist_classified_stable_as_func_of_hull_dist.py

@@ -30,17 +30,19 @@
 
 # %%
 df = pd.read_csv(
-    f"{ROOT}/data/2022-06-11-from-rhys/wren-mp-initial-structures.csv"
+    # f"{ROOT}/data/2022-06-11-from-rhys/wren-mp-initial-structures.csv"
+    f"{ROOT}/models/wrenformer/mp/2022-09-20-wrenformer-e_form-ensemble-1-preds.csv"
 ).set_index("material_id")
 
-df["e_above_hull_mp"] = df_wbm.e_above_hull_mp
+df["e_above_hull"] = df_wbm.e_above_hull_mp2020_corrected_ppd_mp
 
 
 # %%
 nan_counts = df.isna().sum()
 assert all(nan_counts == 0), f"df should not have missing values: {nan_counts}"
 
-target_col = "e_form_target"
+# target_col = "e_form_target"
+target_col = "e_form_per_atom"
 stability_crit: StabilityCriterion = "energy"
 which_energy: WhichEnergy = "true"
 
@@ -57,19 +59,24 @@
 pred_cols = df.filter(regex=r"_pred_\d").columns
 assert len(pred_cols) == 10
 
-ax = hist_classified_stable_as_func_of_hull_dist(
+ax, metrics = hist_classified_stable_as_func_of_hull_dist(
     e_above_hull_pred=df[pred_cols].mean(axis=1) - df[target_col],
-    e_above_hull_true=df.e_above_hull_mp,
+    e_above_hull_true=df.e_above_hull,
     which_energy=which_energy,
     stability_crit=stability_crit,
     std_pred=std_total,
     # stability_threshold=-0.05,
+    # rolling_accuracy=0,
 )
 
 fig = ax.figure
 fig.set_size_inches(10, 9)
 
-ax.legend(loc="center left", frameon=False)
+ax.legend(
+    loc="center left",
+    frameon=False,
+    title=f"Enrichment Factor = {metrics['enrichment']:.3}",
+)
 
 fig_name = f"wren-wbm-hull-dist-hist-{which_energy=}-{stability_crit=}"
 # fig.savefig(f"{ROOT}/figures/{today}-{fig_name}.pdf")

matbench_discovery/plot_scripts/hist_classified_stable_as_func_of_hull_dist_batches.py

@@ -39,8 +39,7 @@
     f"{ROOT}/models/m3gnet/2022-08-16-m3gnet-wbm-IS2RE.json.gz"
 ).set_index("material_id")
 dfs["wrenformer"] = pd.read_csv(
-    f"{ROOT}/models/wrenformer/mp/"
-    "2022-09-20-wrenformer-e_form-ensemble-1-preds-e_form_per_atom.csv"
+    f"{ROOT}/models/wrenformer/mp/2022-09-20-wrenformer-e_form-ensemble-1-preds.csv"
 ).set_index("material_id")
 dfs["bowsr_megnet"] = pd.read_json(
     f"{ROOT}/models/bowsr/2022-09-22-bowsr-megnet-wbm-IS2RE.json.gz"
@@ -78,26 +77,32 @@
     batch_df = df[df.index.str.startswith(f"wbm-step-{batch_idx}-")]
     assert 1e4 < len(batch_df) < 1e5, print(f"{len(batch_df) = :,}")
 
-    hist_classified_stable_as_func_of_hull_dist(
+    ax, metrics = hist_classified_stable_as_func_of_hull_dist(
         e_above_hull_pred=batch_df.e_form_per_atom_pred - batch_df.e_form_per_atom,
         e_above_hull_true=batch_df.e_above_hull_mp,
         which_energy=which_energy,
         stability_crit=stability_crit,
         ax=ax,
     )
 
+    text = f"Enrichment\nFactor = {metrics['enrichment']:.3}"
+    ax.text(0.02, 0.25, text, fontsize=16, transform=ax.transAxes)
+
     title = f"Batch {batch_idx} ({len(batch_df.filter(like='e_').dropna()):,})"
     ax.set(title=title)
 
 
-hist_classified_stable_as_func_of_hull_dist(
+ax, metrics = hist_classified_stable_as_func_of_hull_dist(
     e_above_hull_pred=df.e_form_per_atom_pred - df.e_form_per_atom,
     e_above_hull_true=df.e_above_hull_mp,
     which_energy=which_energy,
     stability_crit=stability_crit,
     ax=axs.flat[-1],
 )
 
+text = f"Enrichment\nFactor = {metrics['enrichment']:.3}"
+ax.text(0.02, 0.25, text, fontsize=16, transform=ax.transAxes)
+
 axs.flat[-1].set(title=f"Combined ({len(df.filter(like='e_').dropna()):,})")
 axs.flat[0].legend(frameon=False, loc="upper left")
 

matbench_discovery/plot_scripts/precision_recall.py

@@ -28,8 +28,7 @@
 ).set_index("material_id")
 
 dfs["wrenformer"] = pd.read_csv(
-    f"{ROOT}/models/wrenformer/mp/"
-    "2022-09-20-wrenformer-e_form-ensemble-1-preds-e_form_per_atom.csv"
+    f"{ROOT}/models/wrenformer/mp/2022-09-20-wrenformer-e_form-ensemble-1-preds.csv"
 ).set_index("material_id")
 
 dfs["bowsr_megnet"] = pd.read_json(

matbench_discovery/plot_scripts/rolling_mae_vs_hull_dist.py

@@ -18,8 +18,7 @@
 
 data_path = (
     f"{ROOT}/data/2022-06-11-from-rhys/wren-mp-initial-structures.csv"
-    # f"{ROOT}/models/wrenformer/mp/"
-    # "2022-09-20-wrenformer-e_form-ensemble-1-preds-e_form_per_atom.csv"
+    # f"{ROOT}/models/wrenformer/mp/2022-09-20-wrenformer-e_form-ensemble-1-preds.csv"
 )
 df = pd.read_csv(data_path).set_index("material_id")
 legend_label = "Wren"

matbench_discovery/plot_scripts/rolling_mae_vs_hull_dist_wbm_batches.py

@@ -21,8 +21,7 @@
 ).set_index("material_id")
 
 df_wrenformer = pd.read_csv(
-    f"{ROOT}/models/wrenformer/mp/"
-    "2022-09-20-wrenformer-e_form-ensemble-1-preds-e_form_per_atom.csv"
+    f"{ROOT}/models/wrenformer/mp/2022-09-20-wrenformer-e_form-ensemble-1-preds.csv"
 ).set_index("material_id")
 
 

matbench_discovery/plots.py

@@ -19,7 +19,7 @@
 AxLine = Literal["x", "y", "xy", ""]
 
 
-# --- define global plot settings
+# --- start global plot settings
 quantity_labels = dict(
     n_atoms="Atom Count",
     n_elems="Element Count",
@@ -55,7 +55,8 @@
 
 
 plt.rc("font", size=14)
-plt.rc("legend", fontsize=16)
+plt.rc("legend", fontsize=16, title_fontsize=16)
+plt.rc("axes", titlesize=16, labelsize=16)
 plt.rc("savefig", bbox="tight", dpi=200)
 plt.rc("figure", dpi=200, titlesize=16)
 plt.rcParams["figure.constrained_layout.use"] = True
@@ -69,11 +70,11 @@ def hist_classified_stable_as_func_of_hull_dist(
     ax: plt.Axes = None,
     which_energy: WhichEnergy = "true",
     stability_crit: StabilityCriterion = "energy",
-    show_mae: bool = False,
-    stability_threshold: float = 0,  # set stability threshold as distance to convex
-    # hull in eV / atom, usually 0 or 0.1 eV
-    x_lim: tuple[float, float] = (-0.4, 0.4),
-) -> plt.Axes:
+    stability_threshold: float = 0,
+    show_threshold: bool = True,
+    x_lim: tuple[float | None, float | None] = (-0.4, 0.4),
+    rolling_accuracy: float = 0.02,
+) -> tuple[plt.Axes, dict[str, float]]:
     """
     Histogram of the energy difference (either according to DFT ground truth [default]
     or model predicted energy) to the convex hull for materials in the WBM data set. The
@@ -85,8 +86,33 @@ def hist_classified_stable_as_func_of_hull_dist(
 
     See fig. S1 in https://science.org/doi/10.1126/sciadv.abn4117.
 
-    NOTE this figure plots hist bars separately which causes aliasing in pdf
-    to resolve this take into Inkscape and merge regions by color
+    Args:
+        e_above_hull_pred (pd.Series): energy difference to convex hull predicted by
+            model, i.e. difference between the model's predicted and true formation
+            energy.
+        e_above_hull_true (pd.Series): energy diff to convex hull according to DFT
+            ground truth.
+        std_pred (pd.Series, optional): standard deviation of the model's predicted
+            formation energy.
+        ax (plt.Axes, optional): matplotlib axes to plot on.
+        which_energy (WhichEnergy, optional): Whether to use the true formation energy
+            or the model's predicted formation energy for the histogram.
+        stability_crit (StabilityCriterion, optional): Whether to add/subtract the
+            model's predicted uncertainty from its energy prediction when measuring
+            predicted stability.
+        stability_threshold (float, optional): set stability threshold as distance to
+            convex hull in eV/atom, usually 0 or 0.1 eV.
+        show_threshold (bool, optional): Whether to plot stability threshold as dashed
+            vertical line.
+        x_lim (tuple[float | None, float | None]): x-axis limits.
+        rolling_accuracy (float): Rolling accuracy window size in eV / atom. Set to 0 to
+            disable. Defaults to 0.01.
+
+    Returns:
+        tuple[plt.Axes, dict[str, float]]: plot axes and classification metrics
+
+    NOTE this figure plots hist bars separately which causes aliasing in pdf. Can be
+    fixed in Inkscape or similar by merging regions by color.
     """
     ax = ax or plt.gca()
 
@@ -153,26 +179,60 @@ def hist_classified_stable_as_func_of_hull_dist(
     #     e_above_hull_true
     # ), f"{n_all} != {len(e_above_hull_true)}"
 
-    # recall = n_true_pos / n_total_pos
-    # f"Prevalence = {null:.2f}\n{precision = :.2f}\n{recall = :.2f}",
-    text = f"Enrichment\nFactor = {precision/null:.3}"
-    if show_mae:
-        MAE = e_above_hull_pred.abs().mean()
-        text += f"\n{MAE = :.3}"
-
-    ax.text(
-        0.98,
-        0.98,
-        text,
-        fontsize=18,
-        verticalalignment="top",
-        horizontalalignment="right",
-        transform=ax.transAxes,
-    )
-
-    ax.set(xlabel=xlabel, ylabel="Number of compounds")
+    ax.set(xlabel=xlabel, ylabel="Number of compounds", xlim=x_lim)
+
+    if rolling_accuracy:
+        # add moving average of the accuracy (computed within 20 meV/atom intervals) as
+        # a function of ΔHd,MP is shown as a blue line (right axis)
+        ax_acc = ax.twinx()
+        ax_acc.set_ylabel("Accuracy", color="darkblue")
+        ax_acc.tick_params(labelcolor="darkblue")
+        ax_acc.set(ylim=(0, 1))
+
+        # --- moving average of the accuracy
+        # compute accuracy within 20 meV/atom intervals
+        bins = np.arange(x_lim[0], x_lim[1], rolling_accuracy)
+        bin_counts = np.histogram(e_above_hull_true, bins)[0]
+        bin_true_pos = np.histogram(true_pos, bins)[0]
+        bin_true_neg = np.histogram(true_neg, bins)[0]
+
+        # compute accuracy
+        bin_accuracies = (bin_true_pos + bin_true_neg) / bin_counts
+        # plot accuracy
+        ax_acc.plot(
+            bins[:-1],
+            bin_accuracies,
+            color="tab:blue",
+            label="Accuracy",
+            linewidth=3,
+        )
+        # ax2.fill_between(
+        #     bin_centers,
+        #     bin_accuracy - bin_accuracy_std,
+        #     bin_accuracy + bin_accuracy_std,
+        #     color="tab:blue",
+        #     alpha=0.2,
+        # )
+
+    if show_threshold:
+        ax.axvline(
+            stability_threshold,
+            color="k",
+            linestyle="--",
+            label="Stability Threshold",
+        )
 
-    return ax
+    recall = n_true_pos / n_total_pos
+
+    return ax, {
+        "enrichment": precision / null,
+        "precision": precision,
+        "recall": recall,
+        "prevalence": null,
+        "accuracy": (n_true_pos + n_true_neg)
+        / (n_true_pos + n_true_neg + n_false_pos + n_false_neg),
+        "f1": 2 * (precision * recall) / (precision + recall),
+    }
 
 
 def rolling_mae_vs_hull_dist(

tests/test_plots.py

@@ -146,7 +146,7 @@ def test_hist_classified_stable_as_func_of_hull_dist(
     else:
         std_total = None
 
-    ax = hist_classified_stable_as_func_of_hull_dist(
+    ax, metrics = hist_classified_stable_as_func_of_hull_dist(
         e_above_hull_pred=df.e_above_hull_pred,
         e_above_hull_true=df.e_above_hull_mp,
         ax=ax,
@@ -160,3 +160,6 @@ def test_hist_classified_stable_as_func_of_hull_dist(
     # assert ax.get_ylim() == pytest.approx((0, 6.3))
     assert ax.get_ylabel() == "Number of compounds"
     assert ax.get_xlabel() == r"$\Delta E_{Hull-MP}$ (eV / atom)"
+
+    assert metrics["precision"] > 0.3
+    assert metrics["recall"] > 0.3