support model_name as array in metadata.yml files to share metadata between multiple types of a model

add chgnet + chgnet_megnet to PredFiles
update model-stats.json, metrics-table.svelte, model-run-times.svelte
show only n_best=8 models by default on model page and sort descending by F1

Author: janosh

matbench_discovery/plots.py

@@ -50,12 +50,13 @@
 )
 model_labels = dict(
     bowsr_megnet="BOWSR + MEGNet",
+    chgnet="CHGNet",
+    chgnet_megnet="CHGNet + MEGNet",
     cgcnn_p="CGCNN+P",
     cgcnn="CGCNN",
     m3gnet_megnet="M3GNet + MEGNet",
     m3gnet="M3GNet",
     megnet="MEGNet",
-    megnet_old="MEGNet Old",
     voronoi_rf="Voronoi RF",
     wrenformer="Wrenformer",
     dft="DFT",

matbench_discovery/preds.py

@@ -29,29 +29,25 @@ class PredFiles(Files):
     _root = f"{ROOT}/models/"
     _key_map = model_labels  # remap model keys below to pretty plot labels (see Files)
 
+    # bowsr optimizer coupled with original megnet
+    bowsr_megnet = "bowsr/2023-01-23-bowsr-megnet-wbm-IS2RE.csv"
+    # default CHGNet model from publication with 400,438 params
+    chgnet = "chgnet/2023-03-04-chgnet-wbm-IS2RE.csv"
+    chgnet_megnet = "chgnet/2023-03-04-chgnet-wbm-IS2RE.csv"
     # CGCnn 10-member ensemble
     cgcnn = "cgcnn/2023-01-26-test-cgcnn-wbm-IS2RE/cgcnn-ensemble-preds.csv"
-
     # cgcnn 10-member ensemble with 5-fold training set perturbations
     cgcnn_p = "cgcnn/2023-02-05-cgcnn-perturb=5.csv"
-
-    # magpie composition+voronoi tessellation structure features + sklearn random forest
-    voronoi_rf = "voronoi/2022-11-27-train-test/e-form-preds-IS2RE.csv"
-
-    # wrenformer 10-member ensemble
-    wrenformer = "wrenformer/2022-11-15-wrenformer-IS2RE-preds.csv"
-
-    # original megnet straight from publication, not re-trained
-    megnet = "megnet/2022-11-18-megnet-wbm-IS2RE/megnet-e-form-preds.csv"
-    megnet_old = "megnet/2022-11-18-megnet-wbm-IS2RE/megnet-e-form-preds.csv"
-
     # original m3gnet straight from publication, not re-trained
     m3gnet = "m3gnet/2022-10-31-m3gnet-wbm-IS2RE.csv"
-
     # m3gnet-relaxed structures fed into megnet for formation energy prediction
     m3gnet_megnet = "m3gnet/2022-10-31-m3gnet-wbm-IS2RE.csv"
-    # bowsr optimizer coupled with original megnet
-    bowsr_megnet = "bowsr/2023-01-23-bowsr-megnet-wbm-IS2RE.csv"
+    # original megnet straight from publication, not re-trained
+    megnet = "megnet/2022-11-18-megnet-wbm-IS2RE/megnet-e-form-preds.csv"
+    # magpie composition+voronoi tessellation structure features + sklearn random forest
+    voronoi_rf = "voronoi/2022-11-27-train-test/e-form-preds-IS2RE.csv"
+    # wrenformer 10-member ensemble
+    wrenformer = "wrenformer/2022-11-15-wrenformer-IS2RE-preds.csv"
 
 
 PRED_FILES = PredFiles()

models/chgnet/2023-03-04-chgnet-wbm-IS2RE.csv

@@ -21,8 +21,7 @@
 from matbench_discovery import today
 from matbench_discovery.data import DATA_FILES, as_dict_handler
 from matbench_discovery.energy import get_e_form_per_atom
-from matbench_discovery.preds import df_wbm as df_summary
-from matbench_discovery.preds import e_form_col
+from matbench_discovery.preds import df_wbm, e_form_col
 
 __author__ = "Janosh Riebesell"
 __date__ = "2023-03-01"
@@ -55,52 +54,41 @@
 
 
 # %%
-df_wbm = pd.read_json(DATA_FILES.wbm_computed_structure_entries).set_index(
+df_cse = pd.read_json(DATA_FILES.wbm_computed_structure_entries).set_index(
     "material_id"
 )
 
-df_wbm["cse"] = [
-    ComputedStructureEntry.from_dict(x) for x in tqdm(df_wbm.computed_structure_entry)
+df_cse["cse"] = [
+    ComputedStructureEntry.from_dict(x) for x in tqdm(df_cse.computed_structure_entry)
 ]
 
 
-# %% transfer chgnet energies and relaxed structures WBM CSEs
+# %% transfer CHGNet energies and relaxed structures WBM CSEs since MP2020 energy
+# corrections applied below are structure-dependent (for oxides and sulfides)
 cse: ComputedStructureEntry
 for row in tqdm(df_chgnet.itertuples(), total=len(df_chgnet)):
     mat_id, struct_dict, chgnet_energy, *_ = row
     chgnet_struct = Structure.from_dict(struct_dict)
-    cse = df_wbm.loc[mat_id, "cse"]
+    cse = df_cse.loc[mat_id, "cse"]
     cse._energy = chgnet_energy  # cse._energy is the uncorrected energy
     cse._structure = chgnet_struct
     df_chgnet.loc[mat_id, "cse"] = cse
 
 
-# %%
-df_chgnet["e_form_per_atom_chgnet_uncorrected"] = [
-    get_e_form_per_atom(cse) for cse in tqdm(df_chgnet.cse)
-]
-
-
-# %% apply energy corrections
+# %% apply energy corrections to CSEs with CHGNet
 out = MaterialsProject2020Compatibility().process_entries(
     df_chgnet.cse, verbose=True, clean=True
 )
 assert len(out) == len(df_chgnet)
 
 
 # %% compute corrected formation energies
-df_chgnet["e_form_per_atom_chgnet"] = [
-    get_e_form_per_atom(cse) for cse in tqdm(df_chgnet.cse)
-]
-
-df_chgnet[e_form_col] = df_summary[e_form_col]
+e_form_chgnet_col = "e_form_per_atom_chgnet"
+df_chgnet[e_form_chgnet_col] = [get_e_form_per_atom(cse) for cse in tqdm(df_chgnet.cse)]
 
 
 # %%
-ax = density_scatter(
-    df=df_chgnet, x="e_form_per_atom_chgnet", y="e_form_per_atom_chgnet_uncorrected"
-)
-ax = density_scatter(df=df_chgnet, x="e_form_per_atom_chgnet", y=e_form_col)
+ax = density_scatter(x=df_wbm[e_form_col], y=df_chgnet[e_form_chgnet_col])
 
 
 # %% load 2019 MEGNet formation energy model
@@ -109,7 +97,7 @@
 
 
 # %% predict formation energies on chgnet relaxed structure with MEGNet
-for material_id, cse in tqdm(df_wbm.cse.items(), total=len(df_wbm)):
+for material_id, cse in tqdm(df_cse.cse.items(), total=len(df_cse)):
     if material_id in megnet_e_form_preds:
         continue
     try:
@@ -119,17 +107,23 @@
     except Exception as exc:
         print(f"Failed to predict {material_id=}: {exc}")
 
-df_chgnet["e_form_per_atom_chgnet_megnet"] = pd.Series(megnet_e_form_preds)
+e_form_megnet_col = "e_form_per_atom_chgnet_megnet"
+# remove legacy MP corrections that MEGNet was trained on and apply newer MP2020
+# corrections instead
+df_chgnet[e_form_megnet_col] = (
+    pd.Series(megnet_e_form_preds)
+    - df_wbm.e_correction_per_atom_mp_legacy
+    + df_wbm.e_correction_per_atom_mp2020
+)
 
 assert (
     n_isna := df_chgnet.e_form_per_atom_chgnet_megnet.isna().sum()
 ) < 10, f"{n_isna=}, expected 7 or similar"
 
 
 # %%
-ax = density_scatter(
-    df=df_chgnet, x="e_form_per_atom_chgnet_megnet", y="e_form_per_atom_chgnet"
-)
+ax = density_scatter(df=df_chgnet, x=e_form_chgnet_col, y=e_form_megnet_col)
+ax = density_scatter(df=df_chgnet, x=e_form_col, y=e_form_megnet_col)
 
 
 # %%

models/chgnet/join_chgnet_results.py

@@ -1,4 +1,4 @@
-model_name: CHGNet
+model_name: [CHGNet, CHGNet + MEGNet]
 model_version: 0.0.1
 matbench_discovery_version: 1.0
 date_added: "2023-03-03"
@@ -36,5 +36,5 @@ trained_on_benchmark: false
 notes:
   description: |
     The Crystal Hamiltonian Graph Neural Network (CHGNet) is a universal GNN-based interatomic potential trained on energies, forces, stresses and magnetic moments from the MP trajectory dataset containing ∼1.5 million inorganic structures.
-    ![CHGNet Pipeline](https://user-images.githubusercontent.com/30958850/222842305-b6ed2468-8773-4e03-9de5-20c8e8de030e.svg)
-  training: Using pre-trained model released with preprint. Training set unreleased until after review.
+    ![CHGNet Pipeline](https://user-images.githubusercontent.com/30958850/222924937-1d09bbce-ee18-4b19-8061-ec689cd15887.svg)
+  training: Using pre-trained model with 400,438 params released with preprint. Training set unreleased at time of writing.

models/chgnet/metadata.yml

@@ -39,11 +39,11 @@
 slurm_vars = slurm_submit(
     job_name=job_name,
     out_dir=out_dir,
-    partition="icelake-himem",
-    account="LEE-SL3-CPU",
+    partition="ampere",
+    account="LEE-SL3-GPU",
     time="3:0:0",
-    array=f"1-{slurm_array_task_count}",
-    slurm_flags=("--mem", str(12_000)),
+    # array=f"1-{slurm_array_task_count}",
+    slurm_flags="--nodes 1 --gpus-per-node 1",
 )
 
 
@@ -104,14 +104,12 @@
     except Exception as error:
         print(f"Failed to relax {material_id}: {error}")
         continue
-    relax_dict = {
+    relax_results[material_id] = {
         "chgnet_structure": relax_result["final_structure"],
         "chgnet_trajectory": relax_result["trajectory"].__dict__,
-        e_pred_col: relax_result["energies"][-1],
+        e_pred_col: relax_result["trajectory"].energies[-1],
     }
 
-    relax_results[material_id] = relax_dict
-
 
 # %%
 df_out = pd.DataFrame(relax_results).T
@@ -123,7 +121,9 @@
 # %%
 df_wbm[e_pred_col] = df_out[e_pred_col]
 table = wandb.Table(
-    dataframe=df_wbm[["uncorrected_energy", e_pred_col, "formula"]].reset_index()
+    dataframe=df_wbm.dropna()[
+        ["uncorrected_energy", e_pred_col, "formula"]
+    ].reset_index()
 )
 
 title = f"CHGNet {task_type} ({len(df_wbm):,})"