move model run/eval scripts from mb_discovery/ to new dir models/

mv mb_discovery/{plot_scripts/plot_funcs.py -> plots.py}

Author: janosh

.gitignore

@@ -20,7 +20,7 @@ job-logs/
 
 # slurm logs
 slurm-*out
-mb_discovery/**/*.csv
+models/**/*.csv
 
 # temporary ignore rule
 paper

mb_discovery/__init__.py

@@ -6,8 +6,6 @@
 PKG_DIR = os.path.dirname(__file__)
 ROOT = os.path.dirname(PKG_DIR)
 
-os.makedirs(f"{PKG_DIR}/plots", exist_ok=True)
-
 
 def chunks(xs: Sequence[Any], n: int) -> Generator[Sequence[Any], None, None]:
     return (xs[i : i + n] for i in range(0, len(xs), n))

mb_discovery/plot_scripts/hist_classified_stable_as_func_of_hull_dist.py

@@ -4,11 +4,11 @@
 import pandas as pd
 
 from mb_discovery import ROOT
-from mb_discovery.plot_scripts import plt
-from mb_discovery.plot_scripts.plot_funcs import (
+from mb_discovery.plots import (
     StabilityCriterion,
     WhichEnergy,
     hist_classified_stable_as_func_of_hull_dist,
+    plt,
 )
 
 __author__ = "Rhys Goodall, Janosh Riebesell"

mb_discovery/plot_scripts/hist_classified_stable_as_func_of_hull_dist_batches.py

@@ -4,11 +4,11 @@
 import pandas as pd
 
 from mb_discovery import ROOT
-from mb_discovery.plot_scripts import plt
-from mb_discovery.plot_scripts.plot_funcs import (
+from mb_discovery.plots import (
     StabilityCriterion,
     WhichEnergy,
     hist_classified_stable_as_func_of_hull_dist,
+    plt,
 )
 
 __author__ = "Rhys Goodall, Janosh Riebesell"

mb_discovery/plot_scripts/precision_recall_vs_calc_count.py

@@ -4,11 +4,7 @@
 import pandas as pd
 
 from mb_discovery import ROOT
-from mb_discovery.plot_scripts import plt
-from mb_discovery.plot_scripts.plot_funcs import (
-    StabilityCriterion,
-    precision_recall_vs_calc_count,
-)
+from mb_discovery.plots import StabilityCriterion, plt, precision_recall_vs_calc_count
 
 __author__ = "Rhys Goodall, Janosh Riebesell"
 __date__ = "2022-06-18"
@@ -32,20 +28,22 @@
 ).set_index("material_id")
 
 dfs["Wrenformer"] = pd.read_csv(
-    f"{ROOT}/data/2022-08-16-wrenformer-preds.csv.bz2"
+    f"{ROOT}/models/wrenformer/mp/"
+    "2022-09-20-wrenformer-e_form-ensemble-1-preds-e_form_per_atom.csv"
 ).set_index("material_id")
 
+print(f"loaded models: {list(dfs)}")
+
 
 # %% download wbm-steps-summary.csv (23.31 MB)
-df_summary = pd.read_csv(
+df_wbm = pd.read_csv(
     "https://figshare.com/files/37542841?private_link=ff0ad14505f9624f0c05"
 ).set_index("material_id")
 
 
+# %%
 stability_crit: StabilityCriterion = "energy"
 
-
-# %%
 for (model_name, df), color in zip(
     dfs.items(),
     ("tab:blue", "tab:orange", "teal", "tab:pink", "black", "red", "turquoise"),

mb_discovery/plot_scripts/rolling_mae_vs_hull_dist.py

@@ -4,8 +4,7 @@
 import pandas as pd
 
 from mb_discovery import ROOT
-from mb_discovery.plot_scripts import plt
-from mb_discovery.plot_scripts.plot_funcs import rolling_mae_vs_hull_dist
+from mb_discovery.plots import plt, rolling_mae_vs_hull_dist
 
 __author__ = "Rhys Goodall, Janosh Riebesell"
 __date__ = "2022-06-18"
@@ -17,8 +16,9 @@
 markers = ["o", "v", "^", "H", "D", ""]
 
 data_path = (
-    f"{ROOT}/data/2022-06-11-from-rhys/wren-mp-initial-structures.csv"
-    # f"{ROOT}/data/2022-08-16-wrenformer-preds.csv.bz2"
+    # f"{ROOT}/data/2022-06-11-from-rhys/wren-mp-initial-structures.csv"
+    f"{ROOT}/models/wrenformer/mp/"
+    "2022-09-20-wrenformer-e_form-ensemble-1-preds-e_form_per_atom.csv"
 )
 df = pd.read_csv(data_path).set_index("material_id")
 legend_label = "Wren"

mb_discovery/plot_scripts/rolling_mae_vs_hull_dist_wbm_batches.py

@@ -4,8 +4,7 @@
 import pandas as pd
 
 from mb_discovery import ROOT
-from mb_discovery.plot_scripts import plt
-from mb_discovery.plot_scripts.plot_funcs import rolling_mae_vs_hull_dist
+from mb_discovery.plots import plt, rolling_mae_vs_hull_dist
 
 __author__ = "Rhys Goodall, Janosh Riebesell"
 __date__ = "2022-06-18"

mb_discovery/plot_scripts/plot_funcs.py mb_discovery/plots.py

@@ -2,21 +2,52 @@
 
 from typing import Any, Literal, Sequence, get_args
 
+import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+import plotly.express as px
+import plotly.io as pio
 import scipy.interpolate
 import scipy.stats
 from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar
 
-from mb_discovery.plot_scripts import plt
-
 __author__ = "Janosh Riebesell"
 __date__ = "2022-08-05"
 
 StabilityCriterion = Literal["energy", "energy+std", "energy-std"]
 WhichEnergy = Literal["true", "pred"]
 
 
+# --- define global plot settings
+px.defaults.labels = {
+    "n_atoms": "Atom Count",
+    "n_elems": "Element Count",
+    "crystal_sys": "Crystal system",
+    "spg_num": "Space group",
+    "n_wyckoff": "Number of Wyckoff positions",
+    "n_sites": "Lattice site count",
+    "energy_per_atom": "Energy (eV/atom)",
+    "e_form": "Formation energy (eV/atom)",
+    "e_above_hull": "Energy above convex hull (eV/atom)",
+    "e_above_hull_pred": "Predicted energy above convex hull (eV/atom)",
+    "e_above_mp_hull": "Energy above MP convex hull (eV/atom)",
+    "e_above_hull_error": "Error in energy above convex hull (eV/atom)",
+}
+
+pio.templates.default = "plotly_white"
+
+# https://github.com/plotly/Kaleido/issues/122#issuecomment-994906924
+# when seeing MathJax "loading" message in exported PDFs, try:
+# pio.kaleido.scope.mathjax = None
+
+
+plt.rc("font", size=14)
+plt.rc("savefig", bbox="tight", dpi=200)
+plt.rc("figure", dpi=200, titlesize=16)
+plt.rcParams["figure.constrained_layout.use"] = True
+# --- end global plot settings
+
+
 def hist_classified_stable_as_func_of_hull_dist(
     e_above_hull_pred: pd.Series,
     e_above_hull_true: pd.Series,

mb_discovery/m3gnet/eda_wbm_pre_vs_post_m3gnet_relaxation.py models/m3gnet/eda_wbm_pre_vs_post_m3gnet_relaxation.py

@@ -14,9 +14,7 @@
 from tqdm import tqdm
 
 from mb_discovery import ROOT, as_dict_handler
-from mb_discovery.plot_scripts.plot_funcs import (
-    hist_classified_stable_as_func_of_hull_dist,
-)
+from mb_discovery.plots import hist_classified_stable_as_func_of_hull_dist
 
 today = f"{datetime.now():%Y-%m-%d}"
 

mb_discovery/m3gnet/join_and_plot_m3gnet_relax_results.py models/m3gnet/join_and_plot_m3gnet_relax_results.py

@@ -20,8 +20,8 @@
 ```sh
 sbatch --partition icelake-himem --account LEE-SL3-CPU --array 1-101 \
     --time 3:0:0 --job-name m3gnet-wbm-relax-RS2RE --mem 12000 \
-    --output mb_discovery/m3gnet/slurm_logs/slurm-%A-%a.out \
-    --wrap "python mb_discovery/m3gnet/slurm_array_m3gnet_relax_wbm.py"
+    --output models/m3gnet/slurm_logs/slurm-%A-%a.out \
+    --wrap "python models/m3gnet/slurm_array_m3gnet_relax_wbm.py"
 ```
 
 --time 2h is probably enough but missing indices are annoying so best be safe.

mb_discovery/m3gnet/slurm_array_m3gnet_relax_wbm.py models/m3gnet/slurm_array_m3gnet_relax_wbm.py

@@ -2,12 +2,13 @@
 from datetime import datetime
 
 import pandas as pd
-from aviary import ROOT
 from aviary.utils import as_dict_handler
 from aviary.wren.utils import get_aflow_label_from_spglib
 from mp_api.client import MPRester
 from tqdm import tqdm
 
+from mb_discovery import ROOT
+
 """
 Download all MP formation and above hull energies on 2022-08-13.
 
@@ -48,7 +49,7 @@
 df["wyckoff"] = [get_aflow_label_from_spglib(x) for x in tqdm(df.structure)]
 
 df.to_json(
-    f"{ROOT}/datasets/{today}-mp-all-energies.json.gz", default_handler=as_dict_handler
+    f"{ROOT}/data/{today}-mp-all-energies.json.gz", default_handler=as_dict_handler
 )
 
-# df = pd.read_json(f"{ROOT}/datasets/2022-08-13-mp-all-energies.json.gz")
+# df = pd.read_json(f"{ROOT}/data/2022-08-13-mp-all-energies.json.gz")

mb_discovery/wrenformer/mp/get_mp_energies.py models/wrenformer/mp/get_mp_energies.py

@@ -7,7 +7,7 @@
 import pytest
 
 from mb_discovery import ROOT
-from mb_discovery.plot_scripts.plot_funcs import (
+from mb_discovery.plots import (
     StabilityCriterion,
     hist_classified_stable_as_func_of_hull_dist,
     precision_recall_vs_calc_count,