set isort known_third_party = wandb

Author: janosh

mb_discovery/__init__.py

@@ -3,7 +3,6 @@
 import os
 from typing import Any, Generator, Sequence
 
-
 PKG_DIR = os.path.dirname(__file__)
 ROOT = os.path.dirname(PKG_DIR)
 

mb_discovery/m3gnet/eda_wbm_pre_vs_post_m3gnet_relaxation.py

@@ -11,7 +11,6 @@
 
 from mb_discovery import ROOT
 
-
 __author__ = "Janosh Riebesell"
 __date__ = "2022-06-18"
 

mb_discovery/m3gnet/join_and_plot_m3gnet_relax_results.py

@@ -18,7 +18,6 @@
     hist_classified_stable_as_func_of_hull_dist,
 )
 
-
 today = f"{datetime.now():%Y-%m-%d}"
 
 

mb_discovery/m3gnet/slurm_array_m3gnet_relax_wbm.py

@@ -14,7 +14,6 @@
 
 from mb_discovery import ROOT, as_dict_handler
 
-
 """
 To slurm submit this file, use
 

mb_discovery/plot_scripts/hist_classified_stable_as_func_of_hull_dist.py

@@ -11,7 +11,6 @@
     hist_classified_stable_as_func_of_hull_dist,
 )
 
-
 __author__ = "Rhys Goodall, Janosh Riebesell"
 __date__ = "2022-06-18"
 

mb_discovery/plot_scripts/hist_classified_stable_as_func_of_hull_dist_batches.py

@@ -11,7 +11,6 @@
     hist_classified_stable_as_func_of_hull_dist,
 )
 
-
 __author__ = "Rhys Goodall, Janosh Riebesell"
 __date__ = "2022-08-25"
 

mb_discovery/plot_scripts/plot_funcs.py

@@ -10,7 +10,6 @@
 
 from mb_discovery.plot_scripts import plt
 
-
 __author__ = "Janosh Riebesell"
 __date__ = "2022-08-05"
 
@@ -365,7 +364,10 @@ def precision_recall_vs_calc_count(
         # previous call
         return ax
 
-    ax.set(xlabel="Number of Calculations", ylabel="Precision and Recall (%)")
+    ax.set(
+        xlabel="Number of compounds sorted by model-predicted stability",
+        ylabel="Precision and Recall (%)",
+    )
 
     ax.set(ylim=(0, 100))
 

mb_discovery/plot_scripts/precision_recall_vs_calc_count.py

@@ -10,7 +10,6 @@
     precision_recall_vs_calc_count,
 )
 
-
 __author__ = "Rhys Goodall, Janosh Riebesell"
 __date__ = "2022-06-18"
 
@@ -28,9 +27,9 @@
     ).set_index("material_id")
     dfs[model_name] = df
 
-# dfs["M3GNet"] = pd.read_json(
-#     f"{ROOT}/data/2022-08-16-m3gnet-wbm-relax-results-IS2RE.json.gz"
-# ).set_index("material_id")
+dfs["M3GNet"] = pd.read_json(
+    f"{ROOT}/data/2022-08-16-m3gnet-wbm-relax-results-IS2RE.json.gz"
+).set_index("material_id")
 
 dfs["Wrenformer"] = pd.read_csv(
     f"{ROOT}/data/2022-08-16-wrenformer-preds.csv.bz2"

mb_discovery/plot_scripts/rolling_mae_vs_hull_dist.py

@@ -7,7 +7,6 @@
 from mb_discovery.plot_scripts import plt
 from mb_discovery.plot_scripts.plot_funcs import rolling_mae_vs_hull_dist
 
-
 __author__ = "Rhys Goodall, Janosh Riebesell"
 __date__ = "2022-06-18"
 

mb_discovery/plot_scripts/rolling_mae_vs_hull_dist_wbm_batches.py

@@ -7,7 +7,6 @@
 from mb_discovery.plot_scripts import plt
 from mb_discovery.plot_scripts.plot_funcs import rolling_mae_vs_hull_dist
 
-
 __author__ = "Rhys Goodall, Janosh Riebesell"
 __date__ = "2022-06-18"
 

mb_discovery/wrenformer/mp/get_mp_energies.py

@@ -7,10 +7,8 @@
 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.
+Download all MP formation and above hull energies on 2022-08-13.
 
 Related EDA of MP formation energies:
 https://github.com/janosh/pymatviz/blob/main/examples/mp_bimodal_e_form.ipynb
@@ -33,17 +31,16 @@
     "symmetry",
     "energy_above_hull",
 ]
-with MPRester() as mpr:
+with MPRester(use_document_model=False) 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 = pd.DataFrame(docs).set_index("material_id")
+df.pop("_id")
 
 df["spacegroup_number"] = df.pop("symmetry").map(lambda x: x.number)
 

mb_discovery/wrenformer/mp/use_trained_wrenformer_ensemble.py mb_discovery/wrenformer/mp/use_ensemble.py

@@ -8,7 +8,6 @@
 import wandb
 from aviary.wrenformer.deploy import deploy_wandb_checkpoints
 
-
 __author__ = "Janosh Riebesell"
 __date__ = "2022-08-15"
 
@@ -18,11 +17,11 @@
 stores predictions to CSV.
 """
 
-MODULE_DIR = os.path.dirname(__file__)
+module_dir = os.path.dirname(__file__)
+today = f"{datetime.now():%Y-%m-%d}"
 
 
 # %%
-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")
@@ -33,12 +32,13 @@
 
 wandb.login()
 wandb_api = wandb.Api()
-runs = wandb_api.runs(
-    "aviary/mp", filters={"tags": {"$in": ["wrenformer-e_form-ensemble-1"]}}
-)
+ensemble_id = "wrenformer-e_form-ensemble-1"
+runs = wandb_api.runs("aviary/mp", filters={"tags": {"$in": [ensemble_id]}})
+
+assert len(runs) == 10, f"Expected 10 runs, got {len(runs)} for {ensemble_id=}"
 
 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")
+df.round(6).to_csv(f"{module_dir}/{today}-{ensemble_id}-preds-{target_col}.csv")

tests/test_plot_funcs.py

@@ -14,7 +14,6 @@
     rolling_mae_vs_hull_dist,
 )
 
-
 DATA_DIR = f"{ROOT}/data/2022-06-11-from-rhys"
 
 df_hull = pd.read_csv(f"{DATA_DIR}/wbm-e-above-mp-hull.csv").set_index("material_id")