fix misnamed variable in rolling_mae_vs_hull_dist() e_above_hull_pred->e_above_hull_error

clean up scripts/rolling_mae_vs_hull_dist{,batches}.py

Author: janosh

matbench_discovery/plots.py

@@ -250,8 +250,8 @@ def hist_classified_stable_vs_hull_dist(
 
 def rolling_mae_vs_hull_dist(
     e_above_hull_true: pd.Series,
-    e_above_hull_pred: pd.Series,
-    half_window: float = 0.02,
+    e_above_hull_error: pd.Series,
+    window: float = 0.04,
     bin_width: float = 0.002,
     x_lim: tuple[float, float] = (-0.2, 0.3),
     ax: plt.Axes = None,
@@ -273,19 +273,25 @@ def rolling_mae_vs_hull_dist(
     rolling_maes = np.zeros_like(bins)
     rolling_stds = np.zeros_like(bins)
     for idx, bin_center in enumerate(bins):
-        low = bin_center - half_window
-        high = bin_center + half_window
+        low = bin_center - window
+        high = bin_center + window
 
         mask = (e_above_hull_true <= high) & (e_above_hull_true > low)
-        rolling_maes[idx] = e_above_hull_pred.loc[mask].abs().mean()
-        rolling_stds[idx] = scipy.stats.sem(e_above_hull_pred.loc[mask].abs())
+        rolling_maes[idx] = e_above_hull_error.loc[mask].abs().mean()
+        rolling_stds[idx] = scipy.stats.sem(e_above_hull_error.loc[mask].abs())
 
     kwargs = dict(linewidth=3) | kwargs
     ax.plot(bins, rolling_maes, **kwargs)
 
     ax.fill_between(
         bins, rolling_maes + rolling_stds, rolling_maes - rolling_stds, alpha=0.3
     )
+    # alternative implementation using pandas.rolling(). drawback: window size can only
+    # be set as number of observations, not fixed-size energy above hull interval.
+    # e_above_hull_error.index = e_above_hull_true  # warning: in-place change
+    # e_above_hull_error.sort_index().abs().rolling(window=8000).mean().plot(
+    #     ax=ax, **kwargs
+    # )
 
     if not is_fresh_ax:
         # return earlier if all plot objects besides the line were already drawn by a
@@ -294,7 +300,7 @@ def rolling_mae_vs_hull_dist(
 
     scale_bar = AnchoredSizeBar(
         ax.transData,
-        2 * half_window,
+        window,
         "40 meV",
         "lower left",
         pad=0.5,

scripts/rolling_mae_vs_hull_dist.py

@@ -1,46 +1,25 @@
 # %%
-import pandas as pd
-
 from matbench_discovery import ROOT, today
-from matbench_discovery.load_preds import df_wbm
+from matbench_discovery.load_preds import load_df_wbm_with_preds
 from matbench_discovery.plots import rolling_mae_vs_hull_dist
 
 __author__ = "Rhys Goodall, Janosh Riebesell"
 __date__ = "2022-06-18"
 
 
 # %%
-data_path = (
-    f"{ROOT}/data/2022-06-11-from-rhys/wren-mp-initial-structures.csv"
-    # f"{ROOT}/models/wrenformer/2022-11-15-wrenformer-IS2RE-preds.csv"
-)
-df = pd.read_csv(data_path).set_index("material_id")
-legend_label = "Wren"
-
-
-# %%
-df["e_above_hull_mp"] = df_wbm.e_above_hull_mp2020_corrected_ppd_mp
-
-assert all(n_nans := df.isna().sum() == 0), f"Found {n_nans} NaNs"
-
-target_col = "e_form_target"
-# target_col = "e_form_per_atom"
-# --- or ---
-# target_col = "e_form_per_atom_target"
-# df["e_form_per_atom_target"] = df.e_form / df.n_sites
-
-# make sure we average the expected number of ensemble member predictions
-assert df.filter(regex=r"_pred_\d").shape[1] == 10
+df_wbm = load_df_wbm_with_preds(models=["Wren", "Wrenformer"]).round(3)
 
-df["e_form_pres_ens"] = df.filter(regex=r"_pred_\d+").mean(axis=1)
-df["e_above_hull_pred"] = df.e_form_pres_ens - df[target_col]
+e_above_hull_col = "e_above_hull_mp2020_corrected_ppd_mp"
+target_col = "e_form_per_atom_mp2020_corrected"
 
 
 # %%
+model_name = "Wrenformer"
 ax = rolling_mae_vs_hull_dist(
-    e_above_hull_pred=df.e_above_hull_pred,
-    e_above_hull_true=df.e_above_hull_mp,
-    label=legend_label,
+    e_above_hull_true=df_wbm[e_above_hull_col],
+    e_above_hull_error=df_wbm[target_col] - df_wbm[model_name],
+    label=model_name,
 )
 
 fig = ax.figure

scripts/rolling_mae_vs_hull_dist_wbm_batches.py

@@ -1,57 +1,35 @@
 # %%
-import pandas as pd
-
 from matbench_discovery import ROOT, today
-from matbench_discovery.load_preds import df_wbm
+from matbench_discovery.load_preds import load_df_wbm_with_preds
 from matbench_discovery.plots import plt, rolling_mae_vs_hull_dist
 
 __author__ = "Rhys Goodall, Janosh Riebesell"
 __date__ = "2022-06-18"
 
 
 # %%
-df_wren = pd.read_csv(
-    f"{ROOT}/data/2022-06-11-from-rhys/wren-mp-initial-structures.csv"
-).set_index("material_id")
-
-df_wrenformer = pd.read_csv(
-    f"{ROOT}/models/wrenformer/2022-11-15-wrenformer-IS2RE-preds.csv"
-).set_index("material_id")
-
-
-# %%
-model_name = "wren"
-df = {"wren": df_wren, "wrenformer": df_wrenformer}[model_name]
-
-df["e_above_hull_mp"] = df_wbm.e_above_hull_mp2020_corrected_ppd_mp
-assert df.e_above_hull_mp.isna().sum() == 0
-
-possible_targets = (
-    "e_form_per_atom_mp2020_corrected e_form_per_atom e_form_target".split()
-)
-target_col = next(filter(lambda x: x in df, possible_targets))
-
-# make sure we average the expected number of ensemble member predictions
-assert df.filter(regex=r"_pred_\d").shape[1] == 10
+df_wbm = load_df_wbm_with_preds(models=["Wren", "Wrenformer"]).round(3)
 
-df["e_above_hull_pred"] = df.filter(regex=r"_pred_\d").mean(axis=1) - df[target_col]
+e_above_hull_col = "e_above_hull_mp2020_corrected_ppd_mp"
+target_col = "e_form_per_atom_mp2020_corrected"
 
 
 # %%
+model_name = "Wrenformer"
 fig, ax = plt.subplots(1, figsize=(10, 9))
 markers = ("o", "v", "^", "H", "D")
 assert len(markers) == 5  # number of WBM rounds of element substitution
 
 for idx, marker in enumerate(markers, 1):
     # select all rows from WBM step=idx
-    df_step = df[df.index.str.startswith(f"wbm-step-{idx}")]
+    df_step = df_wbm[df_wbm.index.str.startswith(f"wbm-step-{idx}")]
 
     title = f"Batch {idx} ({len(df_step.filter(like='e_').dropna()):,})"
     assert 1e4 < len(df_step) < 1e5, print(f"{len(df_step) = :,}")
 
     rolling_mae_vs_hull_dist(
-        e_above_hull_pred=df_step.e_above_hull_pred,
-        e_above_hull_true=df_step.e_above_hull_mp,
+        e_above_hull_error=df_step[target_col] - df_step[model_name],
+        e_above_hull_true=df_step[e_above_hull_col],
         ax=ax,
         label=title,
         marker=marker,
@@ -65,5 +43,5 @@
 ax.set(title=f"{today} model={model_name}")
 
 
-img_name = f"{today}-{model_name}-rolling-mae-vs-hull-dist-wbm-batches"
-# fig.savefig(f"{ROOT}/figures/{img_name}.pdf")
+img_path = f"{ROOT}/figures/{today}-{model_name}-rolling-mae-vs-hull-dist-wbm-batches"
+# fig.savefig(f"{img_path}.pdf")

tests/test_plots.py

@@ -71,25 +71,25 @@ def test_cumulative_precision_recall(
         assert fig.layout.yaxis2.title.text == "Recall"
 
 
-@pytest.mark.parametrize("half_window", (0.02, 0.002))
+@pytest.mark.parametrize("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]
+    window: float, bin_width: float, x_lim: tuple[float, float]
 ) -> None:
     ax = plt.figure().gca()  # new figure ensures test functions use different axes
 
     for (model_name, df), color in zip(
         test_dfs.items(), ("tab:blue", "tab:orange", "tab:pink")
     ):
         ax = rolling_mae_vs_hull_dist(
-            e_above_hull_pred=df.e_above_hull_pred,
+            e_above_hull_error=df.e_above_hull_pred,
             e_above_hull_true=df.e_above_hull_mp,
             color=color,
             label=model_name,
             ax=ax,
             x_lim=x_lim,
-            half_window=half_window,
+            window=window,
             bin_width=bin_width,
         )