fix legend/subplot titles in scripts/scatter_e_above_hull_models.py

add paragraph on chemical diversity of train and test set

Author: janosh

data/wbm/readme.md

@@ -14,7 +14,7 @@ The first integer in each material ID ranging from 1 to 5 and coming right after
 
 Each iteration has varying numbers of materials which are counted by the 2nd integer. Note this 2nd number is not always consecutive. A small number of materials (~0.2%) were removed by the data-cleaning steps detailed below. Don't be surprised to find an ID like `wbm-3-70804` followed by `wbm-3-70807`.
 
-## 🪓   Data processing steps
+## 🪓   Data Processing Steps
 
 The full set of processing steps used to curate the WBM test set from the raw data files (downloaded from URLs listed below) can be found in [`data/wbm/fetch_process_wbm_dataset.py`](https://github.com/janosh/matbench-discovery/blob/site/data/wbm/fetch_process_wbm_dataset.py). Processing involved
 
@@ -45,7 +45,7 @@ The number of materials in each step before and after processing are:
 | before | 61,848 | 52,800 | 79,205 | 40,328 | 23,308 | 257,487 |
 | after  | 61,466 | 52,755 | 79,160 | 40,314 | 23,268 | 256,963 |
 
-## 🔗   Links to raw WBM data files
+## 🔗   Links to raw WBM Data Files
 
 Links to WBM data files have proliferated. This is an attempt to keep track of them.
 
@@ -72,7 +72,9 @@ materialscloud:2021.68 includes a readme file with a description of the dataset,
 
 [wbm paper]: https://nature.com/articles/s41524-020-00481-6
 
-## 📊   Plots
+## 📊   Chemical Diversity
+
+Both the WBM test set and even more so the MP training set are heavily oxide dominated. The WBM test set is about 75% larger than the MP training set and also more chemically diverse, containing a higher fraction of transition metals, post-transition metals and metalloids. Our goal in picking such a large diverse test set is future-proofing. Ideally, this data will provide a challenging materials discovery test bed even for large foundational ML models in the future.
 
 <slot name="wbm-elements-heatmap">
   <img src="./figs/2023-01-08-wbm-elements.svg" alt="Periodic table log heatmap of WBM elements">

scripts/scatter_e_above_hull_models.py

@@ -45,26 +45,25 @@
 
 
 # %%
-def _metric_str(model_name: str) -> str:
-    model_pred = df_wbm[e_above_hull_col] - (df_wbm[e_form_col] - df_wbm[model_name])
-    MAE = (df_wbm[e_above_hull_col] - model_pred).abs().mean()
-    isna = df_wbm[e_above_hull_col].isna() | model_pred.isna()
-    R2 = r2_score(df_wbm[e_above_hull_col][~isna], model_pred[~isna])
-    return f"{model_name} · {MAE=:.2f} · R<sup>2</sup>={R2:.2f}"
+def _metric_str(xs: list[float], ys: list[float]) -> str:
+    # compute MAE and R2 for set of (x, y) pairs
+    isna = np.isnan(xs) | np.isnan(ys)
+    xs, ys = xs[~isna], ys[~isna]
+    MAE = np.abs(xs - ys).mean()
+    R2 = r2_score(xs, ys)
+    return f" · MAE={MAE:.2f} · R<sup>2</sup>={R2:.2f}"
 
 
 def _add_metrics_to_legend(fig: go.Figure) -> None:
     for trace in fig.data:
         # initially hide all traces, let users select which models to compare
         trace.visible = "legendonly"
-        # add MAE and R2 to legend
-        model = trace.name
-        trace.name = _metric_str(model)
+        trace.name = f"{trace.name}{_metric_str(trace.x, trace.y)}"
 
 
 # %% scatter plot of actual vs predicted e_form_per_atom
 fig = px.scatter(
-    df_melt.iloc[::10],
+    df_melt.iloc[::5],
     x=e_form_col,
     y=e_form_preds,
     color=var_name,
@@ -80,13 +79,12 @@ def _add_metrics_to_legend(fig: go.Figure) -> None:
 
 # %%
 img_path = f"{FIGS}/{today}-e-form-scatter-models"
-# fig.write_image(f"{img_path}.pdf")
-save_fig(fig, f"{img_path}.svelte")
+# save_fig(fig, f"{img_path}.svelte")
 
 
 # %% scatter plot of actual vs predicted e_above_hull
 fig = px.scatter(
-    df_melt.iloc[::10],
+    df_melt.iloc[::5],
     x=e_above_hull_col,
     y=e_above_hull_preds,
     color=var_name,
@@ -102,8 +100,7 @@ def _add_metrics_to_legend(fig: go.Figure) -> None:
 
 # %%
 img_path = f"{FIGS}/{today}-e-above-hull-scatter-models"
-# fig.write_image(f"{img_path}.pdf")
-save_fig(fig, f"{img_path}.svelte")
+# save_fig(fig, f"{img_path}.svelte")
 
 
 # %% plot all models in separate subplots
@@ -116,41 +113,75 @@ def _add_metrics_to_legend(fig: go.Figure) -> None:
 )[true_pos * 0 + false_neg * 1 + false_pos * 2 + true_neg * 3]
 
 fig = px.scatter(
-    df_melt.iloc[::10],
+    df_melt.iloc[::50],
     x=e_above_hull_col,
     y=e_above_hull_preds,
     facet_col=var_name,
     facet_col_wrap=3,
+    facet_col_spacing=0.04,
+    facet_row_spacing=0.15,
     hover_data=hover_cols,
     hover_name=id_col,
     color="clf",
     color_discrete_map=dict(zip(classes, ("green", "yellow", "red", "blue"))),
-    opacity=0.4,
+    # opacity=0.4,
+    range_x=[-2, 2],
+    range_y=[-2, 2],
 )
 
+x_title = fig.layout.xaxis.title.text
+y_title = fig.layout.yaxis.title.text
+
 # iterate over subplots and set new title
-for idx, model in enumerate(models, 1):
-    # find index of annotation belonging to model
-    anno_idx = [a.text for a in fig.layout.annotations].index(f"Model={model}")
-    assert anno_idx >= 0, f"could not find annotation for {model}"
+for idx, anno in enumerate(fig.layout.annotations, 1):
+    traces = [t for t in fig.data if t.xaxis == f"x{idx if idx > 1 else ''}"]
+    xs = np.concatenate([t.x for t in traces])
+    ys = np.concatenate([t.y for t in traces])
 
+    model = anno.text.split("=")[1]
     # set new subplot titles (adding MAE and R2)
-    fig.layout.annotations[anno_idx].text = _metric_str(model)
+    fig.layout.annotations[idx - 1].text = f"{model} {_metric_str(xs, ys)}"
 
     # remove x and y axis titles if not on center row or center column
-    if idx != 2:
-        fig.layout[f"xaxis{idx}"].title.text = ""
-    if idx > 1:
-        fig.layout[f"yaxis{idx}"].title.text = ""
+    fig.layout[f"xaxis{idx}"].title.text = ""
+    fig.layout[f"yaxis{idx}"].title.text = ""
 
     # add vertical and horizontal lines at 0
     fig.add_vline(x=0, line=dict(width=1, dash="dash", color="gray"))
     fig.add_hline(y=0, line=dict(width=1, dash="dash", color="gray"))
 
-fig.update_layout(showlegend=False)
+
 fig.update_xaxes(nticks=5)
 fig.update_yaxes(nticks=5)
 
+legend = dict(
+    title="",  # remove legend title
+    itemsizing="constant",  # increase legend marker size
+    orientation="h",
+    x=0.5,  # place legend centered above subplots
+    xanchor="center",
+    y=1.2,
+    yanchor="top",
+)
+fig.layout.legend.update(legend)
+
+axis_titles = dict(xref="paper", yref="paper", showarrow=False)
+fig.add_annotation(
+    x=0.5,
+    y=-0.16,
+    text=x_title,
+    **axis_titles,
+)
+# add y-axis title
+fig.add_annotation(
+    x=-0.06,
+    y=0.5,
+    text=y_title,
+    textangle=-90,
+    **axis_titles,
+)
+
+
 fig.show()
 img_path = f"{STATIC}/{today}-each-scatter-models.png"
-# save_fig(fig, img_path, scale=4, width=1000, height=500)
+save_fig(fig, img_path, scale=4, width=1000, height=500)

site/src/app.d.ts

@@ -2,9 +2,18 @@
 /// <reference types="mdsvex/globals" />
 
 declare module '*.md'
-declare module '*package.json'
+
+declare module '*package.json' {
+  const pkg: Record<string, unknown>
+  export default pkg
+}
 
 declare module '*metadata.yml' {
-  const content: import('$lib/types').ModelMetadata
-  export default content
+  const data: import('$lib').ModelMetadata
+  export default data
+}
+
+declare module '*element-counts.json' {
+  const map: Record<string, number>
+  export default map
 }

site/src/lib/ModelCard.svelte

@@ -2,7 +2,7 @@
   import { repository } from '$site/package.json'
   import Icon from '@iconify/svelte'
   import { pretty_num } from 'sveriodic-table/labels'
-  import type { ModelData, ModelStat } from './types'
+  import type { ModelData, ModelStat } from '.'
 
   export let key: string
   export let data: ModelData

site/src/lib/References.svelte

@@ -1,5 +1,5 @@
 <script lang="ts">
-  import type { Reference } from './types'
+  import type { Reference } from '.'
 
   export let references: Reference[]
 </script>

site/src/lib/index.ts

@@ -2,3 +2,58 @@ export { default as Footer } from './Footer.svelte'
 export { default as ModelCard } from './ModelCard.svelte'
 export { default as Nav } from './Nav.svelte'
 export { default as References } from './References.svelte'
+
+export type ModelData = ModelMetadata & ModelStats
+
+export type ModelMetadata = {
+  model_name: string
+  model_version: string
+  matbench_discovery_version: string
+  date_added: string
+  authors: Author[]
+  repo: string
+  url?: string
+  doi?: string
+  preprint?: string
+  requirements?: Record<string, string>
+  trained_on_benchmark: boolean
+}
+
+export type ModelStats = {
+  MAE: number
+  RMSE: number
+  R2: number
+  Precision: number
+  Recall: number
+  F1: number
+  missing_preds: number
+  missing_percent: number
+  Accuracy: number
+  run_time: number
+  run_time_h: string
+  GPUs: number
+  CPUs: number
+  slurm_jobs: number
+}
+
+export type Author = {
+  name: string
+  email?: string
+  affiliation?: string
+  orcid?: string
+  url?: string
+  twitter?: string
+}
+
+export type Reference = {
+  title: string
+  id: string
+  author: { family: string; given: string }[]
+  DOI: string
+  URL?: string
+  issued: { year: number; month: number; day: number }[]
+  accessed: { year: number; month: number; day: number }[]
+  page: string
+  type: string
+  ISSN?: string
+}

site/src/routes/about-the-test-set/+page.svelte

@@ -10,11 +10,18 @@
   let log = false // log color scale
   const wbm_heat_vals: number[] = Object.values(wbm_elem_counts)
   const mp_heat_vals: number[] = Object.values(mp_elem_counts)
-  const color_map = {
+  const [mp_max, wbm_max] = [Math.max(...mp_heat_vals), Math.max(...wbm_heat_vals)]
+  const mp_color_map = {
     200: `blue`,
-    35_000: `green`,
-    80_000: `yellow`,
-    150_000: `red`,
+    [mp_max / 4]: `green`,
+    [mp_max / 2]: `yellow`,
+    mp_max: `red`,
+  }
+  const wbm_color_map = {
+    200: `blue`,
+    [wbm_max / 4]: `green`,
+    [wbm_max / 2]: `yellow`,
+    wbm_max: `red`,
   }
   let active_mp_elem: ChemicalElement
   let active_wbm_elem: ChemicalElement
@@ -30,7 +37,7 @@
     <span>Log color scale <Toggle bind:checked={log} /></span>
     <PeriodicTable
       heatmap_values={wbm_heat_vals}
-      {color_map}
+      color_map={wbm_color_map}
       {log}
       bind:active_element={active_wbm_elem}
     >
@@ -51,9 +58,10 @@
     </PeriodicTable>
   </svelte:fragment>
   <svelte:fragment slot="mp-elements-heatmap">
+    <span>Log color scale <Toggle bind:checked={log} /></span>
     <PeriodicTable
       heatmap_values={mp_heat_vals}
-      {color_map}
+      color_map={mp_color_map}
       {log}
       bind:active_element={active_mp_elem}
     >

site/src/routes/models/+page.server.ts

@@ -1,9 +1,9 @@
-import type { ModelData, ModelMetadata } from '$lib/types'
+import type { ModelData, ModelMetadata } from '$lib'
 import { dirname } from 'path'
 import type { PageServerLoad } from './$types'
 import model_stats from './2023-01-23-model-stats.json'
 
-export const load: PageServerLoad = async () => {
+export const load: PageServerLoad = () => {
   const yml = import.meta.glob(`$root/models/**/metadata.yml`, {
     eager: true,
   })

site/src/routes/models/+page.svelte

@@ -1,6 +1,6 @@
 <script lang="ts">
+  import type { ModelStats } from '$lib'
   import { ModelCard } from '$lib'
-  import type { ModelStats } from '$lib/types'
   import { RadioButtons } from 'svelte-zoo'
   import { flip } from 'svelte/animate'
   import { fade } from 'svelte/transition'

site/src/routes/paper/+page.svx

@@ -167,7 +167,7 @@ Our benchmark is designed to make [adding future models easy](/how-to-contribute
 
 1. [CGCNN](https://arxiv.org/abs/1710.10324) @xie_crystal_2018
 1. [BOWSR](https://sciencedirect.com/science/article/pii/S1369702121002984) @zuo_accelerating_2021
-1. [Wren](https://arxiv.org/abs/2106.11132) @goodall_rapid_2022
+1. [Wrenformer](https://arxiv.org/abs/2106.11132) @goodall_rapid_2022
 1. [M3GNet](https://arxiv.org/abs/2202.02450) @chen_universal_2022
 1. [MEGNet](https://arxiv.org/abs/1812.05055) @chen_graph_2019
 1. [Voronoi Random Forest](https://journals.aps.org/prb/abstract/10.1103/PhysRevB.96.024104) @goodall_rapid_2022
@@ -181,7 +181,7 @@ Our benchmark is designed to make [adding future models easy](/how-to-contribute
   Classification performance for all models
 </caption>
 
-![Parity plot for each model's energy above hull predictions (based on their formation energy preds) vs DFT ground truth](./figs/2023-01-18-each-scatter-models.png)
+![Parity plot for each model's energy above hull predictions (based on their formation energy preds) vs DFT ground truth](./figs/2023-01-24-each-scatter-models.png)
 
 <figcaption>@label:fig:each-scatter-models Parity plot for each model's energy above hull predictions (based on their formation energy preds) vs DFT ground truth</figcaption>
 
@@ -190,9 +190,9 @@ Our benchmark is designed to make [adding future models easy](/how-to-contribute
 <figcaption>@label:fig:wbm-hull-dist-hist-models Histograms and rolling accuracy of using predicted formation energies for stability classification</figcaption>
 
 <div>
-{#if typeof document !== `undefined`}
-  <CumulativeClfMetrics class="pull-left" />
-{/if}
+  {#if typeof document !== `undefined`}
+    <CumulativeClfMetrics class="pull-left" />
+  {/if}
 </div>
 
 ## Analysis