add models/chgnet/ctk_struct_traj.py and models/wrenformer/analyze_wrenformer.py

Author: janosh

.pre-commit-config.yaml

@@ -7,7 +7,7 @@ default_install_hook_types: [pre-commit, commit-msg]
 
 repos:
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: v0.0.257
+    rev: v0.0.258
     hooks:
       - id: ruff
         args: [--fix]

matbench_discovery/preds.py

@@ -32,22 +32,22 @@ class PredFiles(Files):
     # bowsr optimizer coupled with original megnet
     bowsr_megnet = "bowsr/2023-01-23-bowsr-megnet-wbm-IS2RE.csv"
     # default CHGNet model from publication with 400,438 params
-    chgnet = CHGNet = "chgnet/2023-03-06-chgnet-wbm-IS2RE.csv"
+    chgnet = "chgnet/2023-03-06-chgnet-wbm-IS2RE.csv"
     chgnet_megnet = "chgnet/2023-03-04-chgnet-wbm-IS2RE.csv"
     # CGCnn 10-member ensemble
-    cgcnn = cgcnn = "cgcnn/2023-01-26-test-cgcnn-wbm-IS2RE/cgcnn-ensemble-preds.csv"
+    cgcnn = "cgcnn/2023-01-26-test-cgcnn-wbm-IS2RE/cgcnn-ensemble-preds.csv"
     # cgcnn 10-member ensemble with 5-fold training set perturbations
-    cgcnn_p = CGCNN_P = "cgcnn/2023-02-05-cgcnn-perturb=5.csv"
+    cgcnn_p = "cgcnn/2023-02-05-cgcnn-perturb=5.csv"
     # original m3gnet straight from publication, not re-trained
-    m3gnet = M3GNet = "m3gnet/2022-10-31-m3gnet-wbm-IS2RE.csv"
+    m3gnet = "m3gnet/2022-10-31-m3gnet-wbm-IS2RE.csv"
     # m3gnet-relaxed structures fed into megnet for formation energy prediction
     m3gnet_megnet = "m3gnet/2022-10-31-m3gnet-wbm-IS2RE.csv"
     # original megnet straight from publication, not re-trained
-    megnet = MEGNet = "megnet/2022-11-18-megnet-wbm-IS2RE/megnet-e-form-preds.csv"
+    megnet = "megnet/2022-11-18-megnet-wbm-IS2RE/megnet-e-form-preds.csv"
     # magpie composition+voronoi tessellation structure features + sklearn random forest
-    voronoi_rf = Voronoi_RF = "voronoi/2022-11-27-train-test/e-form-preds-IS2RE.csv"
+    voronoi_rf = "voronoi/2022-11-27-train-test/e-form-preds-IS2RE.csv"
     # wrenformer 10-member ensemble
-    wrenformer = Wrenformer = "wrenformer/2022-11-15-wrenformer-IS2RE-preds.csv"
+    wrenformer = "wrenformer/2022-11-15-wrenformer-IS2RE-preds.csv"
 
 
 PRED_FILES = PredFiles()

models/chgnet/analyze_chgnet.py

@@ -9,55 +9,78 @@
 from pymatgen.core import Structure
 from pymatviz import density_scatter, plot_structure_2d, ptable_heatmap_plotly
 
-from matbench_discovery import plots
+from matbench_discovery import plots as plots
 from matbench_discovery.data import DATA_FILES, df_wbm
 from matbench_discovery.preds import PRED_FILES
 
 __author__ = "Janosh Riebesell"
 __date__ = "2023-03-06"
 
 module_dir = os.path.dirname(__file__)
-del plots  # https://github.com/PyCQA/pyflakes/issues/366
+id_col = "material_id"
 
 
 # %%
-df_chgnet = pd.read_csv(PRED_FILES.CHGNet)
-df_chgnet = df_chgnet.set_index("material_id").add_suffix("_2000")
-df_chgnet_500 = pd.read_csv(PRED_FILES.CHGNet.replace("-06", "-04"))
-df_chgnet_500 = df_chgnet_500.set_index("material_id").add_suffix("_500")
+df_chgnet = pd.read_csv(PRED_FILES.__dict__["CHGNet"])
+df_chgnet = df_chgnet.set_index(id_col).add_suffix("_2000")
+df_chgnet_500 = pd.read_csv(PRED_FILES.__dict__["CHGNet"].replace("-06", "-04"))
+df_chgnet_500 = df_chgnet_500.set_index(id_col).add_suffix("_500")
 df_chgnet[list(df_chgnet_500)] = df_chgnet_500
 df_chgnet["formula"] = df_wbm.formula
 
 e_form_2000 = "e_form_per_atom_chgnet_2000"
 e_form_500 = "e_form_per_atom_chgnet_500"
 
-min_e_diff = 0.35
+min_e_diff = 0.1
+# structures with smaller energy after longer relaxation need many steps
+df_long = df_chgnet.query(f"{e_form_2000} - {e_form_500} < -{min_e_diff}")
+# structures with larger energy after longer relaxation are problematic
 df_bad = df_chgnet.query(f"{e_form_2000} - {e_form_500} > {min_e_diff}")
+# both combined
+df_diff = df_chgnet.query(f"abs({e_form_2000} - {e_form_500}) > {min_e_diff}")
+
+assert len(df_long) + len(df_bad) == len(df_diff)
+
+
+# %%
+density_scatter(df=df_chgnet, x=e_form_500, y=e_form_2000)
 
 
 # %%
-density_scatter(df=df_chgnet, x=e_form_2000, y=e_form_500)
+df_diff.reset_index().plot.scatter(
+    x=e_form_500,
+    y=e_form_2000,
+    hover_name=id_col,
+    hover_data=["formula"],
+    backend="plotly",
+    title=f"{len(df_diff)} structures have > {min_e_diff} eV/atom energy diff after "
+    "longer relaxation",
+)
 
 
 # %%
 fig = ptable_heatmap_plotly(df_bad.formula)
-title = "structures with larger error after longer relaxation"
-fig.layout.title.update(text=f"{len(df_bad)} {title}")
+title = "structures with larger error<br>after longer relaxation"
+fig.layout.title.update(text=f"{len(df_diff)} {title}", x=0.4, y=0.9)
+fig.show()
 
 
 # %%
-df_cse = pd.read_json(DATA_FILES.wbm_initial_structures).set_index("material_id")
+df_cse = pd.read_json(DATA_FILES.wbm_cses_plus_init_structs).set_index(id_col)
+df_cse.loc[df_diff.index].reset_index().to_json(
+    f"{module_dir}/wbm-chgnet-bad-relax.json.gz"
+)
 
 
 # %%
 n_rows, n_cols = 3, 4
 fig, axs = plt.subplots(n_rows, n_cols, figsize=(3 * n_cols, 4 * n_rows))
-n_struct = min(n_rows * n_cols, len(df_bad))
+n_struct = min(n_rows * n_cols, len(df_diff))
 struct_col = "initial_structure"
 
 fig.suptitle(f"{n_struct} {struct_col} {title}", fontsize=16, fontweight="bold", y=1.05)
 for idx, (ax, row) in enumerate(
-    zip(axs.flat, df_cse.loc[df_bad.index].itertuples()), 1
+    zip(axs.flat, df_cse.loc[df_diff.index].itertuples()), 1
 ):
     struct = Structure.from_dict(getattr(row, struct_col))
     plot_structure_2d(struct, ax=ax)

scripts/ctk_structure_viewer.py models/chgnet/ctk_structure_viewer.py

@@ -3,6 +3,8 @@
 import pandas as pd
 from crystal_toolkit.helpers.utils import hook_up_fig_with_struct_viewer
 
+from matbench_discovery.preds import PRED_FILES
+
 __author__ = "Janosh Riebesell"
 __date__ = "2023-03-07"
 
@@ -13,28 +15,24 @@
 Then open http://localhost:8000 in your browser.
 """
 
-df_plot = None
-min_e_diff = 0.1
 e_form_2000 = "e_form_per_atom_chgnet_2000"
 e_form_500 = "e_form_per_atom_chgnet_500"
 
-if df_plot is None:
-    from matbench_discovery.preds import PRED_FILES
+df_chgnet = pd.read_json(PRED_FILES.__dict__["CHGNet"].replace(".csv", ".json.gz"))
+df_chgnet = df_chgnet.set_index("material_id")
 
-    df_chgnet = pd.read_json(PRED_FILES.CHGNet.replace(".csv", ".json.gz"))
-    df_chgnet = df_chgnet.set_index("material_id")
+df_chgnet_2000 = pd.read_csv(PRED_FILES.__dict__["CHGNet"])
+df_chgnet_2000 = df_chgnet_2000.set_index("material_id").add_suffix("_2000")
+df_chgnet[list(df_chgnet_2000)] = df_chgnet_2000
 
-    df_chgnet_2000 = pd.read_csv(PRED_FILES.CHGNet)
-    df_chgnet_2000 = df_chgnet_2000.set_index("material_id").add_suffix("_2000")
-    df_chgnet[list(df_chgnet_2000)] = df_chgnet_2000
+df_chgnet_500 = pd.read_csv(PRED_FILES.__dict__["CHGNet"].replace("-06", "-04"))
+df_chgnet_500 = df_chgnet_500.set_index("material_id").add_suffix("_500")
+df_chgnet[list(df_chgnet_500)] = df_chgnet_500
 
-    df_chgnet_500 = pd.read_csv(PRED_FILES.CHGNet.replace("-06", "-04"))
-    df_chgnet_500 = df_chgnet_500.set_index("material_id").add_suffix("_500")
-    df_chgnet[list(df_chgnet_500)] = df_chgnet_500
-
-    e_form_abs_diff = "e_form_abs_diff"
-    df_chgnet[e_form_abs_diff] = abs(df_chgnet[e_form_2000] - df_chgnet[e_form_500])
-    df_plot = df_chgnet.round(3).query(f"{e_form_abs_diff} > {min_e_diff}")
+e_form_abs_diff = "e_form_abs_diff"
+min_e_diff = 0.1
+df_chgnet[e_form_abs_diff] = abs(df_chgnet[e_form_2000] - df_chgnet[e_form_500])
+df_plot = df_chgnet.round(3).query(f"{e_form_abs_diff} > {min_e_diff}")
 
 
 plot_labels = {
@@ -69,4 +67,4 @@
     # validate_id requires material_id to be hover_name
     validate_id=lambda id: id.startswith(("wbm-", "mp-", "mvc-")),
 )
-app.run_server(debug=True, port=8000)
+app.run(debug=True, port=8000)