add mb_discovery/energy/slurm_e_above_hull.py

mv mb_discovery/compute_formation_energy.py -> mb_discovery/energy/__init__.py

Author: janosh

mb_discovery/build_phase_diagram.py

@@ -13,10 +13,7 @@
 from pymatgen.ext.matproj import MPRester
 
 from mb_discovery import ROOT
-from mb_discovery.compute_formation_energy import (
-    get_elemental_ref_entries,
-    get_form_energy_per_atom,
-)
+from mb_discovery.energy import get_elemental_ref_entries, get_form_energy_per_atom
 
 today = f"{datetime.now():%Y-%m-%d}"
 module_dir = os.path.dirname(__file__)

mb_discovery/compute_formation_energy.py mb_discovery/energy/__init__.py

@@ -1,8 +1,7 @@
 import itertools
 
 import pandas as pd
-from pymatgen.analysis.phase_diagram import Entry
-from pymatgen.entries.computed_entries import ComputedEntry
+from pymatgen.analysis.phase_diagram import Entry, PDEntry
 from tqdm import tqdm
 
 from mb_discovery import ROOT
@@ -47,9 +46,7 @@ def get_elemental_ref_entries(
 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()
+        pd.read_json(mp_elem_refs_path, typ="series").map(PDEntry.from_dict).to_dict()
     )
 except FileNotFoundError:
     mp_elem_reference_entries = None

mb_discovery/plot_scripts/hist_classified_stable_as_func_of_hull_dist_batches.py

@@ -37,7 +37,7 @@
 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")

mb_discovery/plot_scripts/precision_recall_vs_calc_count.py

@@ -2,6 +2,7 @@
 from datetime import datetime
 
 import pandas as pd
+from sklearn.metrics import f1_score
 
 from mb_discovery import ROOT
 from mb_discovery.plots import StabilityCriterion, precision_recall_vs_calc_count
@@ -15,6 +16,7 @@
 # %%
 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")
+rare = "all"
 
 dfs: dict[str, pd.DataFrame] = {}
 for model_name in ("wren", "cgcnn", "voronoi"):
@@ -47,10 +49,9 @@
 # %%
 stability_crit: StabilityCriterion = "energy"
 colors = "tab:blue tab:orange teal tab:pink black red turquoise tab:purple".split()
+F1s: dict[str, float] = {}
 
-for (model_name, df), color in zip(dfs.items(), colors):
-    rare = "all"
-
+for model_name, df in dfs.items():
     # from pymatgen.core import Composition
     # rare = "no-lanthanides"
     # df["contains_rare_earths"] = df.composition.map(
@@ -91,23 +92,46 @@
         assert n_nans < 10, f"{model_name=} has {n_nans=}"
         df = df.dropna()
 
+    F1 = f1_score(df.e_above_mp_hull < 0, df.e_above_hull_pred < 0)
+    F1s[model_name] = F1
+
+
+# %%
+for (model_name, F1), color in zip(sorted(F1s.items(), key=lambda x: x[1]), colors):
+    df = dfs[model_name]
+
     ax = precision_recall_vs_calc_count(
         e_above_hull_error=df.e_above_hull_pred + df.e_above_mp_hull,
         e_above_hull_true=df.e_above_mp_hull,
         color=color,
-        label=model_name,
+        label=f"{model_name} {F1=:.2}",
         intersect_lines="recall_xy",  # or "precision_xy", None, 'all'
         stability_crit=stability_crit,
         std_pred=std_total,
     )
 
+# optimal recall line finds all stable materials without any false positives
+# can be included to confirm all models start out of with near optimal recall
+# and to see how much each model overshoots total n_stable
+n_below_hull = sum(df_hull.e_above_mp_hull < 0)
+ax.plot(
+    [0, n_below_hull],
+    [0, 100],
+    color="green",
+    linestyle="dashed",
+    linewidth=1,
+    label="Optimal Recall",
+)
+
 ax.figure.set_size_inches(10, 9)
 ax.set(xlim=(0, None))
 # 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())
+# x-ticks every 10k materials
+ax.set(xticks=range(0, int(ax.get_xlim()[1]), 10_000))
 
 
 # %%

mb_discovery/plot_scripts/rolling_mae_vs_hull_dist_wbm_batches.py

@@ -19,20 +19,26 @@
     f"{ROOT}/data/2022-06-11-from-rhys/wren-mp-initial-structures.csv"
 ).set_index("material_id")
 
+df_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_wren["e_above_mp_hull"] = df_hull.e_above_mp_hull
-assert df_wren.e_above_mp_hull.isna().sum() == 0
+df_wrenformer["e_above_mp_hull"] = df_hull.e_above_mp_hull
+assert df_wrenformer.e_above_mp_hull.isna().sum() == 0
 
-target_col = "e_form_target"
+target_col = "e_form_per_atom"
+# target_col = "e_form_target"
 
 # make sure we average the expected number of ensemble member predictions
-assert df_wren.filter(regex=r"_pred_\d").shape[1] == 10
+assert df_wrenformer.filter(regex=r"_pred_\d").shape[1] == 10
 
-df_wren["e_above_hull_pred"] = (
-    df_wren.filter(regex=r"_pred_\d").mean(axis=1) - df_wren[target_col]
+df_wrenformer["e_above_hull_pred"] = (
+    df_wrenformer.filter(regex=r"_pred_\d").mean(axis=1) - df_wrenformer[target_col]
 )
 
 
@@ -42,7 +48,7 @@
 assert len(markers) == 5  # number of WBM rounds of element substitution
 
 for idx, marker in enumerate(markers, 1):
-    df = df_wren[df_wren.index.str.startswith(f"wbm-step-{idx}")]
+    df = df_wrenformer[df_wrenformer.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) = :,}")
 
@@ -62,4 +68,4 @@
 
 
 img_path = f"{ROOT}/figures/{today}-rolling-mae-vs-hull-dist-wbm-batches-{rare=}.pdf"
-# plt.savefig(img_path)
+# fig.savefig(img_path)

mb_discovery/plots.py

@@ -320,9 +320,9 @@ def precision_recall_vs_calc_count(
     """
     ax = ax or plt.gca()
 
-    for series in (e_above_hull_error, e_above_hull_true):
-        n_nans = series.isna().sum()
-        assert n_nans == 0, f"{n_nans:,} NaNs in {series.name}"
+    # for series in (e_above_hull_error, e_above_hull_true):
+    #     n_nans = series.isna().sum()
+    #     assert n_nans == 0, f"{n_nans:,} NaNs in {series.name}"
 
     is_fresh_ax = len(ax.lines) == 0
 
@@ -412,8 +412,12 @@ def precision_recall_vs_calc_count(
     ylabel = "Precision and Recall (%)"
     ax.set(ylim=(0, 100), xlabel=xlabel, ylabel=ylabel)
 
-    [precision] = ax.plot((0, 0), (0, 0), "black", linestyle="-")
-    [recall] = ax.plot((0, 0), (0, 0), "black", linestyle=":")
+    [precision] = ax.plot(
+        (0, 0), (0, 0), "black", linestyle="-", linewidth=line_kwargs["linewidth"]
+    )
+    [recall] = ax.plot(
+        (0, 0), (0, 0), "black", linestyle=":", linewidth=line_kwargs["linewidth"]
+    )
     legend = ax.legend(
         [precision, recall],
         ("Precision", "Recall"),

models/bowsr/slurm_array_bowsr_wbm.py

@@ -20,11 +20,13 @@
 To slurm submit this file, use
 
 ```sh
-# slurm will not create logdir automatically and fail if missing
-mkdir -p models/bowsr/slurm_logs
+log_dir=models/bowsr/$(date +"%Y-%m-%d")-bowsr-megnet-wbm
+job_name=bowsr-megnet-wbm-IS2RE
+mkdir -p $log_dir # slurm fails if log_dir is missing
+
 sbatch --partition icelake-himem --account LEE-SL3-CPU --array 1-500 \
-    --time 12:0:0 --job-name bowsr-megnet-wbm-IS2RE --mem 12000 \
-    --output models/bowsr/slurm_logs/slurm-%A-%a.out \
+    --time 12:0:0 --job-name $job_name --mem 12000 \
+    --output $log_dir/slurm-%A-%a.out \
     --wrap "TF_CPP_MIN_LOG_LEVEL=2 python models/bowsr/slurm_array_bowsr_wbm.py"
 ```
 
@@ -50,7 +52,7 @@
 slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
 slurm_array_task_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))
 # set large fallback job array size for fast testing/debugging
-job_array_size = int(os.environ.get("SLURM_ARRAY_TASK_COUNT", 10_000))
+slurm_array_task_count = int(os.environ.get("SLURM_ARRAY_TASK_COUNT", 10_000))
 
 print(f"Job started running {datetime.now():%Y-%m-%d@%H-%M}")
 print(f"{slurm_job_id = }")
@@ -60,7 +62,6 @@
 
 today = f"{datetime.now():%Y-%m-%d}"
 out_dir = f"{module_dir}/{today}-bowsr-megnet-wbm-{task_type}"
-os.makedirs(out_dir, exist_ok=True)
 json_out_path = f"{out_dir}/{slurm_array_task_id}.json.gz"
 
 if os.path.isfile(json_out_path):
@@ -81,6 +82,7 @@
     maml_version=version("maml"),
     slurm_job_id=slurm_job_id,
     slurm_array_task_id=slurm_array_task_id,
+    slurm_array_task_count=slurm_array_task_count,
     data_path=data_path,
     bayes_optim_kwargs=bayes_optim_kwargs,
     optimize_kwargs=optimize_kwargs,
@@ -100,10 +102,10 @@
 
 
 # %%
-print(f"Loading from {data_path=}")
+print(f"Loading from {data_path = }")
 df_wbm = pd.read_json(data_path).set_index("material_id")
 
-df_this_job = np.array_split(df_wbm, job_array_size + 1)[slurm_array_task_id]
+df_this_job = np.array_split(df_wbm, slurm_array_task_count)[slurm_array_task_id - 1]
 
 
 # %%

models/m3gnet/join_m3gnet_relax_results.py

@@ -11,7 +11,7 @@
 from tqdm import tqdm
 
 from mb_discovery import ROOT, as_dict_handler
-from mb_discovery.compute_formation_energy import get_form_energy_per_atom
+from mb_discovery.energy import get_form_energy_per_atom
 
 __author__ = "Janosh Riebesell"
 __date__ = "2022-08-16"

models/m3gnet/slurm_array_m3gnet_relax_wbm.py

@@ -18,11 +18,13 @@
 To slurm submit this file, use
 
 ```sh
-# slurm will not create logdir automatically and fail if missing
-mkdir -p models/m3gnet/slurm_logs
+job_name=m3gnet-wbm-relax-IS2RE
+log_dir=models/m3gnet/$(date +"%Y-%m-%d")-$job_name
+mkdir -p $log_dir # slurm fails if log_dir is missing
+
 sbatch --partition icelake-himem --account LEE-SL3-CPU --array 1-100 \
-    --time 3:0:0 --job-name m3gnet-wbm-relax-IS2RE --mem 12000 \
-    --output models/m3gnet/slurm_logs/slurm-%A-%a.out \
+    --time 3:0:0 --job-name $job_name --mem 12000 \
+    --output $log_dir/slurm-%A-%a.out \
     --wrap "TF_CPP_MIN_LOG_LEVEL=2 python models/m3gnet/slurm_array_m3gnet_relax_wbm.py"
 ```
 
@@ -43,16 +45,15 @@
 slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
 slurm_array_task_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))
 # set large fallback job array size for fast testing/debugging
-job_array_size = int(os.environ.get("SLURM_ARRAY_TASK_COUNT", 10_000))
+slurm_array_task_count = int(os.environ.get("SLURM_ARRAY_TASK_COUNT", 10_000))
 
 print(f"Job started running {datetime.now():%Y-%m-%d@%H-%M}")
 print(f"{slurm_job_id = }")
 print(f"{slurm_array_task_id = }")
 print(f"{version('m3gnet') = }")
 
 today = f"{datetime.now():%Y-%m-%d}"
-out_dir = f"{ROOT}/data/{today}-m3gnet-wbm-relax-{task_type}"
-os.makedirs(out_dir, exist_ok=True)
+out_dir = f"{ROOT}/data/{today}-m3gnet-wbm-{task_type}"
 json_out_path = f"{out_dir}/{slurm_array_task_id}.json.gz"
 
 if os.path.isfile(json_out_path):
@@ -67,12 +68,13 @@
 print(f"Loading from {data_path=}")
 df_wbm = pd.read_json(data_path).set_index("material_id")
 
-df_this_job = np.array_split(df_wbm, job_array_size)[slurm_array_task_id]
+df_this_job = np.array_split(df_wbm, slurm_array_task_count)[slurm_array_task_id - 1]
 
 run_params = dict(
     m3gnet_version=version("m3gnet"),
     slurm_job_id=slurm_job_id,
     slurm_array_task_id=slurm_array_task_id,
+    slurm_array_task_count=slurm_array_task_count,
     data_path=data_path,
     task_type=task_type,
 )