rename job_id->slurm_job_id in wandb run.config, rename 2022-08-16-m3gnet-wbm-relax-results-IS2RE.json.gz to 2022-08-16-m3gnet-wbm-IS2RE.json.gz

Author: janosh

.pre-commit-config.yaml

@@ -12,18 +12,18 @@ repos:
       - id: isort
 
   - repo: https://github.com/psf/black
-    rev: 22.6.0
+    rev: 22.8.0
     hooks:
       - id: black
 
   - repo: https://github.com/pycqa/flake8
-    rev: 4.0.1
+    rev: 5.0.4
     hooks:
       - id: flake8
         additional_dependencies: [flake8-bugbear]
 
   - repo: https://github.com/asottile/pyupgrade
-    rev: v2.34.0
+    rev: v2.38.2
     hooks:
       - id: pyupgrade
         args: [--py39-plus]
@@ -45,19 +45,19 @@ repos:
       - id: trailing-whitespace
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v0.961
+    rev: v0.981
     hooks:
       - id: mypy
         additional_dependencies: [types-pyyaml]
 
   - repo: https://github.com/codespell-project/codespell
-    rev: v2.1.0
+    rev: v2.2.1
     hooks:
       - id: codespell
         stages: [commit, commit-msg]
         exclude_types: [csv, html, json]
 
   - repo: https://github.com/PyCQA/autoflake
-    rev: v1.4
+    rev: v1.6.1
     hooks:
       - id: autoflake

mb_discovery/__init__.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import os
-from typing import Any, Generator, Sequence
+from collections.abc import Generator, Sequence
+from typing import Any
 
 PKG_DIR = os.path.dirname(__file__)
 ROOT = os.path.dirname(PKG_DIR)

mb_discovery/plot_scripts/hist_classified_stable_as_func_of_hull_dist_batches.py

@@ -35,12 +35,15 @@
     f"{ROOT}/data/2022-06-11-from-rhys/wren-mp-initial-structures.csv"
 ).set_index("material_id")
 dfs["m3gnet"] = pd.read_json(
-    f"{ROOT}/models/m3gnet/2022-08-16-m3gnet-wbm-relax-results-IS2RE.json.gz"
+    f"{ROOT}/models/m3gnet/2022-08-16-m3gnet-wbm-IS2RE.json.gz"
 ).set_index("material_id")
 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(
+    f"{ROOT}/models/bowsr/2022-09-22-bowsr-wbm-megnet-IS2RE.json.gz"
+).set_index("material_id")
 
 
 df_hull = pd.read_csv(
@@ -53,9 +56,6 @@
 ).set_index("material_id")
 
 
-dfs["m3gnet"] = dfs.pop("M3Gnet")
-
-
 # %%
 if "wren" in dfs:
     df = dfs["wren"]
@@ -66,14 +66,17 @@
 if "m3gnet" in dfs:
     df = dfs["m3gnet"]
     df["e_form_per_atom_pred"] = df.e_form_ppd_2022_01_25
+if "bowsr_megnet" in dfs:
+    df = dfs["bowsr_megnet"]
+    df["e_form_per_atom_pred"] = df.e_form_per_atom_bowsr
 
 
 # %%
 which_energy: WhichEnergy = "true"
 stability_crit: StabilityCriterion = "energy"
 fig, axs = plt.subplots(2, 3, figsize=(18, 9))
 
-df = dfs[(model_name := "wren")]
+df = dfs[(model_name := "bowsr_megnet")]
 
 df["e_above_mp_hull"] = df_hull.e_above_mp_hull
 df["e_form_per_atom"] = df_wbm.e_form_per_atom
@@ -91,7 +94,7 @@
         ax=ax,
     )
 
-    title = f"Batch {batch_idx} ({len(batch_df):,})"
+    title = f"Batch {batch_idx} ({len(batch_df.filter(like='e_').dropna()):,})"
     ax.set(title=title)
 
 
@@ -103,13 +106,17 @@
     ax=axs.flat[-1],
 )
 
-axs.flat[-1].set(title=f"Combined {batch_idx} ({len(df):,})")
+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"
 )
-# plt.savefig(f"{ROOT}/figures/{img_name}")
+fig.suptitle(img_name.replace("-", "/", 2).replace("-", " "), y=1.07, fontsize=16)
+
+
+# %%
+ax.figure.savefig(f"{ROOT}/figures/{img_name}")
 
 
 # %%

mb_discovery/plot_scripts/precision_recall_vs_calc_count.py

@@ -4,7 +4,7 @@
 import pandas as pd
 
 from mb_discovery import ROOT
-from mb_discovery.plots import StabilityCriterion, plt, precision_recall_vs_calc_count
+from mb_discovery.plots import StabilityCriterion, precision_recall_vs_calc_count
 
 __author__ = "Rhys Goodall, Janosh Riebesell"
 __date__ = "2022-06-18"
@@ -24,14 +24,18 @@
     dfs[model_name] = df
 
 dfs["M3GNet"] = pd.read_json(
-    f"{ROOT}/models/m3gnet/2022-08-16-m3gnet-wbm-relax-results-IS2RE.json.gz"
+    f"{ROOT}/models/m3gnet/2022-08-16-m3gnet-wbm-IS2RE.json.gz"
 ).set_index("material_id")
 
 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(
+    f"{ROOT}/models/bowsr/2022-09-22-bowsr-wbm-megnet-IS2RE.json.gz"
+).set_index("material_id")
+
 print(f"loaded models: {list(dfs)}")
 
 
@@ -43,11 +47,9 @@
 
 # %%
 stability_crit: StabilityCriterion = "energy"
+colors = "tab:blue tab:orange teal tab:pink black red turquoise tab:purple".split()
 
-for (model_name, df), color in zip(
-    dfs.items(),
-    ("tab:blue", "tab:orange", "teal", "tab:pink", "black", "red", "turquoise"),
-):
+for (model_name, df), color in zip(dfs.items(), colors):
     rare = "all"
 
     # from pymatgen.core import Composition
@@ -76,6 +78,8 @@
             # 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:
+            model_preds = df.e_form_per_atom_bowsr
         else:
             raise ValueError(f"Unhandled {model_name = }")
     except AttributeError as exc:
@@ -103,6 +107,7 @@
 # keep this outside loop so all model names appear in legend
 ax.legend(frameon=False, loc="lower right")
 
+
+# %%
 img_path = f"{ROOT}/figures/{today}-precision-recall-vs-calc-count-{rare=}.pdf"
-if False:
-    plt.savefig(img_path)
+ax.figure.savefig(img_path)

mb_discovery/plot_scripts/rolling_mae_vs_hull_dist_wbm_batches.py

@@ -42,8 +42,8 @@
 assert len(markers) == 5  # number of WBM rounds of element substitution
 
 for idx, marker in enumerate(markers, 1):
-    title = f"Batch {idx}"
     df = df_wbm[df_wbm.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) = :,}")
 
     rolling_mae_vs_hull_dist(

mb_discovery/plots.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from typing import Any, Literal, Sequence, get_args
+from collections.abc import Sequence
+from typing import Any, Literal, get_args
 
 import matplotlib.pyplot as plt
 import numpy as np

models/bowsr/slurm_array_bowsr_wbm.py

@@ -47,21 +47,21 @@
 data_path = f"{ROOT}/data/2022-06-26-wbm-cses-and-initial-structures.json.gz"
 
 module_dir = os.path.dirname(__file__)
-job_id = os.environ.get("SLURM_JOB_ID", "debug")
-job_array_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))
+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))
 
 print(f"Job started running {datetime.now():%Y-%m-%d@%H-%M}")
-print(f"{job_id = }")
-print(f"{job_array_id = }")
+print(f"{slurm_job_id = }")
+print(f"{slurm_array_task_id = }")
 print(f"{version('maml') = }")
 print(f"{version('megnet') = }")
 
 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}/{job_array_id}.json.gz"
+json_out_path = f"{out_dir}/{slurm_array_task_id}.json.gz"
 
 if os.path.isfile(json_out_path):
     raise SystemExit(f"{json_out_path = } already exists, exciting early")
@@ -79,8 +79,8 @@
 run_params = dict(
     megnet_version=version("megnet"),
     maml_version=version("maml"),
-    job_id=job_id,
-    job_array_id=job_array_id,
+    slurm_job_id=slurm_job_id,
+    slurm_array_task_id=slurm_array_task_id,
     data_path=data_path,
     bayes_optim_kwargs=bayes_optim_kwargs,
     optimize_kwargs=optimize_kwargs,
@@ -93,7 +93,7 @@
 wandb.init(
     entity="janosh",
     project="matbench-discovery",
-    name=f"bowsr-megnet-wbm-{task_type}-{job_id}-{job_array_id}",
+    name=f"bowsr-megnet-wbm-{task_type}-{slurm_job_id}-{slurm_array_task_id}",
     config=run_params,
 )
 
@@ -102,7 +102,7 @@
 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)[job_array_id]
+df_this_job = np.array_split(df_wbm, job_array_size + 1)[slurm_array_task_id]
 
 
 # %%

models/m3gnet/eda_wbm_pre_vs_post_m3gnet_relaxation.py

@@ -28,10 +28,10 @@
 
 # %%
 df_m3gnet_is2re = pd.read_json(
-    f"{ROOT}/models/m3gnet/2022-08-16-m3gnet-wbm-relax-results-IS2RE.json.gz"
+    f"{ROOT}/models/m3gnet/2022-08-16-m3gnet-wbm-IS2RE.json.gz"
 ).set_index("material_id")
 df_m3gnet_rs2re = pd.read_json(
-    f"{ROOT}/models/m3gnet/2022-08-19-m3gnet-wbm-relax-results-RS2RE.json.gz"
+    f"{ROOT}/models/m3gnet/2022-08-19-m3gnet-wbm-RS2RE.json.gz"
 ).set_index("material_id")
 
 
@@ -226,5 +226,5 @@
 # %% write df back to compressed JSON
 # filter out columns containing 'rs2re'
 # df_m3gnet_is2re.reset_index().filter(regex="^((?!rs2re).)*$").to_json(
-#     f"{ROOT}/models/m3gnet/2022-08-16-m3gnet-wbm-relax-results-IS2RE-2.json.gz"
+#     f"{ROOT}/models/m3gnet/2022-08-16-m3gnet-wbm-IS2RE-2.json.gz"
 # ).set_index("material_id")

models/m3gnet/join_m3gnet_relax_results.py

@@ -118,5 +118,5 @@
 out_path = f"{ROOT}/models/m3gnet/{today}-m3gnet-wbm-relax-{task_type}.json.gz"
 df_m3gnet.reset_index().to_json(out_path, default_handler=as_dict_handler)
 
-# out_path = f"{ROOT}/models/m3gnet/2022-08-16-m3gnet-wbm-relax-results-IS2RE.json.gz"
+# out_path = f"{ROOT}/models/m3gnet/2022-08-16-m3gnet-wbm-IS2RE.json.gz"
 # df_m3gnet = pd.read_json(out_path).set_index("material_id")

models/m3gnet/slurm_array_m3gnet_relax_wbm.py

@@ -40,20 +40,20 @@
 task_type = "IS2RE"
 # task_type = "RS2RE"
 
-job_id = os.environ.get("SLURM_JOB_ID", "debug")
-job_array_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))
+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))
 
 print(f"Job started running {datetime.now():%Y-%m-%d@%H-%M}")
-print(f"{job_id = }")
-print(f"{job_array_id = }")
+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)
-json_out_path = f"{out_dir}/{job_array_id}.json.gz"
+json_out_path = f"{out_dir}/{slurm_array_task_id}.json.gz"
 
 if os.path.isfile(json_out_path):
     raise SystemExit(f"{json_out_path = } already exists, exciting early")
@@ -67,20 +67,20 @@
 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)[job_array_id]
+df_this_job = np.array_split(df_wbm, job_array_size)[slurm_array_task_id]
 
 run_params = dict(
     m3gnet_version=version("m3gnet"),
-    job_id=job_id,
-    job_array_id=job_array_id,
+    slurm_job_id=slurm_job_id,
+    slurm_array_task_id=slurm_array_task_id,
     data_path=data_path,
 )
 if wandb.run is None:
     wandb.login()
 
 wandb.init(
     project="m3gnet",
-    name=f"m3gnet-wbm-relax-{task_type}-{job_id}-{job_array_id}",
+    name=f"m3gnet-wbm-relax-{task_type}-{slurm_job_id}-{slurm_array_task_id}",
     config=run_params,
 )
 

readme.md

@@ -2,7 +2,7 @@
 
 [![Link check](https://github.com/janosh/matbench-discovery/actions/workflows/test.yml/badge.svg)](https://github.com/janosh/matbench-discovery/actions/workflows/test.yml)
 [![pre-commit.ci status](https://results.pre-commit.ci/badge/github/janosh/matbench-discovery/main.svg?badge_token=Qza33izjRxSbegTqeSyDvA)](https://results.pre-commit.ci/latest/github/janosh/matbench-discovery/main?badge_token=Qza33izjRxSbegTqeSyDvA)
-[![Requires Python 3.9+](https://img.shields.io/badge/Python-3.9+-blue.svg)](https://python.org/downloads)
+[![Requires Python 3.9+](https://img.shields.io/badge/Python-3.9+-blue.svg?logo=python)](https://python.org/downloads)
 
 Several new energy models specifically designed to handle unrelaxed structures were published in 2021/22
 

tests/test_plots.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from typing import Any, Sequence
+from collections.abc import Sequence
+from typing import Any
 
 import matplotlib.pyplot as plt
 import pandas as pd