add m3gnet and bowsr-megnet to plot moving_hull_dist_mae_compare_models.py

record energy_model in test_bowsr.py run_params

Author: janosh

matbench_discovery/plot_scripts/hist_classified_stable_vs_hull_dist_batches.py

@@ -60,7 +60,7 @@
     df[pred_col] = df.e_form_per_atom_m3gnet
 if "bowsr_megnet" in dfs:
     df = dfs["bowsr_megnet"]
-    df[pred_col] = df.e_form_per_atom_bowsr
+    df[pred_col] = df.e_form_per_atom_bowsr_megnet
 if "wrenformer" in dfs:
     pred_col = "e_form_per_atom_mp2020_corrected_pred_ens"
 

matbench_discovery/plot_scripts/precision_recall.py

@@ -43,7 +43,7 @@
 colors = "tab:blue tab:orange teal tab:pink black red turquoise tab:purple".split()
 F1s: dict[str, float] = {}
 
-for model_name, df in dfs.items():
+for model_name, df in sorted(dfs.items()):
     if "std" in stability_crit:
         # TODO column names to compute standard deviation from are currently hardcoded
         # needs to be updated when adding non-aviary models with uncertainty estimation
@@ -63,8 +63,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
+        elif model_name == "bowsr_megnet":
+            model_preds = df.e_form_per_atom_bowsr_megnet
         else:
             raise ValueError(f"Unhandled {model_name = }")
     except AttributeError as exc:

models/bowsr/join_bowsr_results.py

@@ -33,7 +33,6 @@
 for file_path in tqdm(file_paths):
     if file_path in dfs:
         continue
-    # keep whole dataframe in memory
     df = pd.read_json(file_path).set_index("material_id")
 
     df["bowsr_structure"] = df.structure_bowsr.map(Structure.from_dict)
@@ -51,20 +50,14 @@
 data_path = f"{ROOT}/data/wbm/2022-10-19-wbm-summary.csv"
 df_wbm = pd.read_csv(data_path).set_index("material_id")
 
-df_bowsr["e_form_wbm"] = df_wbm.e_form_per_atom
 
-print(f"{len(df_bowsr) - len(df_wbm) = :,} = {len(df_bowsr):,} - {len(df_wbm):,}")
-
-
-# %%
-df_bowsr.hist(bins=200, figsize=(18, 12))
-df_bowsr.isna().sum()
+print(f"{len(df_bowsr):,} - {len(df_wbm):,} = {len(df_bowsr) - len(df_wbm) = :,}")
 
 
 # %%
 pymatviz.density_scatter(
-    df_bowsr.dropna().e_form_per_atom_bowsr,
-    df_bowsr.dropna().e_form_wbm,
+    x=df_bowsr.e_form_per_atom_bowsr_megnet,
+    y=df_bowsr.e_form_wbm,
 )
 
 

models/bowsr/test_bowsr.py

@@ -35,8 +35,8 @@
 slurm_array_task_count = 500
 timestamp = f"{datetime.now():%Y-%m-%d@%H-%M-%S}"
 today = timestamp.split("@")[0]
-slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
-job_name = f"bowsr-megnet-wbm-{task_type}-{slurm_job_id}"
+energy_model = "megnet"
+job_name = f"bowsr-{energy_model}-wbm-{task_type}"
 out_dir = f"{module_dir}/{today}-{job_name}"
 
 data_path = f"{ROOT}/data/wbm/2022-10-19-wbm-init-structs.json.bz2"
@@ -64,7 +64,7 @@
 print(f"{data_path = }")
 print(f"{out_path = }")
 print(f"{version('maml') = }")
-print(f"{version('megnet') = }")
+print(f"{version(energy_model) = }")
 
 
 if os.path.isfile(out_path):
@@ -94,7 +94,8 @@
     data_path=data_path,
     df=dict(shape=str(df_this_job.shape), columns=", ".join(df_this_job)),
     maml_version=version("maml"),
-    megnet_version=version("megnet"),
+    energy_model=energy_model,
+    energy_model_version=version(energy_model),
     optimize_kwargs=optimize_kwargs,
     task_type=task_type,
     slurm_max_job_time=slurm_max_job_time,
@@ -103,12 +104,11 @@
 if wandb.run is None:
     wandb.login()
 
-# getting wandb: 429 encountered ({"error":"rate limit exceeded"}), retrying request
-# https://community.wandb.ai/t/753/14
+slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
 wandb.init(
     entity="janosh",
     project="matbench-discovery",
-    name=f"{job_name}-{slurm_array_task_id}",
+    name=f"{job_name}-{slurm_job_id}-{slurm_array_task_id}",
     config=run_params,
 )
 
@@ -146,11 +146,11 @@
 
     structure_bowsr, energy_bowsr = bayes_optimizer.get_optimized_structure_and_energy()
 
-    results = dict(
-        e_form_per_atom_bowsr=model.predict_energy(structure),
-        structure_bowsr=structure_bowsr,
-        energy_bowsr=energy_bowsr,
-    )
+    results = {
+        f"e_form_per_atom_bowsr_{energy_model}": model.predict_energy(structure),
+        "structure_bowsr": structure_bowsr,
+        "energy_bowsr": energy_bowsr,
+    }
 
     relax_results[material_id] = results