|
| 1 | +# %% |
| 2 | +import numpy as np |
| 3 | +import plotly.graph_objects as go |
| 4 | +from pymatviz.utils import add_identity_line, save_fig |
| 5 | +from sklearn.metrics import r2_score |
| 6 | + |
| 7 | +from matbench_discovery import FIGS, PAPER, today |
| 8 | +from matbench_discovery.data import PRED_FILENAMES, load_df_wbm_with_preds |
| 9 | +from matbench_discovery.energy import classify_stable |
| 10 | +from matbench_discovery.plots import px |
| 11 | + |
| 12 | +__author__ = "Janosh Riebesell" |
| 13 | +__date__ = "2022-11-28" |
| 14 | + |
| 15 | + |
| 16 | +# %% |
| 17 | +print(f"loadable models: {list(PRED_FILENAMES)}") |
| 18 | +models = sorted( |
| 19 | + "CGCNN, Voronoi RF, Wrenformer, MEGNet, M3GNet, BOWSR MEGNet".split(", ") |
| 20 | +) |
| 21 | +df_wbm = load_df_wbm_with_preds(models=models).round(3) |
| 22 | + |
| 23 | +e_form_col = "e_form_per_atom_mp2020_corrected" |
| 24 | +e_above_hull_col = "e_above_hull_mp2020_corrected_ppd_mp" |
| 25 | +id_col = "material_id" |
| 26 | +legend = dict(x=1, y=0, xanchor="right", yanchor="bottom", title=None) |
| 27 | + |
| 28 | + |
| 29 | +# %% |
| 30 | +e_form_preds = "e_form_per_atom_pred" |
| 31 | +e_above_hull_preds = "e_above_hull_pred" |
| 32 | +var_name = "Model" |
| 33 | +hover_cols = (id_col, e_form_col, e_above_hull_col, "formula") |
| 34 | + |
| 35 | +df_melt = df_wbm.melt( |
| 36 | + id_vars=hover_cols, |
| 37 | + value_vars=models, |
| 38 | + var_name=var_name, |
| 39 | + value_name=e_form_preds, |
| 40 | +) |
| 41 | + |
| 42 | +df_melt[e_above_hull_preds] = ( |
| 43 | + df_melt[e_above_hull_col] - df_melt[e_form_col] + df_melt[e_form_preds] |
| 44 | +) |
| 45 | + |
| 46 | + |
| 47 | +# %% |
| 48 | +def _metric_str(model_name: str) -> str: |
| 49 | + MAE = (df_wbm[e_form_col] - df_wbm[model_name]).abs().mean() |
| 50 | + R2 = r2_score(*df_wbm[[e_form_col, model_name]].dropna().to_numpy().T) |
| 51 | + return f"{model_name} · {MAE=:.2} · R<sup>2</sup>={R2:.2}" |
| 52 | + |
| 53 | + |
| 54 | +def _add_metrics_to_legend(fig: go.Figure) -> None: |
| 55 | + for trace in fig.data: |
| 56 | + # initially hide all traces, let users select which models to compare |
| 57 | + trace.visible = "legendonly" |
| 58 | + # add MAE and R2 to legend |
| 59 | + model = trace.name |
| 60 | + trace.name = _metric_str(model) |
| 61 | + |
| 62 | + |
| 63 | +# %% scatter plot of actual vs predicted e_form_per_atom |
| 64 | +fig = px.scatter( |
| 65 | + df_melt.iloc[::10], |
| 66 | + x=e_form_col, |
| 67 | + y=e_form_preds, |
| 68 | + color=var_name, |
| 69 | + hover_data=hover_cols, |
| 70 | + hover_name=id_col, |
| 71 | +) |
| 72 | + |
| 73 | +_add_metrics_to_legend(fig) |
| 74 | +fig.update_layout(legend=legend) |
| 75 | +add_identity_line(fig) |
| 76 | +fig.show() |
| 77 | + |
| 78 | + |
| 79 | +# %% |
| 80 | +img_path = f"{FIGS}/{today}-e-form-scatter-models" |
| 81 | +# fig.write_image(f"{img_path}.pdf") |
| 82 | +save_fig(fig, f"{img_path}.svelte") |
| 83 | + |
| 84 | + |
| 85 | +# %% scatter plot of actual vs predicted e_above_hull |
| 86 | +fig = px.scatter( |
| 87 | + df_melt.iloc[::10], |
| 88 | + x=e_above_hull_col, |
| 89 | + y=e_above_hull_preds, |
| 90 | + color=var_name, |
| 91 | + hover_data=hover_cols, |
| 92 | + hover_name=id_col, |
| 93 | +) |
| 94 | + |
| 95 | +_add_metrics_to_legend(fig) |
| 96 | +fig.update_layout(legend=legend) |
| 97 | +add_identity_line(fig) |
| 98 | +fig.show() |
| 99 | + |
| 100 | + |
| 101 | +# %% |
| 102 | +img_path = f"{FIGS}/{today}-e-above-hull-scatter-models" |
| 103 | +# fig.write_image(f"{img_path}.pdf") |
| 104 | +save_fig(fig, f"{img_path}.svelte") |
| 105 | + |
| 106 | + |
| 107 | +# %% plot all models in separate subplots |
| 108 | +true_pos, false_neg, false_pos, true_neg = classify_stable( |
| 109 | + df_melt[e_above_hull_col], df_melt[e_above_hull_preds] |
| 110 | +) |
| 111 | + |
| 112 | +df_melt["clf"] = np.array( |
| 113 | + classes := ["true positive", "false negative", "false positive", "true negative"] |
| 114 | +)[true_pos * 0 + false_neg * 1 + false_pos * 2 + true_neg * 3] |
| 115 | + |
| 116 | +fig = px.scatter( |
| 117 | + df_melt.iloc[::10], |
| 118 | + x=e_above_hull_col, |
| 119 | + y=e_above_hull_preds, |
| 120 | + facet_col=var_name, |
| 121 | + facet_col_wrap=3, |
| 122 | + hover_data=hover_cols, |
| 123 | + hover_name=id_col, |
| 124 | + color="clf", |
| 125 | + color_discrete_map=dict(zip(classes, ("green", "yellow", "red", "blue"))), |
| 126 | + opacity=0.4, |
| 127 | +) |
| 128 | + |
| 129 | +# iterate over subplots and set new title |
| 130 | +for idx, model in enumerate(models, 1): |
| 131 | + |
| 132 | + # add MAE and R2 to subplot title |
| 133 | + MAE = (df_wbm[e_form_col] - df_wbm[model]).abs().mean() |
| 134 | + R2 = r2_score(*df_wbm[[e_form_col, model]].dropna().to_numpy().T) |
| 135 | + # find index of annotation belonging to model |
| 136 | + anno_idx = [a.text for a in fig.layout.annotations].index(f"Model={model}") |
| 137 | + assert anno_idx >= 0, f"could not find annotation for {model}" |
| 138 | + # set new title |
| 139 | + fig.layout.annotations[anno_idx].text = _metric_str(model) |
| 140 | + # remove x and y axis titles if not on center row or center column |
| 141 | + if idx != 2: |
| 142 | + fig.layout[f"xaxis{idx}"].title.text = "" |
| 143 | + if idx > 1: |
| 144 | + fig.layout[f"yaxis{idx}"].title.text = "" |
| 145 | + # add vertical and horizontal lines at 0 |
| 146 | + fig.add_vline(x=0, line=dict(width=1, dash="dash", color="gray")) |
| 147 | + fig.add_hline(y=0, line=dict(width=1, dash="dash", color="gray")) |
| 148 | + |
| 149 | +id_line = add_identity_line(fig, ret_shape=True) |
| 150 | +fig.update_layout(showlegend=False) |
| 151 | +fig.update_xaxes(nticks=5) |
| 152 | +fig.update_yaxes(nticks=5) |
| 153 | + |
| 154 | +fig.show() |
| 155 | +img_path = f"{PAPER}/{today}-e-form-scatter-models.png" |
| 156 | +save_fig(fig, img_path, scale=4, width=1000, height=500) |
0 commit comments