fix slurm_array_m3gnet_relax_wbm.py wandb project m3gnet->matbench-discovery

update
data_path = f"{ROOT}/data/2022-06-26-wbm-cses-and-initial-structures.json.gz"
to
data_path = f"{ROOT}/data/wbm/2022-10-19-wbm-cses+init-structs.json.bz2"

Author: janosh

models/bowsr/slurm_array_bowsr_wbm.py

@@ -21,9 +21,7 @@
 __date__ = "2022-08-15"
 
 """
-To slurm submit this file, run:
-
-python path/to/file.py slurm-submit
+To slurm submit this file: python path/to/file.py slurm-submit
 
 Requires MEGNet and MAML installation: pip install megnet maml
 """
@@ -112,7 +110,9 @@
 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, slurm_array_task_count)[slurm_array_task_id - 1]
+df_this_job: pd.DataFrame = np.array_split(df_wbm, slurm_array_task_count)[
+    slurm_array_task_id - 1
+]
 
 
 # %%

models/m3gnet/join_m3gnet_relax_results.py

@@ -21,8 +21,8 @@
 
 # %%
 module_dir = os.path.dirname(__file__)
-task_type = "RS2RE"
-date = "2022-08-19"
+task_type = "IS2RE"
+date = "2022-10-31"
 glob_pattern = f"{date}-m3gnet-wbm-relax-{task_type}/*.json.gz"
 file_paths = sorted(glob(f"{module_dir}/{glob_pattern}"))
 print(f"Found {len(file_paths):,} files for {glob_pattern = }")
@@ -38,8 +38,7 @@
         continue
     try:
         # keep whole dataframe in memory
-        df = pd.read_json(file_path)
-        df.index = df.index.str.replace("_", "-")
+        df = pd.read_json(file_path).set_index("material_id")
         df.index.name = "material_id"
         col_map = dict(
             final_structure="m3gnet_structure", trajectory="m3gnet_trajectory"
@@ -52,24 +51,18 @@
         df["volume"] = df.m3gnet_structure.map(lambda x: x.volume)
         df["n_sites"] = df.m3gnet_structure.map(len)
         dfs[file_path] = df.drop(columns=["m3gnet_trajectory"])
-    except (ValueError, FileNotFoundError):
-        # pandas v1.5+ correctly raises FileNotFoundError, below raises ValueError
+    except FileNotFoundError:
         continue
 
 
 # %%
 df_m3gnet = pd.concat(dfs.values())
-if any(df_m3gnet.index.str.contains("_")):
-    df_m3gnet.index = df_m3gnet.index.str.replace("_", "-")
 
 
 # %%
-pd_entries_m3gnet = [
-    PDEntry(row.m3gnet_structure.composition, row.m3gnet_energy)
-    for row in df_m3gnet.itertuples()
-]
-df_m3gnet["e_form_m3gnet_from_ppd"] = [
-    get_e_form_per_atom(entry) for entry in pd_entries_m3gnet
+df_m3gnet["e_form_per_atom_m3gnet"] = [
+    get_e_form_per_atom(PDEntry(row.m3gnet_structure.composition, row.m3gnet_energy))
+    for row in tqdm(df_m3gnet.itertuples(), total=len(df_m3gnet), disable=None)
 ]
 df_m3gnet.isna().sum()
 

models/m3gnet/slurm_array_m3gnet_relax_wbm.py

@@ -17,9 +17,7 @@
 from matbench_discovery.slurm import slurm_submit_python
 
 """
-To slurm submit this file, run:
-
-python path/to/file.py slurm-submit
+To slurm submit this file: python path/to/file.py slurm-submit
 
 Requires M3GNet installation: pip install m3gnet
 """
@@ -69,11 +67,13 @@
 
 
 # %%
-data_path = f"{ROOT}/data/2022-06-26-wbm-cses-and-initial-structures.json.gz"
+data_path = f"{ROOT}/data/wbm/2022-10-19-wbm-cses+init-structs.json.bz2"
 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, slurm_array_task_count)[slurm_array_task_id - 1]
+df_this_job: pd.DataFrame = np.array_split(df_wbm, slurm_array_task_count)[
+    slurm_array_task_id - 1
+]
 
 run_params = dict(
     data_path=data_path,
@@ -89,7 +89,7 @@
     wandb.login()
 
 wandb.init(
-    project="m3gnet",
+    project="matbench-discovery",
     name=f"{job_name}-{slurm_job_id}-{slurm_array_task_id}",
     config=run_params,
 )