make 2nd arg elemental_ref_entries to get_form_energy_per_atom() optional

now defaults to mp_elem_reference_entries

Author: janosh

mb_discovery/build_phase_diagram.py

@@ -89,22 +89,14 @@
     json.dump(elemental_ref_entries, file, default=lambda x: x.as_dict())
 
 
-# %% load MP elemental reference entries to compute formation energies
-mp_elem_refs_path = f"{ROOT}/data/2022-09-19-mp-elemental-reference-entries.json"
-mp_reference_entries = (
-    pd.read_json(mp_elem_refs_path, typ="series").map(ComputedEntry.from_dict).to_dict()
-)
-
-
 df_mp = pd.read_json(f"{ROOT}/data/2022-08-13-mp-all-energies.json.gz").set_index(
     "material_id"
 )
 
 
 # %%
 df_mp["our_mp_e_form"] = [
-    get_form_energy_per_atom(all_mp_computed_entries[mp_id], mp_reference_entries)
-    for mp_id in df_mp.index
+    get_form_energy_per_atom(all_mp_computed_entries[mp_id]) for mp_id in df_mp.index
 ]
 
 

mb_discovery/compute_formation_energy.py

@@ -1,8 +1,12 @@
 import itertools
 
+import pandas as pd
 from pymatgen.analysis.phase_diagram import Entry
+from pymatgen.entries.computed_entries import ComputedEntry
 from tqdm import tqdm
 
+from mb_discovery import ROOT
+
 
 def get_elemental_ref_entries(
     entries: list[Entry], verbose: bool = False
@@ -38,12 +42,46 @@ def get_elemental_ref_entries(
     return elemental_ref_entries
 
 
+# contains all MP elemental reference entries to compute formation energies
+# produced by get_elemental_ref_entries() in build_phase_diagram.py
+mp_elem_refs_path = f"{ROOT}/data/2022-09-19-mp-elemental-reference-entries.json"
+try:
+    mp_elem_reference_entries = (
+        pd.read_json(mp_elem_refs_path, typ="series")
+        .map(ComputedEntry.from_dict)
+        .to_dict()
+    )
+except FileNotFoundError:
+    mp_elem_reference_entries = None
+
+
 def get_form_energy_per_atom(
-    entry: Entry, elemental_ref_entries: dict[str, Entry]
+    entry: Entry, elemental_ref_entries: dict[str, Entry] = None
 ) -> float:
     """Get the formation energy of a composition from a list of entries and elemental
     reference energies.
+
+    Args:
+        entry (Entry): pymatgen Entry (PDEntry, ComputedEntry or ComputedStructureEntry)
+            to compute formation energy of.
+        elemental_ref_entries (dict[str, Entry], optional): Must be a complete set of
+            terminal (i.e. elemental) reference entries containing the lowest energy
+            phase for each element present in entry. Defaults to MP elemental reference
+            entries as collected on 2022-09-19 get_elemental_ref_entries(). This was
+            tested to give the same formation energies as computed by MP.
+
+    Returns:
+        float: formation energy in eV/atom.
     """
+    if elemental_ref_entries is None:
+        if mp_elem_reference_entries is None:
+            raise ValueError(
+                f"Couldn't load {mp_elem_refs_path=}, you must pass "
+                f"{elemental_ref_entries=} explicitly."
+            )
+
+        elemental_ref_entries = mp_elem_reference_entries
+
     comp = entry.composition
     form_energy = entry.energy - sum(
         comp[el] * elemental_ref_entries[str(el)].energy_per_atom

mb_discovery/plot_scripts/hist_classified_stable_as_func_of_hull_dist_batches.py

@@ -65,7 +65,7 @@
     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
+    df["e_form_per_atom_pred"] = df.e_form_m3gnet
 if "bowsr_megnet" in dfs:
     df = dfs["bowsr_megnet"]
     df["e_form_per_atom_pred"] = df.e_form_per_atom_bowsr
@@ -76,7 +76,8 @@
 stability_crit: StabilityCriterion = "energy"
 fig, axs = plt.subplots(2, 3, figsize=(18, 9))
 
-df = dfs[(model_name := "bowsr_megnet")]
+model_name = "m3gnet"
+df = dfs[model_name]
 
 df["e_above_mp_hull"] = df_hull.e_above_mp_hull
 df["e_form_per_atom"] = df_wbm.e_form_per_atom
@@ -109,14 +110,13 @@
 axs.flat[-1].set(title=f"Combined ({len(df.filter(like='e_').dropna()):,})")
 axs.flat[0].legend(frameon=False, loc="upper left")
 
-img_name = (
-    f"{today}-{model_name}-wbm-hull-dist-hist-{which_energy=}-{stability_crit=}.pdf"
-)
-fig.suptitle(img_name.replace("-", "/", 2).replace("-", " "), y=1.07, fontsize=16)
+img_name = f"{today}-{model_name}-wbm-hull-dist-hist-{which_energy=}-{stability_crit=}"
+suptitle = img_name.replace("-", "/", 2).replace("-", " ")
+fig.suptitle(suptitle, y=1.07, fontsize=16)
 
 
 # %%
-ax.figure.savefig(f"{ROOT}/figures/{img_name}")
+ax.figure.savefig(f"{ROOT}/figures/{img_name}.pdf")
 
 
 # %%

mb_discovery/plot_scripts/precision_recall_vs_calc_count.py

@@ -17,22 +17,21 @@
 df_hull = pd.read_csv(f"{DATA_DIR}/wbm-e-above-mp-hull.csv").set_index("material_id")
 
 dfs: dict[str, pd.DataFrame] = {}
-for model_name in ("Wren", "CGCNN", "Voronoi"):
-    df = pd.read_csv(
-        f"{DATA_DIR}/{model_name.lower()}-mp-initial-structures.csv"
-    ).set_index("material_id")
+for model_name in ("wren", "cgcnn", "voronoi"):
+    csv_path = f"{DATA_DIR}/{model_name}-mp-initial-structures.csv"
+    df = pd.read_csv(csv_path).set_index("material_id")
     dfs[model_name] = df
 
-dfs["M3GNet"] = pd.read_json(
+dfs["m3gnet"] = pd.read_json(
     f"{ROOT}/models/m3gnet/2022-08-16-m3gnet-wbm-IS2RE.json.gz"
 ).set_index("material_id")
 
-dfs["Wrenformer"] = pd.read_csv(
+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")
 
-dfs["BOWSR Megnet"] = pd.read_json(
+dfs["bowsr_megnet"] = pd.read_json(
     f"{ROOT}/models/bowsr/2022-09-22-bowsr-wbm-megnet-IS2RE.json.gz"
 ).set_index("material_id")
 
@@ -69,16 +68,16 @@
         std_total = None
 
     try:
-        if model_name == "M3GNet":
+        if model_name == "m3gnet":
             model_preds = df.e_form_m3gnet
-        elif "Wrenformer" in model_name:
+        elif "wrenformer" in model_name:
             model_preds = df.e_form_per_atom_pred_ens
         elif len(pred_cols := df.filter(like="e_form_pred").columns) >= 1:
             # Voronoi+RF has single prediction column, Wren and CGCNN each have 10
             # other cases are unexpected
             assert len(pred_cols) in (1, 10), f"{model_name=} has {len(pred_cols)=}"
             model_preds = df[pred_cols].mean(axis=1)
-        elif "BOWSR" in model_name:
+        elif "bowsr" in model_name:
             model_preds = df.e_form_per_atom_bowsr
         else:
             raise ValueError(f"Unhandled {model_name = }")
@@ -107,7 +106,9 @@
 # keep this outside loop so all model names appear in legend
 ax.legend(frameon=False, loc="lower right")
 
+img_name = f"{today}-precision-recall-vs-calc-count-{rare=}"
+ax.set(title=img_name.replace("-", "/", 2).replace("-", " ").title())
+
 
 # %%
-img_path = f"{ROOT}/figures/{today}-precision-recall-vs-calc-count-{rare=}.pdf"
-ax.figure.savefig(img_path)
+ax.figure.savefig(f"{ROOT}/figures/{img_name}.pdf")

mb_discovery/plot_scripts/rolling_mae_vs_hull_dist_wbm_batches.py

@@ -15,24 +15,24 @@
 # %%
 rare = "all"
 
-df_wbm = pd.read_csv(
+df_wren = pd.read_csv(
     f"{ROOT}/data/2022-06-11-from-rhys/wren-mp-initial-structures.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_wbm["e_above_mp_hull"] = df_hull.e_above_mp_hull
-assert df_wbm.e_above_mp_hull.isna().sum() == 0
+df_wren["e_above_mp_hull"] = df_hull.e_above_mp_hull
+assert df_wren.e_above_mp_hull.isna().sum() == 0
 
 target_col = "e_form_target"
 
 # make sure we average the expected number of ensemble member predictions
-assert df_wbm.filter(regex=r"_pred_\d").shape[1] == 10
+assert df_wren.filter(regex=r"_pred_\d").shape[1] == 10
 
-df_wbm["e_above_hull_pred"] = (
-    df_wbm.filter(regex=r"_pred_\d").mean(axis=1) - df_wbm[target_col]
+df_wren["e_above_hull_pred"] = (
+    df_wren.filter(regex=r"_pred_\d").mean(axis=1) - df_wren[target_col]
 )
 
 
@@ -42,7 +42,7 @@
 assert len(markers) == 5  # number of WBM rounds of element substitution
 
 for idx, marker in enumerate(markers, 1):
-    df = df_wbm[df_wbm.index.str.startswith(f"wbm-step-{idx}")]
+    df = df_wren[df_wren.index.str.startswith(f"wbm-step-{idx}")]
     title = f"Batch {idx} ({len(df.filter(like='e_').dropna()):,})"
     assert 1e4 < len(df) < 1e5, print(f"{len(df) = :,}")
 

mb_discovery/plots.py

@@ -20,20 +20,30 @@
 
 
 # --- define global plot settings
-px.defaults.labels = {
-    "n_atoms": "Atom Count",
-    "n_elems": "Element Count",
-    "crystal_sys": "Crystal system",
-    "spg_num": "Space group",
-    "n_wyckoff": "Number of Wyckoff positions",
-    "n_sites": "Lattice site count",
-    "energy_per_atom": "Energy (eV/atom)",
-    "e_form": "Formation energy (eV/atom)",
-    "e_above_hull": "Energy above convex hull (eV/atom)",
-    "e_above_hull_pred": "Predicted energy above convex hull (eV/atom)",
-    "e_above_mp_hull": "Energy above MP convex hull (eV/atom)",
-    "e_above_hull_error": "Error in energy above convex hull (eV/atom)",
-}
+quantity_labels = dict(
+    n_atoms="Atom Count",
+    n_elems="Element Count",
+    crystal_sys="Crystal system",
+    spg_num="Space group",
+    n_wyckoff="Number of Wyckoff positions",
+    n_sites="Lattice site count",
+    energy_per_atom="Energy (eV/atom)",
+    e_form="Formation energy (eV/atom)",
+    e_above_hull="Energy above convex hull (eV/atom)",
+    e_above_hull_pred="Predicted energy above convex hull (eV/atom)",
+    e_above_mp_hull="Energy above MP convex hull (eV/atom)",
+    e_above_hull_error="Error in energy above convex hull (eV/atom)",
+)
+model_labels = dict(
+    wren="Wren",
+    wrenformer="Wrenformer",
+    m3gnet="M3GNet",
+    bowsr_megnet="BOWSR + MEGNet",
+    cgcnn="CGCNN",
+    voronoi="Voronoi",
+    wbm="WBM",
+)
+px.defaults.labels = quantity_labels | model_labels
 
 pio.templates.default = "plotly_white"
 

models/bowsr/slurm_array_bowsr_wbm.py

@@ -84,6 +84,7 @@
     data_path=data_path,
     bayes_optim_kwargs=bayes_optim_kwargs,
     optimize_kwargs=optimize_kwargs,
+    task_type=task_type,
 )
 if wandb.run is None:
     wandb.login()

models/m3gnet/join_m3gnet_relax_results.py

@@ -1,20 +1,17 @@
 # %%
 from __future__ import annotations
 
-import gzip
-import io
 import os
-import pickle
-import urllib.request
 from datetime import datetime
 from glob import glob
 
 import pandas as pd
-from pymatgen.analysis.phase_diagram import PatchedPhaseDiagram, PDEntry
+from pymatgen.analysis.phase_diagram import PDEntry
 from pymatgen.core import Structure
 from tqdm import tqdm
 
 from mb_discovery import ROOT, as_dict_handler
+from mb_discovery.compute_formation_energy import get_form_energy_per_atom
 
 __author__ = "Janosh Riebesell"
 __date__ = "2022-08-16"
@@ -67,21 +64,12 @@
 
 
 # %%
-# 2022-01-25-ppd-mp+wbm.pkl.gz (235 MB)
-ppd_pickle_url = "https://figshare.com/files/36669624"
-zipped_file = urllib.request.urlopen(ppd_pickle_url)
-
-ppd_mp_wbm: PatchedPhaseDiagram = pickle.load(
-    io.BytesIO(gzip.decompress(zipped_file.read()))
-)
-
-
 pd_entries_m3gnet = [
     PDEntry(row.m3gnet_structure.composition, row.m3gnet_energy)
     for row in df_m3gnet.itertuples()
 ]
 df_m3gnet["e_form_m3gnet_from_ppd"] = [
-    ppd_mp_wbm.get_form_energy_per_atom(x) for x in pd_entries_m3gnet
+    get_form_energy_per_atom(entry) for entry in pd_entries_m3gnet
 ]
 
 
@@ -93,22 +81,6 @@
 df_m3gnet["e_above_mp_hull"] = df_hull.e_above_mp_hull
 
 
-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_m3gnet["e_form_wbm"] = df_wbm.e_form_per_atom
-df_m3gnet["wbm_energy"] = df_wbm.energy
-
-pd_entries_wbm = [
-    PDEntry(row.m3gnet_structure.composition, row.wbm_energy)
-    for row in df_m3gnet.itertuples()
-]
-df_m3gnet["e_form_ppd_2022_01_25"] = [
-    ppd_mp_wbm.get_form_energy_per_atom(x) for x in pd_entries_wbm
-]
-
-
 # %%
 df_m3gnet.hist(bins=200, figsize=(18, 12))
 df_m3gnet.isna().sum()

models/m3gnet/slurm_array_m3gnet_relax_wbm.py

@@ -74,6 +74,7 @@
     slurm_job_id=slurm_job_id,
     slurm_array_task_id=slurm_array_task_id,
     data_path=data_path,
+    task_type=task_type,
 )
 if wandb.run is None:
     wandb.login()