add mb_discovery/wrenformer/{mp,m3gnet_train_set}

Author: janosh

.gitignore

@@ -19,3 +19,4 @@ job-logs/
 
 # slurm logs
 slurm-*out
+mb_discovery/**/*.csv

mb_discovery/plot_scripts/precision_recall_vs_calc_count.py

@@ -33,7 +33,7 @@
 # ).set_index("material_id")
 
 dfs["Wrenformer"] = pd.read_csv(
-    f"{ROOT}/data/2022-08-16-wrenformer-ensemble-predictions.csv.bz2"
+    f"{ROOT}/data/2022-08-16-wrenformer-preds.csv.bz2"
 ).set_index("material_id")
 
 

mb_discovery/plot_scripts/rolling_mae_vs_hull_dist.py

@@ -19,7 +19,7 @@
 
 data_path = (
     f"{ROOT}/data/2022-06-11-from-rhys/wren-mp-initial-structures.csv"
-    # f"{ROOT}/data/2022-08-16-wrenformer-ensemble-predictions.csv.bz2"
+    # f"{ROOT}/data/2022-08-16-wrenformer-preds.csv.bz2"
 )
 df = pd.read_csv(data_path).set_index("material_id")
 legend_label = "Wren"

mb_discovery/wrenformer/mp/get_mp_energies.py

@@ -0,0 +1,56 @@
+# %%
+from datetime import datetime
+
+import pandas as pd
+from aviary import ROOT
+from aviary.utils import as_dict_handler
+from aviary.wren.utils import get_aflow_label_from_spglib
+from mp_api.client import MPRester
+
+
+"""
+Download all MP formation and above hull energies on 2022-08-13 for training a
+Wrenformer ensemble.
+
+Related EDA of MP formation energies:
+https://github.com/janosh/pymatviz/blob/main/examples/mp_bimodal_e_form.ipynb
+"""
+
+__author__ = "Janosh Riebesell"
+__date__ = "2022-08-13"
+
+today = f"{datetime.now():%Y-%m-%d}"
+
+
+# %% query all MP formation energies on 2022-08-13
+fields = [
+    "material_id",
+    "task_ids",
+    "formula_pretty",
+    "formation_energy_per_atom",
+    "energy_per_atom",
+    "structure",
+    "symmetry",
+    "energy_above_hull",
+]
+with MPRester() as mpr:
+    docs = mpr.summary.search(fields=fields)
+
+print(f"{today}: {len(docs) = :,}")
+# 2022-08-13: len(docs) = 146,323
+
+
+# %%
+df = pd.DataFrame(
+    [{key: getattr(doc, key, None) for key in fields} for doc in docs]
+).set_index("material_id")
+
+df["spacegroup_number"] = df.pop("symmetry").map(lambda x: x.number)
+
+df["wyckoff"] = df.structure.map(get_aflow_label_from_spglib)
+
+df.to_json(
+    f"{ROOT}/datasets/{today}-mp-all-energies.json.gz", default_handler=as_dict_handler
+)
+
+# df = pd.read_json(f"{ROOT}/datasets/2022-08-13-mp-all-energies.json.gz")

mb_discovery/wrenformer/mp/use_trained_wrenformer_ensemble.py

@@ -0,0 +1,44 @@
+# %%
+from __future__ import annotations
+
+import os
+from datetime import datetime
+
+import pandas as pd
+import wandb
+from aviary.wrenformer.deploy import deploy_wandb_checkpoints
+
+
+__author__ = "Janosh Riebesell"
+__date__ = "2022-08-15"
+
+"""
+Script that downloads checkpoints for an ensemble of Wrenformer models trained on the MP
+formation energies, then makes predictions on some dataset, prints ensemble metrics and
+stores predictions to CSV.
+"""
+
+MODULE_DIR = os.path.dirname(__file__)
+
+
+# %%
+today = f"{datetime.now():%Y-%m-%d}"
+# download wbm-steps-summary.csv (23.31 MB)
+data_path = "https://figshare.com/files/36714216?private_link=ff0ad14505f9624f0c05"
+df = pd.read_csv(data_path).set_index("material_id")
+
+
+target_col = "e_form_per_atom"
+df[target_col] = df.e_form / df.n_sites
+
+wandb.login()
+wandb_api = wandb.Api()
+runs = wandb_api.runs(
+    "aviary/mp", filters={"tags": {"$in": ["wrenformer-e_form-ensemble-1"]}}
+)
+
+df, ensemble_metrics = deploy_wandb_checkpoints(
+    runs, df, input_col="wyckoff", target_col=target_col
+)
+
+df.to_csv(f"{MODULE_DIR}/{today}-wrenformer-preds-{target_col}.csv")