add backend=plotly | matplotlib to hist_classified_stable_vs_hull_dist()

janosh · janosh · commit 0fb755053e9d · 2023-06-19T20:29:22.000-07:00
diff --git a/matbench_discovery/plots.py b/matbench_discovery/plots.py
@@ -6,6 +6,7 @@
 import numpy as np
 import pandas as pd
 import plotly.express as px
+import plotly.graph_objs as go
 import plotly.io as pio
 import scipy.interpolate
 import scipy.stats
@@ -19,7 +20,7 @@
 
 WhichEnergy = Literal["true", "pred"]
 AxLine = Literal["x", "y", "xy", ""]
-
+Backend = Literal["matplotlib", "plotly"]
 
 # --- start global plot settings
 quantity_labels = dict(
@@ -53,8 +54,11 @@
     dft="DFT",
 )
 px.defaults.labels = quantity_labels | model_labels
-
-pio.templates.default = "plotly_white"
+pastel_layout = dict(
+    colorway=px.colors.qualitative.Pastel, margin=dict(l=40, r=30, t=60, b=30)
+)
+pio.templates["pastel"] = dict(layout=pastel_layout)
+pio.templates.default = "plotly_white+pastel"
 
 # https://github.com/plotly/Kaleido/issues/122#issuecomment-994906924
 # when seeing MathJax "loading" message in exported PDFs, try:
@@ -79,7 +83,9 @@ def hist_classified_stable_vs_hull_dist(
     show_threshold: bool = True,
     x_lim: tuple[float | None, float | None] = (-0.4, 0.4),
     rolling_accuracy: float | None = 0.02,
-) -> tuple[plt.Axes, dict[str, float]]:
+    backend: Backend = "plotly",
+    ylabel: str = "Number of materials",
+) -> tuple[plt.Axes | go.Figure, dict[str, float]]:
     """
     Histogram of the energy difference (either according to DFT ground truth [default]
     or model predicted energy) to the convex hull for materials in the WBM data set. The
@@ -106,16 +112,16 @@ def hist_classified_stable_vs_hull_dist(
             vertical line.
         x_lim (tuple[float | None, float | None]): x-axis limits.
         rolling_accuracy (float): Rolling accuracy window size in eV / atom. Set to None
-            or 0 to disable. Defaults to 0.01.
+            or 0 to disable. Defaults to 0.02, meaning 20 meV / atom.
+        backend ('matplotlib' | 'plotly'], optional): Which plotting backend to use.
+            Changes the return type.
 
     Returns:
         tuple[plt.Axes, dict[str, float]]: plot axes and classification metrics
 
     NOTE this figure plots hist bars separately which causes aliasing in pdf. Can be
     fixed in Inkscape or similar by merging regions by color.
     """
-    ax = ax or plt.gca()
-
     true_pos, false_neg, false_pos, true_neg = classify_stable(
         e_above_hull_true, e_above_hull_pred, stability_threshold
     )
@@ -131,90 +137,105 @@ def hist_classified_stable_vs_hull_dist(
     eah_false_neg = e_above_hull[false_neg]
     eah_false_pos = e_above_hull[false_pos]
     eah_true_neg = e_above_hull[true_neg]
-    xlabel = dict(
-        true="$E_\\mathrm{above\\ hull}$ (eV / atom)",
-        pred="$E_\\mathrm{above\\ hull\\ pred}$ (eV / atom)",
-    )[which_energy]
-
-    ax.hist(
-        [eah_true_pos, eah_false_neg, eah_false_pos, eah_true_neg],
-        bins=200,
-        range=x_lim,
-        alpha=0.5,
-        color=["tab:green", "tab:orange", "tab:red", "tab:blue"],
-        label=[
-            "True Positives",
-            "False Negatives",
-            "False Positives",
-            "True Negatives",
-        ],
-        stacked=True,
-    )
-
     n_true_pos, n_false_pos, n_true_neg, n_false_neg = map(
-        len, (eah_true_pos, eah_false_pos, eah_true_neg, eah_false_neg)
+        sum, (true_pos, false_pos, true_neg, false_neg)
     )
     # null = (tp + fn) / (tp + tn + fp + fn)
     precision = n_true_pos / (n_true_pos + n_false_pos)
 
-    # assert (n_all := n_true_pos + n_false_pos + n_true_neg + n_false_neg) == len(
-    #     e_above_hull_true
-    # ), f"{n_all} != {len(e_above_hull_true)}"
-
-    ax.set(xlabel=xlabel, ylabel="Number of compounds", xlim=x_lim)
-
-    if rolling_accuracy:
-        # add moving average of the accuracy (computed within 20 meV/atom intervals) as
-        # a function of ΔHd,MP is shown as a blue line (right axis)
-        ax_acc = ax.twinx()
-        ax_acc.set_ylabel("Accuracy", color="darkblue")
-        ax_acc.tick_params(labelcolor="darkblue")
-        ax_acc.set(ylim=(0, 1))
-
-        # --- moving average of the accuracy
-        # compute accuracy within 20 meV/atom intervals
-        bins = np.arange(x_lim[0], x_lim[1], rolling_accuracy)
-        bin_counts = np.histogram(e_above_hull_true, bins)[0]
-        bin_true_pos = np.histogram(eah_true_pos, bins)[0]
-        bin_true_neg = np.histogram(eah_true_neg, bins)[0]
-
-        # compute accuracy
-        bin_accuracies = (bin_true_pos + bin_true_neg) / bin_counts
-        # plot accuracy
-        ax_acc.plot(
-            bins[:-1],
-            bin_accuracies,
-            color="tab:blue",
-            label="Accuracy",
-            linewidth=3,
+    xlabel = dict(
+        true=r"$E_\mathrm{above\ hull}\;\mathrm{(eV / atom)}$",
+        pred=r"$E_\mathrm{above\ hull\ pred}\;\mathrm{(eV / atom)}$",
+    )[which_energy]
+    labels = ["True Positives", "False Negatives", "False Positives", "True Negatives"]
+
+    if backend == "matplotlib":
+        ax = ax or plt.gca()
+        ax.hist(
+            [eah_true_pos, eah_false_neg, eah_false_pos, eah_true_neg],
+            bins=200,
+            range=x_lim,
+            alpha=0.5,
+            color=["tab:green", "tab:orange", "tab:red", "tab:blue"],
+            label=labels,
+            stacked=True,
         )
-        # ax2.fill_between(
-        #     bin_centers,
-        #     bin_accuracy - bin_accuracy_std,
-        #     bin_accuracy + bin_accuracy_std,
-        #     color="tab:blue",
-        #     alpha=0.2,
-        # )
-
-    if show_threshold:
+        ax.set(xlabel=xlabel, ylabel=ylabel, xlim=x_lim)
+
         ax.axvline(
             stability_threshold,
-            color="k",
+            color="black",
             linestyle="--",
             label="Stability Threshold",
         )
 
-    recall = n_true_pos / n_total_pos
+        if rolling_accuracy:
+            # add moving average of the accuracy computed within given window
+            # as a function of e_above_hull shown as blue line (right axis)
+            ax_acc = ax.twinx()
+            ax_acc.set_ylabel("Accuracy", color="darkblue")
+            ax_acc.tick_params(labelcolor="darkblue")
+            ax_acc.set(ylim=(0, 1))
+
+            # --- moving average of the accuracy
+            # compute accuracy within 20 meV/atom intervals
+            bins = np.arange(x_lim[0], x_lim[1], rolling_accuracy)
+            bin_counts = np.histogram(e_above_hull_true, bins)[0]
+            bin_true_pos = np.histogram(eah_true_pos, bins)[0]
+            bin_true_neg = np.histogram(eah_true_neg, bins)[0]
+
+            # compute accuracy
+            bin_accuracies = (bin_true_pos + bin_true_neg) / bin_counts
+            # plot accuracy
+            ax_acc.plot(
+                bins[:-1],
+                bin_accuracies,
+                color="tab:blue",
+                label="Accuracy",
+                linewidth=3,
+            )
+            # ax2.fill_between(
+            #     bin_centers,
+            #     bin_accuracy - bin_accuracy_std,
+            #     bin_accuracy + bin_accuracy_std,
+            #     color="tab:blue",
+            #     alpha=0.2,
+            # )
+
+    if backend == "plotly":
+        clf = (true_pos * 1 + false_neg * 2 + false_pos * 3 + true_neg * 4).map(
+            dict(zip(range(1, 5), labels))
+        )
+        df = pd.DataFrame(dict(e_above_hull=e_above_hull, clf=clf))
 
-    return ax, {
-        "enrichment": precision / null,
-        "precision": precision,
-        "recall": recall,
-        "prevalence": null,
-        "accuracy": (n_true_pos + n_true_neg)
+        ax = px.histogram(
+            df, x="e_above_hull", color="clf", nbins=20000, range_x=x_lim, opacity=0.9
+        )
+        ax.update_layout(
+            dict(xaxis_title=xlabel, yaxis_title=ylabel),
+            legend=dict(title=None, yanchor="top", y=1, xanchor="right", x=1),
+        )
+
+        ax.add_vline(stability_threshold, line=dict(dash="dash", width=1))
+        ax.add_annotation(
+            text="Stability threshold",
+            x=stability_threshold,
+            y=1.1,
+            yref="paper",
+            font=dict(size=14, color="gray"),
+            showarrow=False,
+        )
+
+    recall = n_true_pos / n_total_pos
+    return ax, dict(
+        enrichment=precision / null,
+        precision=precision,
+        recall=recall,
+        prevalence=null,
+        accuracy=(n_true_pos + n_true_neg)
         / (n_true_pos + n_true_neg + n_false_pos + n_false_neg),
-        "f1": 2 * (precision * recall) / (precision + recall),
-    }
+        f1=2 * (precision * recall) / (precision + recall),
+    )
 
 
 def rolling_mae_vs_hull_dist(
@@ -432,7 +453,7 @@ def cumulative_clf_metric(
 
 
 def wandb_scatter(table: wandb.Table, fields: dict[str, str], **kwargs: Any) -> None:
-    """Log a parity scatter plot using custom vega spec to WandB.
+    """Log a parity scatter plot using custom Vega spec to WandB.
 
     Args:
         table (wandb.Table): WandB data table.
diff --git a/scripts/hist_classified_stable_vs_hull_dist.py b/scripts/hist_classified_stable_vs_hull_dist.py
@@ -1,5 +1,5 @@
 # %%
-from matbench_discovery import today
+from matbench_discovery import ROOT, today
 from matbench_discovery.load_preds import load_df_wbm_with_preds
 from matbench_discovery.plots import WhichEnergy, hist_classified_stable_vs_hull_dist
 
@@ -50,15 +50,19 @@
     which_energy=which_energy,
     # stability_threshold=-0.05,
     rolling_accuracy=0,
+    # backend="matplotlib",
 )
+if hasattr(ax, "legend"):
+    legend_title = f"Enrichment Factor = {metrics['enrichment']:.3}"
+    ax.legend(loc="upper left", frameon=False, title=legend_title)
 
-fig = ax.figure
-fig.set_size_inches(10, 9)
-
-legend_title = f"Enrichment Factor = {metrics['enrichment']:.3}"
-ax.legend(loc="center left", frameon=False, title=legend_title)
+ax
 
 
 # %%
-fig_name = f"{today}-wren-wbm-hull-dist-hist-{which_energy=}"
-# fig.savefig(f"{ROOT}/figures/{fig_name}.pdf")
+fig_name = f"{ROOT}/figures/{today}-wren-wbm-hull-dist-hist-{which_energy=}.pdf"
+if hasattr(ax, "write_image"):
+    # fig.write_image(fig_name)
+    ax.write_html(fig_name.replace(".pdf", ".html"))
+else:
+    ax.figure.savefig(fig_name)
diff --git a/scripts/hist_classified_stable_vs_hull_dist_batches.py b/scripts/hist_classified_stable_vs_hull_dist_batches.py
@@ -56,6 +56,7 @@
     + (df_wbm[model_name] - df_wbm[target_col]),
     which_energy=which_energy,
     ax=axs.flat[-1],
+    backend="matplotlib",
 )
 
 text = f"Enrichment\nFactor = {metrics['enrichment']:.3}"
diff --git a/scripts/hist_classified_stable_vs_hull_dist_models.py b/scripts/hist_classified_stable_vs_hull_dist_models.py
@@ -1,7 +1,10 @@
 # %%
+from plotly.subplots import make_subplots
+
 from matbench_discovery import ROOT, today
 from matbench_discovery.load_preds import load_df_wbm_with_preds
 from matbench_discovery.plots import (
+    Backend,
     WhichEnergy,
     hist_classified_stable_vs_hull_dist,
     plt,
@@ -30,30 +33,49 @@
 
 # %%
 which_energy: WhichEnergy = "true"
-fig, axs = plt.subplots(3, 3, figsize=(18, 12))
-
 model_name = "Wrenformer"
 
-for model_name, ax in zip(models, axs.flat, strict=True):
+backend: Backend = "matplotlib"
+if backend == "matplotlib":
+    fig, axs = plt.subplots(3, 3, figsize=(18, 12))
+else:
+    fig = make_subplots(rows=3, cols=3)
+
 
+for idx, model_name in enumerate(models):
     ax, metrics = hist_classified_stable_vs_hull_dist(
         e_above_hull_true=df_wbm[e_above_hull_col],
         e_above_hull_pred=df_wbm[e_above_hull_col]
         + (df_wbm[model_name] - df_wbm[target_col]),
         which_energy=which_energy,
-        ax=ax,
+        ax=axs.flat[idx],
+        backend=backend,
     )
-
-    text = f"Enrichment\nFactor = {metrics['enrichment']:.3}"
-    ax.text(0.02, 0.25, text, fontsize=16, transform=ax.transAxes)
-
     title = f"{model_name} ({len(df_wbm[model_name].dropna()):,})"
-    ax.set(title=title)
-
+    text = f"Enrichment\nFactor = {metrics['enrichment']:.3}"
 
-# axs.flat[0].legend(frameon=False, loc="upper left")
+    if backend == "matplotlib":
+        ax.text(0.02, 0.25, text, fontsize=16, transform=ax.transAxes)
+        ax.set(title=title)
+
+    else:
+        ax.add_annotation(text=text, x=0.5, y=0.5, showarrow=False)
+        ax.update_xaxes(title_text=title)
+
+        for trace in ax.data:
+            fig.append_trace(trace, row=idx % 3 + 1, col=idx // 3 + 1)
+
+if backend == "matplotlib":
+    fig.suptitle(f"{today} {which_energy=}", y=1.07, fontsize=16)
+    plt.figlegend(
+        *ax.get_legend_handles_labels(),
+        ncol=10,
+        loc="lower center",
+        bbox_to_anchor=(0.5, -0.05),
+        frameon=False,
+    )
 
-fig.suptitle(f"{today} {which_energy=}", y=1.07, fontsize=16)
+fig
 
 
 # %%
diff --git a/tests/test_plots.py b/tests/test_plots.py

Original file line number	Diff line number	Diff line change
`@@ -56,6 +56,7 @@`
`56`	`56`	`+ (df_wbm[model_name] - df_wbm[target_col]),`
`57`	`57`	`which_energy=which_energy,`
`58`	`58`	`ax=axs.flat[-1],`
	`59`	`+ backend="matplotlib",`
`59`	`60`	`)`
`60`	`61`
`61`	`62`	`text = f"Enrichment\nFactor = {metrics['enrichment']:.3}"`