matbench_discovery/__init__.py define CHECKPOINT_DIR = f"{ROOT}/wandb/checkpoints"

test_cgcnn remove reset_index from cg_data.df.reset_index(drop=True)

Author: janosh

.gitignore

@@ -23,5 +23,5 @@ models/**/*.csv
 
 # temporary ignore rules
 paper
-meeting-notes
 models/voronoi/*.zip
+site

matbench_discovery/__init__.py

@@ -8,6 +8,7 @@
 
 ROOT = os.path.dirname(os.path.dirname(__file__))
 DEBUG = "slurm-submit" not in sys.argv and "SLURM_JOB_ID" not in os.environ
+CHECKPOINT_DIR = f"{ROOT}/wandb/checkpoints"
 
 timestamp = f"{datetime.now():%Y-%m-%d@%H-%M-%S}"
 today = timestamp.split("@")[0]

models/cgcnn/test_cgcnn.py

@@ -14,7 +14,7 @@
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
-from matbench_discovery import DEBUG, ROOT, today
+from matbench_discovery import CHECKPOINT_DIR, DEBUG, ROOT, today
 from matbench_discovery.load_preds import df_wbm
 from matbench_discovery.plots import wandb_scatter
 from matbench_discovery.slurm import slurm_submit
@@ -23,9 +23,9 @@
 __date__ = "2022-08-15"
 
 """
-Script that downloads checkpoints for an ensemble of Wrenformer models trained on the MP
+Script that downloads checkpoints for an ensemble of CGCNN models trained on all MP
 formation energies, then makes predictions on some dataset, prints ensemble metrics and
-stores predictions to CSV.
+saves predictions to CSV.
 """
 
 task_type = "RS2RE"
@@ -54,7 +54,7 @@
 else:
     raise ValueError(f"Unexpected {task_type=}")
 
-df = pd.read_json(data_path).set_index("material_id", drop=False)
+df = pd.read_json(data_path).set_index("material_id")
 
 target_col = "e_form_per_atom_mp2020_corrected"
 df[target_col] = df_wbm[target_col]
@@ -88,7 +88,7 @@
     task_type=task_type,
     target_col=target_col,
     input_col=input_col,
-    filters=filters,
+    wandb_run_filters=filters,
     slurm_vars=slurm_vars,
 )
 
@@ -99,15 +99,16 @@
     df,
     task_dict={target_col: "regression"},
     structure_col=input_col,
-    identifiers=("material_id", "formula_from_cse"),
+    identifiers=["formula_from_cse"],
 )
 data_loader = DataLoader(
     cg_data, batch_size=1024, shuffle=False, collate_fn=collate_batch
 )
 df, ensemble_metrics = predict_from_wandb_checkpoints(
     runs,
     # dropping isolated-atom structs means len(cg_data.df) < len(df)
-    df=cg_data.df.reset_index(drop=True).drop(columns=input_col),
+    cache_dir=CHECKPOINT_DIR,
+    df=cg_data.df.drop(columns=input_col),
     target_col=target_col,
     model_cls=CrystalGraphConvNet,
     data_loader=data_loader,
@@ -122,6 +123,6 @@
 MAE = ensemble_metrics.MAE.mean()
 R2 = ensemble_metrics.R2.mean()
 
-title = rf"CGCNN {task_type} ensemble={len(runs)} {MAE=:.4} {R2=:.4}"
+title = f"CGCNN {task_type} ensemble={len(runs)} {MAE=:.4} {R2=:.4}"
 
 wandb_scatter(table, fields=dict(x=target_col, y=pred_col), title=title)

models/wrenformer/test_wrenformer.py

@@ -11,17 +11,17 @@
 from aviary.wrenformer.data import df_to_in_mem_dataloader
 from aviary.wrenformer.model import Wrenformer
 
-from matbench_discovery import DEBUG, ROOT, today
+from matbench_discovery import CHECKPOINT_DIR, DEBUG, ROOT, today
 from matbench_discovery.plots import wandb_scatter
 from matbench_discovery.slurm import slurm_submit
 
 __author__ = "Janosh Riebesell"
 __date__ = "2022-08-15"
 
 """
-Download WandB checkpoints for an ensemble of Wrenformer models trained on MP
+Download WandB checkpoints for an ensemble of Wrenformer models trained on all MP
 formation energies, then makes predictions on some dataset, prints ensemble metrics and
-stores predictions to CSV.
+saves predictions to CSV.
 """
 
 task_type = "IS2RE"
@@ -74,7 +74,7 @@
     task_type=task_type,
     target_col=target_col,
     input_col=input_col,
-    filters=filters,
+    wandb_run_filters=filters,
     slurm_vars=slurm_vars,
 )
 
@@ -84,6 +84,7 @@
 # %%
 data_loader = df_to_in_mem_dataloader(
     df=df,
+    cache_dir=CHECKPOINT_DIR,
     target_col=target_col,
     batch_size=1024,
     input_col=input_col,
@@ -108,6 +109,6 @@
 MAE = ensemble_metrics.MAE.mean()
 R2 = ensemble_metrics.R2.mean()
 
-title = rf"Wrenformer {task_type} ensemble={len(runs)} {MAE=:.4} {R2=:.4}"
+title = f"Wrenformer {task_type} ensemble={len(runs)} {MAE=:.4} {R2=:.4}"
 
 wandb_scatter(table, fields=dict(x=target_col, y=pred_col), title=title)