add df metadata to slurm_array_{m3gnet,bowsr}_wbm.py run_params

janosh · Jun 20, 2023 · f127a9e · f127a9e
1 parent 444fb7f
commit f127a9e
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 19 deletions.
diff --git a/models/bowsr/slurm_array_bowsr_wbm.py b/models/bowsr/slurm_array_bowsr_wbm.py
@@ -71,6 +71,15 @@
     raise SystemExit(f"{out_path = } already exists, exciting early")
 
 
+# %%
+print(f"Loading from {data_path = }")
+df_wbm = pd.read_json(data_path).set_index("material_id")
+
+df_this_job: pd.DataFrame = np.array_split(df_wbm, slurm_array_task_count)[
+    slurm_array_task_id - 1
+]
+
+
 # %%
 bayes_optim_kwargs = dict(
     relax_coords=True,
@@ -83,6 +92,7 @@
 run_params = dict(
     bayes_optim_kwargs=bayes_optim_kwargs,
     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"),
     optimize_kwargs=optimize_kwargs,
@@ -104,15 +114,6 @@
 )
 
 
-# %%
-print(f"Loading from {data_path = }")
-df_wbm = pd.read_json(data_path).set_index("material_id")
-
-df_this_job: pd.DataFrame = np.array_split(df_wbm, slurm_array_task_count)[
-    slurm_array_task_id - 1
-]
-
-
 # %%
 model = MEGNet()
 relax_results: dict[str, dict[str, Any]] = {}

diff --git a/models/cgcnn/slurm_train_cgcnn_ensemble.py b/models/cgcnn/slurm_train_cgcnn_ensemble.py
@@ -88,8 +88,8 @@
 run_params = dict(
     data_path=data_path,
     batch_size=batch_size,
-    train_df=dict(shape=train_data.df.shape, columns=", ".join(train_df)),
-    test_df=dict(shape=test_data.df.shape, columns=", ".join(test_df)),
+    train_df=dict(shape=str(train_data.df.shape), columns=", ".join(train_df)),
+    test_df=dict(shape=str(test_data.df.shape), columns=", ".join(test_df)),
 )
 
 

diff --git a/models/m3gnet/slurm_array_m3gnet_wbm.py b/models/m3gnet/slurm_array_m3gnet_wbm.py
@@ -80,6 +80,7 @@
     slurm_array_task_count=slurm_array_task_count,
     task_type=task_type,
     slurm_max_job_time=slurm_max_job_time,
+    df=dict(shape=str(df_this_job.shape), columns=", ".join(df_this_job)),
     **slurm_vars,
 )
 if wandb.run is None:

diff --git a/models/wrenformer/mp/use_wrenformer_ensemble.py b/models/wrenformer/mp/use_wrenformer_ensemble.py
@@ -46,14 +46,6 @@
 assert target_col in df, f"{target_col=} not in {list(df)}"
 assert input_col in df, f"{input_col=} not in {list(df)}"
 
-wandb.login()
-wandb_api = wandb.Api()
-runs = wandb_api.runs(
-    "janosh/matbench-discovery", filters={"tags": {"$in": [ensemble_id]}}
-)
-
-assert len(runs) == 10, f"Expected 10 runs, got {len(runs)} for {ensemble_id=}"
-
 data_loader = df_to_in_mem_dataloader(
     df=df,
     target_col=target_col,
@@ -63,6 +55,22 @@
     shuffle=False,  # False is default but best be explicit
 )
 
+
+# %%
+wandb.login()
+wandb_api = wandb.Api()
+runs = wandb_api.runs(
+    "janosh/matbench-discovery",
+    filters={
+        "$and": [{"created_at": {"$gt": "2022-11-10", "$lt": "2022-11-11"}}],
+        "display_name": "wrenformer-robust-mp-formation_energy_per_atom-epochs=300",
+    },
+)
+
+assert len(runs) == 10, f"Expected 10 runs, got {len(runs)} for {ensemble_id=}"
+
+
+# %%
 df, ensemble_metrics = predict_from_wandb_checkpoints(
     runs, data_loader=data_loader, df=df, model_cls=Wrenformer
 )