create no-structure CSVs of m3gnet and bowsr-megnet preds for faster loading used in load_df_wbm_with_preds()

add test_glob_to_df() and dummy_df fixture

Author: janosh

matbench_discovery/plot_scripts/precision_recall.py

@@ -10,8 +10,8 @@
 
 # %%
 models = (
-    "Wren, CGCNN IS2RE, CGCNN RS2RE, Voronoi IS2RE, Voronoi RS2RE, "
-    "Wrenformer, MEGNet"
+    "Wren, CGCNN IS2RE, CGCNN RS2RE, Voronoi RF, "
+    "Wrenformer, MEGNet, M3GNet, BOWSR MEGNet"
 ).split(", ")
 
 df_wbm = load_df_wbm_with_preds(models=models).round(3)

models/bowsr/join_bowsr_results.py

@@ -35,7 +35,7 @@
 
 
 # %%
-df_bowsr = pd.concat(dfs.values())
+df_bowsr = pd.concat(dfs.values()).round(6)
 
 
 # %% compare against WBM formation energy targets to make sure we got sensible results
@@ -57,5 +57,8 @@
 out_path = f"{ROOT}/models/bowsr/{today}-bowsr-megnet-wbm-{task_type}.json.gz"
 df_bowsr.reset_index().to_json(out_path, default_handler=lambda x: x.as_dict())
 
-# out_path = f"{ROOT}/models/bowsr/2022-08-16-bowsr-megnet-wbm-IS2RE.json.gz"
-# df_bowsr = pd.read_json(out_path).set_index("material_id")
+# save energy and formation energy as CSV for fast loading
+df_bowsr.select_dtypes("number").to_csv(out_path.replace(".json.gz", ".csv"))
+
+# in_path = f"{ROOT}/models/bowsr/2022-11-22-bowsr-megnet-wbm-IS2RE.json.gz"
+# df_bowsr = pd.read_json(in_path).set_index("material_id")

models/m3gnet/join_m3gnet_results.py

@@ -27,30 +27,24 @@
 
 
 # %%
-# 2022-08-16 tried multiprocessing.Pool() to load files in parallel but was somehow
-# slower than serial loading
 for file_path in tqdm(file_paths):
     if file_path in dfs:
         continue
     df = pd.read_json(file_path).set_index("material_id")
-    df.index.name = "material_id"
-    col_map = dict(final_structure="structure_m3gnet", trajectory="m3gnet_trajectory")
-    df = df.rename(columns=col_map)
-    df.reset_index().to_json(file_path)
-    df[f"m3gnet_energy_{task_type}"] = df.m3gnet_trajectory.map(
-        lambda x: x["energies"][-1][0]
-    )
+    df[f"m3gnet_energy_{task_type}"] = [
+        x["energies"][-1][0] for x in df.m3gnet_trajectory
+    ]
     # drop trajectory to save memory
-    dfs[file_path] = df.drop(columns=["m3gnet_trajectory"])
+    dfs[file_path] = df.drop(columns="m3gnet_trajectory")
 
 
 # %%
-df_m3gnet = pd.concat(dfs.values())
+df_m3gnet = pd.concat(dfs.values()).round(6)
 
 
 # %%
 df_m3gnet["e_form_per_atom_m3gnet"] = [
-    get_e_form_per_atom(PDEntry(row.structure_m3gnet.composition, row.m3gnet_energy))
+    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()
@@ -60,5 +54,7 @@
 out_path = f"{ROOT}/models/m3gnet/{today}-m3gnet-wbm-{task_type}.json.gz"
 df_m3gnet.reset_index().to_json(out_path, default_handler=as_dict_handler)
 
-# out_path = f"{ROOT}/models/m3gnet/2022-10-31-m3gnet-wbm-IS2RE.json.gz"
-# df_m3gnet = pd.read_json(out_path).set_index("material_id")
+df_m3gnet.select_dtypes("number").to_csv(out_path.replace(".json.gz", ".csv"))
+
+# in_path = f"{ROOT}/models/m3gnet/2022-10-31-m3gnet-wbm-IS2RE.json.gz"
+# df_m3gnet = pd.read_json(in_path).set_index("material_id")

models/voronoi/join_voronoi_features.py

@@ -22,8 +22,6 @@
 
 
 # %%
-# 2022-08-16 tried multiprocessing.Pool() to load files in parallel but was somehow
-# slower than serial loading
 for file_path in tqdm(file_paths):
     if file_path in dfs:
         continue
@@ -32,7 +30,7 @@
 
 
 # %%
-df_features = pd.concat(dfs.values())
+df_features = pd.concat(dfs.values()).round(6)
 
 ax = df_features.isna().sum().value_counts().T.plot.bar()
 ax.set(xlabel="# NaNs", ylabel="# columns", title="NaNs per column")

models/voronoi/train_test_voronoi_rf.py

@@ -10,7 +10,7 @@
 from sklearn.pipeline import Pipeline
 
 from matbench_discovery import DEBUG, ROOT, today
-from matbench_discovery.plot_scripts import df_wbm
+from matbench_discovery.plot_scripts import df_wbm, glob_to_df
 from matbench_discovery.plots import wandb_log_scatter
 from matbench_discovery.slurm import slurm_submit
 from models.voronoi import featurizer
@@ -41,24 +41,29 @@
 
 
 # %%
-train_path = f"{module_dir}/2022-11-25-features-mp.csv.bz2"
-print(f"{train_path=}")
-df_train = pd.read_csv(train_path).set_index("material_id")
+train_path = f"{module_dir}/2022-11-25-features-mp/voronoi-features-mp-*.csv.bz2"
+df_train = glob_to_df(train_path).set_index("material_id")
 print(f"{df_train.shape=}")
 
 mp_energies_path = f"{ROOT}/data/mp/2022-08-13-mp-energies.json.gz"
 df_mp = pd.read_json(mp_energies_path).set_index("material_id")
 train_target_col = "formation_energy_per_atom"
-df_train[train_target_col] = df_mp[train_target_col]
-
 
 test_path = f"{module_dir}/2022-11-18-features-wbm-{task_type}.csv.bz2"
-print(f"{test_path=}")
 df_test = pd.read_csv(test_path).set_index("material_id")
 print(f"{df_test.shape=}")
 
 test_target_col = "e_form_per_atom_mp2020_corrected"
-df_test[test_target_col] = df_wbm[test_target_col]
+
+
+for df, df_tar, col in (
+    (df_train, df_mp, train_target_col),
+    (df_test, df_wbm, test_target_col),
+):
+    df[train_target_col] = df_tar[train_target_col]
+    nans = df_tar[col].isna().sum()
+    assert nans == 0, f"{nans} NaNs in {col} targets"
+
 model_name = "Voronoi RandomForestRegressor"
 
 run_params = dict(