change ModelCard heading color scale from PuOr to Cividis

define reusable col names in matbench_discovery/__init__.py

Author: janosh

data/mp/build_phase_diagram.py

@@ -17,7 +17,7 @@
 from pymatgen.ext.matproj import MPRester
 from tqdm import tqdm
 
-from matbench_discovery import ROOT, today
+from matbench_discovery import ROOT, id_col, today
 from matbench_discovery.data import DATA_FILES
 from matbench_discovery.energy import get_e_form_per_atom, get_elemental_ref_entries
 
@@ -30,7 +30,7 @@
 # save all ComputedStructureEntries to disk
 # mp-15590 appears twice so we drop_duplicates()
 df = pd.DataFrame(all_mp_computed_structure_entries, columns=["entry"])
-df.index.name = "material_id"
+df.index.name = id_col
 df.index = [e.entry_id for e in df.entry]
 df.reset_index().to_json(
     f"{module_dir}/{today}-mp-computed-structure-entries.json.gz",
@@ -40,7 +40,7 @@
 
 # %%
 data_path = f"{module_dir}/2023-02-07-mp-computed-structure-entries.json.gz"
-df = pd.read_json(data_path).set_index("material_id")
+df = pd.read_json(data_path).set_index(id_col)
 
 # drop the structure, just load ComputedEntry, makes the PPD faster to build and load
 mp_computed_entries = [ComputedEntry.from_dict(dct) for dct in tqdm(df.entry)]
@@ -63,9 +63,7 @@
 
 
 # %% build phase diagram with both MP entries + WBM entries
-df_wbm = pd.read_json(DATA_FILES.wbm_computed_structure_entries).set_index(
-    "material_id"
-)
+df_wbm = pd.read_json(DATA_FILES.wbm_computed_structure_entries).set_index(id_col)
 
 # using ComputedStructureEntry vs ComputedEntry here is important as CSEs receive
 # more accurate energy corrections that take into account peroxide/superoxide nature
@@ -104,7 +102,7 @@
     json.dump(elemental_ref_entries, file, default=lambda x: x.as_dict())
 
 
-df_mp = pd.read_csv(DATA_FILES.mp_energies, na_filter=False).set_index("material_id")
+df_mp = pd.read_csv(DATA_FILES.mp_energies, na_filter=False).set_index(id_col)
 
 
 # %%

data/mp/get_mp_energies.py

@@ -8,7 +8,7 @@
 from pymatviz.utils import annotate_metrics
 from tqdm import tqdm
 
-from matbench_discovery import STABILITY_THRESHOLD, today
+from matbench_discovery import STABILITY_THRESHOLD, id_col, today
 from matbench_discovery.data import DATA_FILES
 
 """
@@ -26,7 +26,7 @@
 
 # %%
 fields = {
-    "material_id",
+    id_col,
     "formula_pretty",
     "formation_energy_per_atom",
     "energy_per_atom",
@@ -46,7 +46,7 @@
 
 
 # %%
-df = pd.DataFrame(docs).set_index("material_id")
+df = pd.DataFrame(docs).set_index(id_col)
 
 df_spg = pd.json_normalize(df.pop("symmetry"))[["number", "symbol"]]
 df["spacegroup_symbol"] = df_spg.symbol.to_numpy()
@@ -56,7 +56,7 @@
 
 
 # %%
-df_cse = pd.read_json(DATA_FILES.mp_computed_structure_entries).set_index("material_id")
+df_cse = pd.read_json(DATA_FILES.mp_computed_structure_entries).set_index(id_col)
 
 struct_col = "structure"
 df_cse[struct_col] = [
@@ -76,7 +76,7 @@
 assert (spg_nums.sort_index() == df_spg["number"].sort_index()).all()
 
 df.to_csv(DATA_FILES.mp_energies)
-# df = pd.read_csv(DATA_FILES.mp_energies, na_filter=False).set_index("material_id")
+# df = pd.read_csv(DATA_FILES.mp_energies, na_filter=False).set_index(id_col)
 
 
 # %% reproduce fig. 1b from https://arxiv.org/abs/2001.10591 (as data consistency check)

data/wbm/compare_cse_vs_ce_mp_2020_corrections.py

@@ -10,7 +10,7 @@
 from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry
 from tqdm import tqdm
 
-from matbench_discovery import ROOT, today
+from matbench_discovery import ROOT, id_col, today
 from matbench_discovery.data import DATA_FILES, df_wbm
 from matbench_discovery.energy import get_e_form_per_atom
 from matbench_discovery.plots import plt
@@ -22,9 +22,7 @@
 """
 
 
-df_cse = pd.read_json(DATA_FILES.wbm_computed_structure_entries).set_index(
-    "material_id"
-)
+df_cse = pd.read_json(DATA_FILES.wbm_computed_structure_entries).set_index(id_col)
 
 cses = [
     ComputedStructureEntry.from_dict(dct)

data/wbm/eda.py

@@ -12,7 +12,7 @@
 )
 from pymatviz.io import save_fig
 
-from matbench_discovery import PDF_FIGS, ROOT, SITE_FIGS, STABILITY_THRESHOLD
+from matbench_discovery import PDF_FIGS, ROOT, SITE_FIGS, STABILITY_THRESHOLD, id_col
 from matbench_discovery import plots as plots
 from matbench_discovery.data import DATA_FILES, df_wbm
 from matbench_discovery.energy import mp_elem_reference_entries
@@ -180,14 +180,14 @@
 # many models struggle on the halogens in per-element error periodic table heatmaps
 # https://janosh.github.io/matbench-discovery/models
 df_2d_tsne = pd.read_csv(f"{module_dir}/tsne/one-hot-112-composition-2d.csv.gz")
-df_2d_tsne = df_2d_tsne.set_index("material_id")
+df_2d_tsne = df_2d_tsne.set_index(id_col)
 
 df_3d_tsne = pd.read_csv(f"{module_dir}/tsne/one-hot-112-composition-3d.csv.gz")
 model = "Wrenformer"
 df_3d_tsne = pd.read_csv(
     f"{module_dir}/tsne/one-hot-112-composition+{model}-each-err-3d-metric=eucl.csv.gz"
 )
-df_3d_tsne = df_3d_tsne.set_index("material_id")
+df_3d_tsne = df_3d_tsne.set_index(id_col)
 
 df_wbm[list(df_2d_tsne)] = df_2d_tsne
 df_wbm[list(df_3d_tsne)] = df_3d_tsne
@@ -205,7 +205,7 @@
     x="2d t-SNE 1",
     y="2d t-SNE 2",
     color=color_col,
-    hover_name="material_id",
+    hover_name=id_col,
     hover_data=("formula", each_true_col),
     range_color=(0, clr_range_max),
 )
@@ -219,7 +219,7 @@
     y="3d t-SNE 2",
     z="3d t-SNE 3",
     color=color_col,
-    custom_data=["material_id", "formula", each_true_col, color_col],
+    custom_data=[id_col, "formula", each_true_col, color_col],
     range_color=(0, clr_range_max),
 )
 fig.data[0].hovertemplate = (

data/wbm/fetch_process_wbm_dataset.py

@@ -18,7 +18,7 @@
 from pymatviz.io import save_fig
 from tqdm import tqdm
 
-from matbench_discovery import SITE_FIGS, today
+from matbench_discovery import SITE_FIGS, id_col, today
 from matbench_discovery.data import DATA_FILES
 from matbench_discovery.energy import get_e_form_per_atom
 from matbench_discovery.plots import pio
@@ -156,7 +156,7 @@ def increment_wbm_material_id(wbm_id: str) -> str:
 
 
 df_wbm.index = df_wbm.index.map(increment_wbm_material_id)
-df_wbm.index.name = "material_id"
+df_wbm.index.name = id_col
 assert df_wbm.index[0] == "wbm-1-1"
 assert df_wbm.index[-1] == "wbm-5-23308"
 
@@ -296,13 +296,13 @@ def increment_wbm_material_id(wbm_id: str) -> str:
     "e_form": "e_form_per_atom_wbm",
     "e_hull": "e_above_hull_wbm",
     "gap": "bandgap_pbe",
-    "id": "material_id",
+    "id": id_col,
 }
 # WBM summary was shared twice, once on google drive, once on materials cloud
 # download both and check for consistency
 df_summary = pd.read_csv(
     f"{module_dir}/raw/wbm-summary.txt", sep="\t", names=col_map.values()
-).set_index("material_id")
+).set_index(id_col)
 
 df_summary_bz2 = pd.read_csv(
     f"{mat_cloud_url}&filename=summary.txt.bz2", sep="\t"
@@ -618,7 +618,7 @@ def fix_bad_struct_index_mismatch(material_id: str) -> str:
 suggest = "not found, run scripts/compute_struct_fingerprints.py to generate"
 fp_diff_col = "site_stats_fingerprint_init_final_norm_diff"
 try:
-    df_fp = pd.read_json(fingerprints_path).set_index("material_id")
+    df_fp = pd.read_json(fingerprints_path).set_index(id_col)
     df_summary[fp_diff_col] = df_fp[fp_diff_col]
 except FileNotFoundError:
     print(f"{fingerprints_path=} {suggest}")
@@ -633,11 +633,11 @@ def fix_bad_struct_index_mismatch(material_id: str) -> str:
 # %% only here to load data for later inspection
 if False:
     df_summary = pd.read_csv(f"{module_dir}/2022-10-19-wbm-summary.csv.gz").set_index(
-        "material_id"
+        id_col
     )
     df_wbm = pd.read_json(
         f"{module_dir}/2022-10-19-wbm-computed-structure-entries+init-structs.json.bz2"
-    ).set_index("material_id")
+    ).set_index(id_col)
 
     df_wbm["cse"] = [
         ComputedStructureEntry.from_dict(dct)

matbench_discovery/__init__.py

@@ -7,6 +7,7 @@
 SITE_FIGS = f"{ROOT}/site/src/figs"  # directory for interactive figures
 SITE_MODELS = f"{ROOT}/site/src/routes/models"  # directory to write model analysis
 FIGSHARE = f"{ROOT}/data/figshare"
+SCRIPTS = f"{ROOT}/scripts"
 PDF_FIGS = f"{ROOT}/paper/figs"  # directory for light-themed PDF figures
 
 for directory in [SITE_FIGS, SITE_MODELS, FIGSHARE, PDF_FIGS]:
@@ -31,3 +32,8 @@
     warnings.filterwarnings(
         action="ignore", category=UserWarning, module="pymatgen", lineno=lineno
     )
+
+id_col = "material_id"
+init_struct_col = "initial_structure"
+struct_col = "structure"
+e_form_col = "formation_energy_per_atom"

matbench_discovery/data.py

@@ -16,7 +16,7 @@
 from pymatgen.analysis.phase_diagram import PatchedPhaseDiagram
 from tqdm import tqdm
 
-from matbench_discovery import FIGSHARE
+from matbench_discovery import FIGSHARE, id_col
 
 # repo URL to raw files on GitHub
 RAW_REPO_URL = "https://github.com/janosh/matbench-discovery/raw"
@@ -119,8 +119,8 @@ def load(
         print(f"\n\nvariable dump:\n{file=},\n{reader=}\n{kwargs=}")
         raise
 
-    if "material_id" in df:
-        df = df.set_index("material_id")
+    if id_col in df:
+        df = df.set_index(id_col)
     if hydrate:
         for col in df:
             if not isinstance(df[col].iloc[0], dict):
@@ -256,4 +256,4 @@ def _on_not_found(self, key: str, msg: str) -> None:  # type: ignore[override]
 
 
 df_wbm = load("wbm_summary")
-df_wbm["material_id"] = df_wbm.index
+df_wbm[id_col] = df_wbm.index

matbench_discovery/preds.py

@@ -7,6 +7,7 @@
 from tqdm import tqdm
 
 from matbench_discovery import ROOT, STABILITY_THRESHOLD
+from matbench_discovery import id_col as default_id_col
 from matbench_discovery.data import Files, df_wbm, glob_to_df
 from matbench_discovery.metrics import stable_metrics
 from matbench_discovery.plots import (
@@ -83,7 +84,7 @@ class PredFiles(Files):
 def load_df_wbm_with_preds(
     models: Sequence[str] = (*PRED_FILES,),
     pbar: bool = True,
-    id_col: str = "material_id",
+    id_col: str = default_id_col,
     **kwargs: Any,
 ) -> pd.DataFrame:
     """Load WBM summary dataframe with model predictions from disk.

models/alignn/test_alignn.py

@@ -17,7 +17,7 @@
 from sklearn.metrics import r2_score
 from tqdm import tqdm
 
-from matbench_discovery import today
+from matbench_discovery import id_col, today
 from matbench_discovery.data import DATA_FILES, df_wbm
 from matbench_discovery.plots import wandb_scatter
 from matbench_discovery.preds import e_form_col
@@ -36,7 +36,6 @@
 task_type = "IS2RE"
 target_col = e_form_col
 input_col = "initial_structure"
-id_col = "material_id"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 job_name = f"{model_name}-wbm-{task_type}"
 out_dir = os.getenv("SBATCH_OUTPUT", f"{module_dir}/{today}-{job_name}")

models/alignn/train_alignn.py

@@ -18,7 +18,7 @@
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
-from matbench_discovery import today
+from matbench_discovery import e_form_col, id_col, struct_col, today
 from matbench_discovery.data import DATA_FILES
 from matbench_discovery.slurm import slurm_submit
 
@@ -30,10 +30,8 @@
 
 # %%
 model_name = "alignn-mp-e_form"
-target_col = "formation_energy_per_atom"
-struct_col = "structure"
+target_col = e_form_col
 input_col = "atoms"
-id_col = "material_id"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 job_name = f"train-{model_name}"
 
@@ -48,7 +46,6 @@
 
 slurm_vars = slurm_submit(
     job_name=job_name,
-    # partition="perlmuttter",
     account="matgen",
     time="4:0:0",
     out_dir=out_dir,

models/alignn_ff/alignn_ff_relax.py

@@ -9,7 +9,7 @@
 from pymatgen.io.jarvis import JarvisAtomsAdaptor
 from tqdm import tqdm
 
-from matbench_discovery import today
+from matbench_discovery import id_col, init_struct_col, today
 from matbench_discovery.data import DATA_FILES, df_wbm
 from matbench_discovery.preds import e_form_col as target_col
 
@@ -29,8 +29,7 @@
 # model_name = "mp_e_form_alignn"  # pre-trained by NIST
 model_name = f"{out_dir}/best-model.pth"
 task_type = "IS2RE"
-input_col = "initial_structure"
-id_col = "material_id"
+input_col = init_struct_col
 job_name = f"{model_name}-wbm-{task_type}"
 out_path = (
     f"{out_dir}/{'alignn-relaxed-structs' if batch == 0 else f'{batch=}'}.json.gz"

models/alignn_ff/test_alignn_ff.py

@@ -18,7 +18,7 @@
 from sklearn.metrics import r2_score
 from tqdm import tqdm
 
-from matbench_discovery import today
+from matbench_discovery import init_struct_col, today
 from matbench_discovery.data import DATA_FILES, df_wbm
 from matbench_discovery.plots import wandb_scatter
 from matbench_discovery.preds import e_form_col as target_col
@@ -33,8 +33,7 @@
 n_splits = 100
 # model_name = "mp_e_form_alignnn"  # pre-trained by NIST
 task_type = "IS2RE"
-input_col = "initial_structure"
-id_col = "material_id"
+input_col = init_struct_col
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_name = f"alignn-ff-wbm-{task_type}"
 job_name = f"{model_name}-relaxed-wbm-{task_type}"

models/bowsr/join_bowsr_results.py

@@ -8,7 +8,7 @@
 import pymatviz
 from tqdm import tqdm
 
-from matbench_discovery.data import DATA_FILES
+from matbench_discovery.data import DATA_FILES, id_col
 
 __author__ = "Janosh Riebesell"
 __date__ = "2022-09-22"
@@ -30,14 +30,14 @@
 for file_path in tqdm(file_paths):
     if file_path in dfs:
         continue
-    dfs[file_path] = pd.read_json(file_path).set_index("material_id")
+    dfs[file_path] = pd.read_json(file_path).set_index(id_col)
 
 
 df_bowsr = pd.concat(dfs.values()).round(4)
 
 
 # %% compare against WBM formation energy targets to make sure we got sensible results
-df_wbm = pd.read_csv(DATA_FILES.wbm_summary).set_index("material_id")
+df_wbm = pd.read_csv(DATA_FILES.wbm_summary).set_index(id_col)
 
 
 print(
@@ -75,4 +75,4 @@
 
 
 # in_path = f"{module_dir}/2023-01-23-bowsr-megnet-wbm-IS2RE.json.gz"
-# df_bowsr = pd.read_json(in_path).set_index("material_id")
+# df_bowsr = pd.read_json(in_path).set_index(id_col)