add models/voronoi/train_voronoi_rf.py

move featurizer def from models/voronoi/featurize_mp_wbm.py to new models/voronoi/__init__.py

Author: janosh

models/cgcnn/train_cgcnn.py

@@ -15,7 +15,7 @@
 from matbench_discovery.slurm import slurm_submit
 
 """
-Train a Wrenformer ensemble of size n_folds on target_col of data_path.
+Train a CGCNN ensemble of size n_folds on target_col of data_path.
 """
 
 __author__ = "Janosh Riebesell"
@@ -25,7 +25,7 @@
 # %%
 epochs = 300
 target_col = "formation_energy_per_atom"
-run_name = f"cgcnn-robust-{target_col}-{epochs=}"
+run_name = f"cgcnn-robust-{target_col}"
 print(f"{run_name=}")
 robust = "robust" in run_name.lower()
 n_folds = 10

models/m3gnet/test_m3gnet.py

@@ -32,8 +32,7 @@
 # set large job array size for fast testing/debugging
 slurm_array_task_count = 100
 slurm_mem_per_node = 12000
-slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
-job_name = f"m3gnet-wbm-{task_type}-{slurm_job_id}"
+job_name = f"m3gnet-wbm-{task_type}"
 out_dir = f"{module_dir}/{today}-{job_name}"
 
 slurm_vars = slurm_submit(
@@ -85,9 +84,10 @@
 if wandb.run is None:
     wandb.login()
 
+slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
 wandb.init(
     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,
 )
 

models/voronoi/__init__.py

@@ -0,0 +1,24 @@
+import matminer.featurizers.composition as feat_comp
+import matminer.featurizers.structure as feat_struct
+from matminer.featurizers.base import MultipleFeaturizer
+
+# Create the featurizer: Ward et al. use a variety of different featurizers
+# https://journals.aps.org/prb/abstract/10.1103/PhysRevB.96.024104
+featurizers = [
+    feat_struct.SiteStatsFingerprint.from_preset("CoordinationNumber_ward-prb-2017"),
+    feat_struct.StructuralHeterogeneity(),
+    feat_struct.ChemicalOrdering(),
+    feat_struct.MaximumPackingEfficiency(),
+    feat_struct.SiteStatsFingerprint.from_preset(
+        "LocalPropertyDifference_ward-prb-2017"
+    ),
+    feat_struct.StructureComposition(feat_comp.Stoichiometry()),
+    feat_struct.StructureComposition(feat_comp.ElementProperty.from_preset("magpie")),
+    feat_struct.StructureComposition(feat_comp.ValenceOrbital(props=["frac"])),
+    feat_struct.StructureComposition(feat_comp.IonProperty(fast=True)),
+]
+featurizer = MultipleFeaturizer(featurizers)
+
+# multiprocessing seems to be the cause of OOM errors on large structures even when
+# taking only small slice of the data and launching slurm jobs with --mem 100G
+featurizer.set_n_jobs(1)

models/voronoi/featurize_mp_wbm.py models/voronoi/voronoi_featurize_dataset.py

@@ -3,26 +3,28 @@
 import warnings
 from datetime import datetime
 
-import matminer.featurizers.composition as feat_comp
-import matminer.featurizers.structure as feat_struct
 import numpy as np
 import pandas as pd
 import wandb
-from matminer.featurizers.base import MultipleFeaturizer
 from pymatgen.core import Structure
 from tqdm import tqdm
 
 from matbench_discovery import ROOT, as_dict_handler
 from matbench_discovery.slurm import slurm_submit
+from models.voronoi import featurizer
 
 today = f"{datetime.now():%Y-%m-%d}"
 module_dir = os.path.dirname(__file__)
 
 
-data_path = f"{ROOT}/data/mp/2022-09-16-mp-computed-structure-entries.json.gz"
-# data_path = f"{ROOT}/data/wbm/2022-10-19-wbm-init-structs.json.bz2"
-input_col = "initial_structure"
-data_name = "wbm" if "wbm" in data_path else "mp"
+data_name = "mp"  # "mp"
+if data_name == "wbm":
+    data_path = f"{ROOT}/data/wbm/2022-10-19-wbm-init-structs.json.bz2"
+    input_col = "initial_structure"
+elif data_name == "mp":
+    data_path = f"{ROOT}/data/mp/2022-09-16-mp-computed-structure-entries.json.gz"
+    input_col = "structure"
+
 slurm_array_task_count = 10
 job_name = f"voronoi-features-{data_name}"
 log_dir = f"{module_dir}/{today}-{job_name}"
@@ -39,7 +41,8 @@
 
 # %%
 slurm_array_task_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))
-run_name = f"{job_name}-{slurm_array_task_id}"
+slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
+run_name = f"{job_name}-{slurm_job_id}-{slurm_array_task_id}"
 out_path = f"{log_dir}/{run_name}.csv.bz2"
 
 if os.path.isfile(out_path):
@@ -57,7 +60,7 @@
     struct_dicts = df_this_job.initial_structure
 
 df_this_job[input_col] = [
-    Structure.from_dict(x) for x in tqdm(df_this_job.initial_structure, disable=None)
+    Structure.from_dict(x) for x in tqdm(struct_dicts, disable=None)
 ]
 
 
@@ -79,27 +82,6 @@
 )
 
 
-# %% Create the featurizer: Ward et al. use a variety of different featurizers
-# https://journals.aps.org/prb/abstract/10.1103/PhysRevB.96.024104
-featurizers = [
-    feat_struct.SiteStatsFingerprint.from_preset("CoordinationNumber_ward-prb-2017"),
-    feat_struct.StructuralHeterogeneity(),
-    feat_struct.ChemicalOrdering(),
-    feat_struct.MaximumPackingEfficiency(),
-    feat_struct.SiteStatsFingerprint.from_preset(
-        "LocalPropertyDifference_ward-prb-2017"
-    ),
-    feat_struct.StructureComposition(feat_comp.Stoichiometry()),
-    feat_struct.StructureComposition(feat_comp.ElementProperty.from_preset("magpie")),
-    feat_struct.StructureComposition(feat_comp.ValenceOrbital(props=["frac"])),
-    feat_struct.StructureComposition(feat_comp.IonProperty(fast=True)),
-]
-featurizer = MultipleFeaturizer(featurizers)
-# multiprocessing seems to be the cause of OOM errors on large structures even when
-# taking only small slice of the data and launching slurm jobs with --mem 100G
-featurizer.set_n_jobs(1)
-
-
 # %% prints lots of pymatgen warnings
 # > No electronegativity for Ne. Setting to NaN. This has no physical meaning, ...
 warnings.filterwarnings(action="ignore", category=UserWarning, module="pymatgen")

models/wrenformer/train_wrenformer.py

@@ -24,7 +24,7 @@
 # data_path = f"{ROOT}/data/2022-08-25-m3gnet-trainset-mp-2021-struct-energy.json.gz"
 # target_col = "mp_energy_per_atom"
 data_name = "m3gnet-trainset" if "m3gnet" in data_path else "mp"
-run_name = f"train-wrenformer-robust-{data_name}-{target_col}-{epochs=}"
+run_name = f"train-wrenformer-robust-{data_name}-{target_col}"
 n_folds = 10
 timestamp = f"{datetime.now():%Y-%m-%d@%H-%M-%S}"
 today = timestamp.split("@")[0]
@@ -65,8 +65,9 @@
     test_df=dict(shape=test_df.shape, columns=", ".join(test_df)),
 )
 
+slurm_job_id = os.environ.get("SLURM_JOB_ID", "debug")
 train_wrenformer(
-    run_name=run_name,
+    run_name=f"{run_name}-{slurm_job_id}",
     train_df=train_df,
     test_df=test_df,
     target_col=target_col,