make model switching easy in hist_classified_stable_as_func_of_hull_dist_batches.py

Author: janosh

mb_discovery/plot_scripts/hist_classified_stable_as_func_of_hull_dist.py

@@ -40,7 +40,7 @@
 df["e_above_mp_hull"] = df_hull.e_above_mp_hull
 
 # download wbm-steps-summary.csv (23.31 MB)
-df_summary = pd.read_csv(
+df_wbm = pd.read_csv(
     "https://figshare.com/files/37570234?private_link=ff0ad14505f9624f0c05"
 ).set_index("material_id")
 

mb_discovery/plot_scripts/hist_classified_stable_as_func_of_hull_dist_batches.py

@@ -2,6 +2,7 @@
 from datetime import datetime
 
 import pandas as pd
+import pymatviz
 
 from mb_discovery import ROOT
 from mb_discovery.plots import (
@@ -29,48 +30,73 @@
 
 
 # %%
-df = pd.read_csv(
+dfs = {}
+dfs["wren"] = pd.read_csv(
     f"{ROOT}/data/2022-06-11-from-rhys/wren-mp-initial-structures.csv"
 ).set_index("material_id")
+dfs["m3gnet"] = pd.read_json(
+    f"{ROOT}/models/m3gnet/2022-08-16-m3gnet-wbm-relax-results-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"
+).set_index("material_id")
+
 
 df_hull = pd.read_csv(
     f"{ROOT}/data/2022-06-11-from-rhys/wbm-e-above-mp-hull.csv"
 ).set_index("material_id")
 
-df["e_above_mp_hull"] = df_hull.e_above_mp_hull
-
 # download wbm-steps-summary.csv (23.31 MB)
-df_summary = pd.read_csv(
+df_wbm = pd.read_csv(
     "https://figshare.com/files/37570234?private_link=ff0ad14505f9624f0c05"
 ).set_index("material_id")
 
 
+dfs["m3gnet"] = dfs.pop("M3Gnet")
+
+
+# %%
+if "wren" in dfs:
+    df = dfs["wren"]
+    pred_cols = df.filter(regex=r"_pred_\d").columns
+    # make sure we average the expected number of ensemble member predictions
+    assert len(pred_cols) == 10
+    df["e_form_per_atom_pred"] = df[pred_cols].mean(axis=1)
+if "m3gnet" in dfs:
+    df = dfs["m3gnet"]
+    df["e_form_per_atom_pred"] = df.e_form_ppd_2022_01_25
+
+
 # %%
 which_energy: WhichEnergy = "true"
 stability_crit: StabilityCriterion = "energy"
-df["wbm_batch"] = df.index.str.split("-").str[2]
 fig, axs = plt.subplots(2, 3, figsize=(18, 9))
 
-# make sure we average the expected number of ensemble member predictions
-pred_cols = df.filter(regex=r"_pred_\d").columns
-assert len(pred_cols) == 10
+df = dfs[(model_name := "wren")]
 
+df["e_above_mp_hull"] = df_hull.e_above_mp_hull
+df["e_form_per_atom"] = df_wbm.e_form_per_atom
+
+
+for batch_idx, ax in zip(range(1, 6), axs.flat):
+    batch_df = df[df.index.str.startswith(f"wbm-step-{batch_idx}-")]
+    assert 1e4 < len(batch_df) < 1e5, print(f"{len(batch_df) = :,}")
 
-for (batch_idx, batch_df), ax in zip(df.groupby("wbm_batch"), axs.flat):
     hist_classified_stable_as_func_of_hull_dist(
-        e_above_hull_pred=batch_df[pred_cols].mean(axis=1) - batch_df.e_form_target,
+        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_mp_hull,
         which_energy=which_energy,
         stability_crit=stability_crit,
         ax=ax,
     )
 
-    title = f"Batch {batch_idx} ({len(df):,})"
+    title = f"Batch {batch_idx} ({len(batch_df):,})"
     ax.set(title=title)
 
 
 hist_classified_stable_as_func_of_hull_dist(
-    e_above_hull_pred=df[pred_cols].mean(axis=1),
+    e_above_hull_pred=df.e_form_per_atom_pred - df.e_form_per_atom,
     e_above_hull_true=df.e_above_mp_hull,
     which_energy=which_energy,
     stability_crit=stability_crit,
@@ -80,5 +106,17 @@
 axs.flat[-1].set(title=f"Combined {batch_idx} ({len(df):,})")
 axs.flat[0].legend(frameon=False, loc="upper left")
 
-img_name = f"{today}-wren-wbm-hull-dist-hist-{which_energy=}-{stability_crit=}.pdf"
+img_name = (
+    f"{today}-{model_name}-wbm-hull-dist-hist-{which_energy=}-{stability_crit=}.pdf"
+)
 # plt.savefig(f"{ROOT}/figures/{img_name}")
+
+
+# %%
+pymatviz.density_scatter(
+    dfs["wren"].dropna().e_form_per_atom_pred, dfs["wren"].dropna().e_form_per_atom
+)
+
+pymatviz.density_scatter(
+    dfs["m3gnet"].dropna().e_form_per_atom_pred, dfs["m3gnet"].dropna().e_form_per_atom
+)

mb_discovery/plot_scripts/rolling_mae_vs_hull_dist_wbm_batches.py

@@ -44,6 +44,7 @@
 for idx, marker in enumerate(markers, 1):
     title = f"Batch {idx}"
     df = df_wbm[df_wbm.index.str.startswith(f"wbm-step-{idx}")]
+    assert 1e4 < len(df) < 1e5, print(f"{len(df) = :,}")
 
     rolling_mae_vs_hull_dist(
         e_above_hull_pred=df.e_above_hull_pred,

mb_discovery/plots.py

@@ -129,16 +129,15 @@ def hist_classified_stable_as_func_of_hull_dist(
         stacked=True,
     )
 
-    n_true_pos, n_false_pos, n_true_neg, n_false_neg = (
-        len(true_pos),
-        len(false_pos),
-        len(true_neg),
-        len(false_neg),
+    n_true_pos, n_false_pos, n_true_neg, n_false_neg = map(
+        len, (true_pos, false_pos, true_neg, false_neg)
     )
     # null = (tp + fn) / (tp + tn + fp + fn)
     precision = n_true_pos / (n_true_pos + n_false_pos)
 
-    assert n_true_pos + n_false_pos + n_true_neg + n_false_neg == len(e_above_hull_true)
+    # assert (n_all := n_true_pos + n_false_pos + n_true_neg + n_false_neg) == len(
+    #     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}",

models/bowsr/join_bowsr_results.py

@@ -6,11 +6,11 @@
 from glob import glob
 
 import pandas as pd
+import pymatviz
 from pymatgen.core import Structure
 from tqdm import tqdm
 
 from mb_discovery import ROOT, as_dict_handler
-from mb_discovery.plots import hist_classified_stable_as_func_of_hull_dist
 
 __author__ = "Janosh Riebesell"
 __date__ = "2022-09-22"
@@ -22,7 +22,7 @@
 module_dir = os.path.dirname(__file__)
 task_type = "IS2RE"
 date = "2022-09-22"
-glob_pattern = f"{date}-bowsr-wbm-{task_type}/*.json.gz"
+glob_pattern = f"{date}-bowsr-megnet-wbm-{task_type}/*.json.gz"
 file_paths = sorted(glob(f"{module_dir}/{glob_pattern}"))
 print(f"Found {len(file_paths):,} files for {glob_pattern = }")
 
@@ -37,12 +37,7 @@
         continue
     # keep whole dataframe in memory
     df = pd.read_json(file_path).set_index("material_id")
-    col_map = dict(
-        structure_pred="structure_bowsr",
-        energy_pred="energy_bowsr",
-        e_form_per_atom_pred="e_form_per_atom_bowsr",
-    )
-    df = df.rename(columns=col_map)
+
     df["structure_bowsr"] = df.structure_bowsr.map(Structure.from_dict)
     df["formula"] = df.structure_bowsr.map(lambda x: x.formula)
     df["volume"] = df.structure_bowsr.map(lambda x: x.volume)
@@ -54,39 +49,31 @@
 df_bowsr = pd.concat(dfs.values())
 
 
-# %%
+# %% compare against WBM formation energy targets to make sure we got sensible results
 df_wbm = pd.read_csv(  # download wbm-steps-summary.csv (23.31 MB)
     "https://figshare.com/files/37570234?private_link=ff0ad14505f9624f0c05"
 ).set_index("material_id")
 
 df_bowsr["e_form_wbm"] = df_wbm.e_form_per_atom
 
+print(f"{len(df_bowsr) - len(df_wbm) = :,} = {len(df_bowsr):,} - {len(df_wbm):,}")
+
 
 # %%
 df_bowsr.hist(bins=200, figsize=(18, 12))
 df_bowsr.isna().sum()
 
 
 # %%
-out_path = f"{ROOT}/models/bowsr/{today}-bowsr-wbm-{task_type}.json.gz"
-df_bowsr.reset_index().to_json(out_path, default_handler=as_dict_handler)
-
-out_path = f"{ROOT}/models/bowsr/2022-08-16-bowsr-wbm-IS2RE.json.gz"
-df_bowsr = pd.read_json(out_path).set_index("material_id")
+pymatviz.density_scatter(
+    df_bowsr.dropna().e_form_per_atom_bowsr,
+    df_bowsr.dropna().e_form_wbm,
+)
 
 
 # %%
-df_hull = pd.read_csv(
-    f"{ROOT}/data/2022-06-11-from-rhys/wbm-e-above-mp-hull.csv"
-).set_index("material_id")
-df_bowsr["e_above_mp_hull"] = df_hull.e_above_mp_hull
-df_bowsr["e_above_hull_pred"] = (  # TODO fix this incorrect e_above_hull_pred
-    df_bowsr["e_form_per_atom_bowsr"] - df_bowsr["e_above_mp_hull"]
-)
-
-ax_hull_dist_hist = hist_classified_stable_as_func_of_hull_dist(
-    e_above_hull_pred=df_bowsr.e_above_hull_pred,
-    e_above_hull_true=df_bowsr.e_above_mp_hull,
-)
+out_path = f"{ROOT}/models/bowsr/{today}-bowsr-megnet-wbm-{task_type}.json.gz"
+df_bowsr.reset_index().to_json(out_path, default_handler=as_dict_handler)
 
-# ax_hull_dist_hist.figure.savefig(f"{ROOT}/plots/{today}-bowsr-wbm-hull-dist-hist.pdf")
+# out_path = f"{ROOT}/models/bowsr/2022-08-16-bowsr-megnet-wbm-IS2RE.json.gz"
+# df_bowsr = pd.read_json(out_path).set_index("material_id")

models/m3gnet/join_m3gnet_relax_results.py

@@ -15,7 +15,6 @@
 from tqdm import tqdm
 
 from mb_discovery import ROOT, as_dict_handler
-from mb_discovery.plots import hist_classified_stable_as_func_of_hull_dist
 
 __author__ = "Janosh Riebesell"
 __date__ = "2022-08-16"
@@ -86,7 +85,7 @@
 ]
 
 
-# %%
+# %% compare against WBM formation energy targets to make sure we got sensible results
 df_hull = pd.read_csv(
     f"{ROOT}/data/2022-06-11-from-rhys/wbm-e-above-mp-hull.csv"
 ).set_index("material_id")
@@ -119,18 +118,5 @@
 out_path = f"{ROOT}/models/m3gnet/{today}-m3gnet-wbm-relax-{task_type}.json.gz"
 df_m3gnet.reset_index().to_json(out_path, default_handler=as_dict_handler)
 
-out_path = f"{ROOT}/models/m3gnet/2022-08-16-m3gnet-wbm-relax-results-IS2RE.json.gz"
-df_m3gnet = pd.read_json(out_path).set_index("material_id")
-
-
-# %%
-df_m3gnet["e_above_hull_pred"] = (  # TODO fix this incorrect e_above_hull_pred
-    df_m3gnet["e_form_m3gnet_from_ppd"] - df_m3gnet["e_above_mp_hull"]
-)
-
-ax_hull_dist_hist = hist_classified_stable_as_func_of_hull_dist(
-    e_above_hull_pred=df_m3gnet.e_above_hull_pred,
-    e_above_hull_true=df_m3gnet.e_above_mp_hull,
-)
-
-# ax_hull_dist_hist.figure.savefig(f"{ROOT}/plots/{today}-m3gnet-wbm-hull-dist-hist.pdf")
+# out_path = f"{ROOT}/models/m3gnet/2022-08-16-m3gnet-wbm-relax-results-IS2RE.json.gz"
+# df_m3gnet = pd.read_json(out_path).set_index("material_id")