compare IS2RE and RS2RE energies in eda_wbm_pre_vs_post_m3gnet_relaxation.py

add figures/2022-08-22-m3gnet-energy-per-atom-scatter-is2re-vs-rs2re.png

Author: janosh

mb_discovery/m3gnet/eda_wbm_pre_vs_post_m3gnet_relaxation.py

@@ -7,6 +7,7 @@
 from pymatgen.core import Structure
 from pymatgen.util.coord import pbc_diff
 from pymatviz.utils import add_identity_line
+from sklearn.metrics import r2_score
 
 from mb_discovery import ROOT
 
@@ -27,37 +28,39 @@
 
 
 # %%
-df_m3gnet = pd.read_json(
-    f"{ROOT}/data/2022-08-16-m3gnet-wbm-relax-results.json.gz"
+df_m3gnet_is2re = pd.read_json(
+    f"{ROOT}/data/2022-08-16-m3gnet-wbm-relax-results-IS2RE.json.gz"
+).set_index("material_id")
+df_m3gnet_rs2re = pd.read_json(
+    f"{ROOT}/data/2022-08-19-m3gnet-wbm-relax-results-RS3RE.json.gz"
 ).set_index("material_id")
-
-print("Number of WBM crystals for which we have M3GNet results:")
-print(f"{len(df_m3gnet):,} / {len(df_wbm):,} = {len(df_m3gnet)/len(df_wbm):.1%}")
 
 
 # %% spread M3GNet post-pseudo-relaxation lattice params into separate columns
-df_m3gnet["final_energy"] = df_m3gnet.trajectory.map(lambda x: x["energies"][-1][0])
+df_m3gnet_is2re["final_energy"] = df_m3gnet_is2re.trajectory.map(
+    lambda x: x["energies"][-1][0]
+)
 
 df_m3gnet_lattice = pd.json_normalize(
-    df_m3gnet.initial_structure.map(lambda x: x["lattice"])
+    df_m3gnet_is2re.initial_structure.map(lambda x: x["lattice"])
 ).add_prefix("m3gnet_")
-df_m3gnet[df_m3gnet_lattice.columns] = df_m3gnet_lattice.to_numpy()
-df_m3gnet
+df_m3gnet_is2re[df_m3gnet_lattice.columns] = df_m3gnet_lattice.to_numpy()
+df_m3gnet_is2re
 
 
 # %% spread WBM initial and final lattice params into separate columns
-df_m3gnet["final_wbm_structure"] = df_wbm.cse.map(lambda x: x["structure"])
+df_m3gnet_is2re["final_wbm_structure"] = df_wbm.cse.map(lambda x: x["structure"])
 df_wbm_final_lattice = pd.json_normalize(
-    df_m3gnet.final_wbm_structure.map(lambda x: x["lattice"])
+    df_m3gnet_is2re.final_wbm_structure.map(lambda x: x["lattice"])
 ).add_prefix("final_wbm_")
-df_m3gnet[df_wbm_final_lattice.columns] = df_wbm_final_lattice.to_numpy()
+df_m3gnet_is2re[df_wbm_final_lattice.columns] = df_wbm_final_lattice.to_numpy()
 
 
-df_m3gnet["initial_wbm_structure"] = df_wbm.initial_structure
+df_m3gnet_is2re["initial_wbm_structure"] = df_wbm.initial_structure
 df_wbm_initial_lattice = pd.json_normalize(
-    df_m3gnet.initial_structure.map(lambda x: x["lattice"])
+    df_m3gnet_is2re.initial_structure.map(lambda x: x["lattice"])
 ).add_prefix("initial_wbm_")
-df_m3gnet[df_wbm_initial_lattice.columns] = df_wbm_initial_lattice.to_numpy()
+df_m3gnet_is2re[df_wbm_initial_lattice.columns] = df_wbm_initial_lattice.to_numpy()
 
 
 # %%
@@ -78,7 +81,7 @@
 
 # %%
 px.histogram(
-    df_m3gnet.filter(like="volume"),
+    df_m3gnet_is2re.filter(like="volume"),
     nbins=500,
     barmode="overlay",
     opacity=0.5,
@@ -88,10 +91,10 @@
 
 # %%
 fig = px.scatter(
-    df_m3gnet.round(1),
+    df_m3gnet_is2re.round(1),
     x="final_wbm_volume",
     y=["initial_wbm_volume", "m3gnet_volume"],
-    hover_data=[df_m3gnet.index],
+    hover_data=[df_m3gnet_is2re.index],
 )
 add_identity_line(fig)
 fig.update_layout(
@@ -102,60 +105,63 @@
 
 # %% histogram of alpha lattice angles (similar results for beta and gamma)
 fig = px.histogram(
-    df_m3gnet.filter(like="alpha"), nbins=1000, barmode="overlay", log_y=True
+    df_m3gnet_is2re.filter(like="alpha"), nbins=1000, barmode="overlay", log_y=True
 )
 fig.show()
 
 
 # %%
 px.histogram(
-    df_m3gnet.filter(regex="_c$"),
+    df_m3gnet_is2re.filter(regex="_c$"),
     nbins=1000,
     log_y=True,
     barmode="overlay",
     opacity=0.5,
 )
 
 
-# %%
-df_m3gnet["final_m3gnet_structure"] = df_m3gnet.final_structure.map(Structure.from_dict)
-df_m3gnet["initial_wbm_structure"] = df_m3gnet.initial_wbm_structure.map(
+# %% compute mean absolute PBC difference between initial and final fractional
+# coordinates of crystal sites
+df_m3gnet_is2re["m3gnet_structure"] = df_m3gnet_is2re.m3gnet_structure.map(
+    Structure.from_dict
+)
+df_m3gnet_is2re["initial_wbm_structure"] = df_m3gnet_is2re.initial_wbm_structure.map(
     Structure.from_dict
 )
-df_m3gnet["final_wbm_structure"] = df_m3gnet.final_wbm_structure.map(
+df_m3gnet_is2re["final_wbm_structure"] = df_m3gnet_is2re.final_wbm_structure.map(
     Structure.from_dict
 )
 
 
-df_m3gnet["m3gnet_pbc_diffs"] = [
+df_m3gnet_is2re["m3gnet_pbc_diffs"] = [
     abs(
         pbc_diff(
             row.initial_wbm_structure.frac_coords,
-            row.final_m3gnet_structure.frac_coords,
+            row.m3gnet_structure.frac_coords,
         )
     ).mean()
-    for row in df_m3gnet.itertuples()
+    for row in df_m3gnet_is2re.itertuples()
 ]
 
 
-df_m3gnet["wbm_pbc_diffs"] = [
+df_m3gnet_is2re["wbm_pbc_diffs"] = [
     abs(
         pbc_diff(
             row.initial_wbm_structure.frac_coords,
             row.final_wbm_structure.frac_coords,
         )
     ).mean()
-    for row in df_m3gnet.itertuples()
+    for row in df_m3gnet_is2re.itertuples()
 ]
 
-df_m3gnet["m3gnet_to_final_wbm_pbc_diffs"] = [
+df_m3gnet_is2re["m3gnet_to_final_wbm_pbc_diffs"] = [
     abs(
         pbc_diff(
-            row.final_m3gnet_structure.frac_coords,
+            row.m3gnet_structure.frac_coords,
             row.final_wbm_structure.frac_coords,
         )
     ).mean()
-    for row in df_m3gnet.itertuples()
+    for row in df_m3gnet_is2re.itertuples()
 ]
 
 
@@ -164,13 +170,60 @@
     "and M3GNet"
 )
 
-wbm_pbc_diffs_mean = df_m3gnet.wbm_pbc_diffs.mean()
+wbm_pbc_diffs_mean = df_m3gnet_is2re.wbm_pbc_diffs.mean()
 print(f"{wbm_pbc_diffs_mean = :.3}")
 
-m3gnet_pbc_diffs_mean = df_m3gnet.m3gnet_pbc_diffs.mean()
+m3gnet_pbc_diffs_mean = df_m3gnet_is2re.m3gnet_pbc_diffs.mean()
 print(f"{m3gnet_pbc_diffs_mean = :.3}")
 
-m3gnet_to_final_wbm_pbc_diffs_mean = df_m3gnet.m3gnet_to_final_wbm_pbc_diffs.mean()
+m3gnet_to_final_wbm_pbc_diffs_mean = (
+    df_m3gnet_is2re.m3gnet_to_final_wbm_pbc_diffs.mean()
+)
 print(f"{m3gnet_to_final_wbm_pbc_diffs_mean = :.3}")
 
 print(f"{wbm_pbc_diffs_mean / m3gnet_pbc_diffs_mean = :.3}")
+
+
+# %%
+# plt_fig = df_m3gnet_is2re.plot.scatter(
+#     x="e_m3gnet_per_atom_rs2re", y="e_m3gnet_per_atom_is2re"
+# )
+# df_m3gnet_is2re.filter(like="m3gnet_energy").hist(bins=100)
+
+df_m3gnet_is2re["m3gnet_energy_rs2re"] = df_m3gnet_rs2re.m3gnet_energy
+
+for task_type in ["is2re", "rs2re"]:
+    e_per_atom = df_m3gnet_is2re[f"m3gnet_energy_{task_type}"] / df_m3gnet_is2re.n_sites
+
+    df_m3gnet_is2re[f"e_m3gnet_per_atom_{task_type}"] = e_per_atom
+
+fig = px.scatter(
+    df_m3gnet_is2re,
+    x="e_m3gnet_per_atom_rs2re",
+    y="e_m3gnet_per_atom_is2re",
+    render_mode="webgl",
+)
+add_identity_line(fig)
+
+len_overlap = df_m3gnet_is2re.filter(like="e_m3gnet_per_atom_").dropna().shape[0]
+x_vals, y_vals = df_m3gnet_is2re.filter(like="e_m3gnet_per_atom_").dropna().values.T
+
+MAE = abs(x_vals - y_vals).mean()
+R2 = r2_score(x_vals, y_vals)
+
+title = f"data size = {len_overlap:,} \t {MAE = :.2} \t {R2 = :.4}"
+fig.update_layout(title=dict(text=title, x=0.5))
+
+# 250k scatter points require exporting to PNG, interactive version freezes the
+# notebook server
+fig.show(renderer="png", scale=2)
+fig.write_image(
+    f"{ROOT}/figures/{today}-m3gnet-energy-per-atom-scatter-is2re-vs-rs2re.png", scale=2
+)
+
+
+# %% write df back to compressed JSON
+# filter out columns containing 'rs2re'
+# df_m3gnet_is2re.reset_index().filter(regex="^((?!rs2re).)*$").to_json(
+#     f"{ROOT}/data/2022-08-16-m3gnet-wbm-relax-results-IS2RE-2.json.gz"
+# ).set_index("material_id")

mb_discovery/m3gnet/join_and_plot_m3gnet_relax_results.py

@@ -23,8 +23,9 @@
 
 
 # %%
-# glob_pattern = "2022-08-16-m3gnet-relax-wbm-IS3RE/*.json.gz"
-glob_pattern = "2022-08-19-m3gnet-relax-wbm-RS3RE/*.json.gz"
+task_type = "RS3RE"
+date = "2022-08-19"
+glob_pattern = f"{date}-m3gnet-relax-wbm-{task_type}/*.json.gz"
 file_paths = sorted(glob(f"{ROOT}/data/{glob_pattern}"))
 print(f"Found {len(file_paths):,} files for {glob_pattern = }")
 
@@ -78,7 +79,7 @@
     PDEntry(row.m3gnet_structure.composition, row.m3gnet_energy)
     for row in df_m3gnet.itertuples()
 ]
-df_m3gnet["e_form_m3gnet"] = [
+df_m3gnet["e_form_m3gnet_from_ppd"] = [
     ppd_mp_wbm.get_form_energy_per_atom(x) for x in pd_entries_m3gnet
 ]
 
@@ -113,14 +114,14 @@
 
 
 # %%
-out_path = f"{ROOT}/data/{today}-m3gnet-wbm-relax-results.json.gz"
+out_path = f"{ROOT}/data/{today}-m3gnet-relax-wbm-{task_type}.json.gz"
 df_m3gnet.reset_index().to_json(out_path, default_handler=as_dict_handler)
 
 
 # %%
 ax_hull_dist_hist = hist_classify_stable_as_func_of_hull_dist(
-    formation_energy_targets=df_m3gnet.e_form_ppd,
-    formation_energy_preds=df_m3gnet.e_form_m3gnet,
+    formation_energy_targets=df_m3gnet.e_form_ppd_2022_01_25,
+    formation_energy_preds=df_m3gnet.e_form_m3gnet_from_ppd,
     e_above_hull_vals=df_m3gnet.e_above_mp_hull,
 )