Refactor ptable plotters and add ptable_heatmap with diagonally-split tiles

.github/workflows/svgo.yml

@@ -7,6 +7,8 @@ on:
 
 jobs:
   tests:
+    # don't run on PRs from forks
+    if: github.event.pull_request.head.repo.fork == false
     runs-on: ubuntu-latest
     steps:
       - name: Check out repo

citation.cff

@@ -5,16 +5,23 @@ authors:
   - family-names: Riebesell
     given-names: Janosh
     affiliation: University of Cambridge
-    email: janosh@lbl.gov
+    email: janosh[email protected]
     orcid: https://orcid.org/0000-0001-5233-3462
+    github: janosh
   - family-names: Goodall
     given-names: Rhys
     affiliation: University of Cambridge
     orcid: https://orcid.org/0000-0002-6589-1700
+    github: comprhys
   - family-names: Baird
     given-names: Sterling G.
     affiliation: University of Utah
     orcid: https://orcid.org/0000-0002-4491-6876
+    github: sgbaird
+  - family-names: Yang
+    given-names: Haoyu (Daniel)
+    email: [email protected]
+    github: DanielYang59
 license: MIT
 license-url: https://github.com/janosh/pymatviz/blob/main/license"
 repository-code: https://github.com/janosh/pymatviz

examples/_generate_assets.py

@@ -1,5 +1,6 @@
 # %%
 import json
+import random
 from glob import glob
 
 import matplotlib.pyplot as plt
@@ -37,8 +38,10 @@
     ptable_heatmap,
     ptable_heatmap_plotly,
     ptable_heatmap_ratio,
+    ptable_heatmap_splits,
     ptable_hists,
-    ptable_plots,
+    ptable_lines,
+    ptable_scatters,
 )
 from pymatviz.relevance import precision_recall_curve, roc_curve
 from pymatviz.sankey import sankey_from_2_df_cols
@@ -48,7 +51,7 @@
 from pymatviz.utils import TEST_FILES, df_ptable
 
 
-# %% configure matplotlib and load test data
+# %% Configure matplotlib and load test data
 plt.rc("font", size=14)
 plt.rc("savefig", bbox="tight", dpi=200)
 plt.rc("axes", titlesize=16, titleweight="bold")
@@ -58,7 +61,7 @@
 px.defaults.template = "pymatviz_white"
 pio.templates.default = "pymatviz_white"
 
-# random classification data
+# Random classification data
 np.random.seed(42)
 rand_clf_size = 100
 y_binary = np.random.choice([0, 1], size=rand_clf_size)
@@ -67,7 +70,7 @@
 )
 
 
-# random regression data
+# Random regression data
 rand_regression_size = 500
 y_true = np.random.normal(5, 4, rand_regression_size)
 y_pred = 1.2 * y_true - 2 * np.random.normal(0, 1, rand_regression_size)
@@ -164,7 +167,7 @@
 
 
 # %% Histograms laid out in as a periodic table
-# generate random parity data with y \approx x with some noise
+# Generate random parity data with y \approx x with some noise
 data_dict = {
     elem.symbol: np.random.randn(100) + np.random.randn(100) for elem in Element
 }
@@ -179,20 +182,57 @@
     elem.symbol: [
         np.random.randint(0, 20, 10),
         np.random.randint(0, 20, 10),
-        np.random.randint(0, 20, 10),
+        # np.random.randint(0, 20, 10),  # TODO: allow 3rd dim
     ]
     for elem in Element
 }
 
-fig = ptable_plots(
+fig = ptable_scatters(
     data_dict,
-    colormap="coolwarm",
-    cbar_title="Periodic Table Scatter Plots",
-    plot_kwds=dict(marker="o", linestyle=""),
+    # colormap="coolwarm",
+    # cbar_title="Periodic Table Scatter Plots",
+    child_args=dict(marker="o", linestyle=""),
+    symbol_pos=(0.5, 1.2),
+    symbol_kwargs=dict(fontsize=14),
 )
 save_and_compress_svg(fig, "ptable-scatters")
 
 
+# %% Line plots laid out as a periodic table
+data_dict = {
+    elem.symbol: [
+        np.linspace(0, 10, 10),
+        np.sin(2 * np.pi * np.linspace(0, 10, 10)) + np.random.normal(0, 0.2, 10),
+    ]
+    for elem in Element
+}
+
+fig = ptable_lines(
+    data_dict,
+    symbol_pos=(0.5, 1.2),
+    symbol_kwargs=dict(fontsize=14),
+)
+save_and_compress_svg(fig, "ptable-lines")
+
+
+# %% Evenly-split tile plots laid out as a periodic table
+data_dict = {
+    elem.symbol: [
+        random.randint(0, 10),
+        random.randint(10, 20),
+    ]
+    for elem in Element
+}
+
+fig = ptable_heatmap_splits(
+    data=data_dict,
+    colormap="coolwarm",
+    start_angle=135,
+    cbar_title="Periodic Table Evenly-Split Heatmap Plots",
+)
+save_and_compress_svg(fig, "ptable-heatmap-splits")
+
+
 # %% Uncertainty Plots
 ax = qq_gaussian(y_pred, y_true, y_std, identity_line={"line_kwds": {"color": "red"}})
 save_and_compress_svg(ax, "normal-prob-plot")
@@ -261,7 +301,7 @@
 
 
 # %% Correlation Plots
-# plot eigenvalue distribution of a pure-noise correlation matrix
+# Plot eigenvalue distribution of a pure-noise correlation matrix
 # i.e. the correlation matrix contains no significant correlations
 # beyond the spurious correlation that occurs randomly
 n_rows, n_cols = 500, 1000
@@ -272,7 +312,7 @@
 ax = marchenko_pastur(corr_mat, gamma=n_cols / n_rows)
 save_and_compress_svg(ax, "marchenko-pastur")
 
-# plot eigenvalue distribution of a correlation matrix with significant
+# Plot eigenvalue distribution of a correlation matrix with significant
 # (i.e. non-noise) eigenvalue
 n_rows, n_cols = 50, 400
 linear_matrix = np.arange(n_rows * n_cols).reshape(n_rows, n_cols) / n_cols
@@ -282,7 +322,7 @@
 ax = marchenko_pastur(corr_mat, gamma=n_cols / n_rows)
 save_and_compress_svg(ax, "marchenko-pastur-significant-eval")
 
-# plot eigenvalue distribution of a rank-deficient correlation matrix
+# Plot eigenvalue distribution of a rank-deficient correlation matrix
 n_rows, n_cols = 600, 500
 rand_tall_mat = np.random.normal(0, 1, size=(n_rows, n_cols))
 

examples/dataset_exploration/wbm/explore_wbm.py

@@ -10,10 +10,10 @@
 __date__ = "2022-08-18"
 
 
-# %% download wbm-steps-summary.csv (23.31 MB)
-df_wbm = pd.read_csv(
-    "https://figshare.com/files/37542841?private_link=ff0ad14505f9624f0c05"
-).set_index("material_id", drop=False)
+# %% download wbm-summary.csv (12 MB)
+df_wbm = pd.read_csv("https://figshare.com/ndownloader/files/44225498").set_index(
+    "material_id", drop=False
+)
 
 df_wbm["batch_idx"] = df_wbm.index.str.split("-").str[2].astype(int)
 df_wbm["spg_num"] = df_wbm.wyckoff.str.split("_").str[2].astype(int)

examples/diatomics/plot.py

@@ -16,7 +16,7 @@
 from matplotlib.backends.backend_pdf import PdfPages
 from pymatgen.core import Element
 
-from pymatviz import ptable_plots
+from pymatviz import ptable_plots  # type: ignore[attr-defined]  # TODO:
 from pymatviz.io import save_fig
 from pymatviz.utils import df_ptable
 

examples/matbench_dielectric_eda.ipynb

@@ -33,7 +33,7 @@
     "from tqdm import tqdm\n",
     "\n",
     "from pymatviz import ptable_heatmap, spacegroup_hist, spacegroup_sunburst\n",
-    "from pymatviz.utils import get_crystal_sys\n",
+    "from pymatviz.utils import crystal_sys_from_spg_num\n",
     "\n",
     "\n",
     "__author__ = \"Janosh Riebesell\"\n",
@@ -64,7 +64,7 @@
     "df_diel[[\"spg_symbol\", \"spg_num\"]] = [\n",
     "    struct.get_space_group_info() for struct in tqdm(df_diel.structure)\n",
     "]\n",
-    "df_diel[\"crys_sys\"] = df_diel.spg_num.map(get_crystal_sys)"
+    "df_diel[\"crys_sys\"] = df_diel.spg_num.map(crystal_sys_from_spg_num)"
    ]
   },
   {

examples/matbench_perovskites_eda.ipynb

@@ -26,23 +26,21 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import matplotlib.pyplot as plt\n",
     "import pandas as pd\n",
+    "import plotly.express as px\n",
     "import plotly.io as pio\n",
     "\n",
     "# see https://github.com/CompRhys/aviary/blob/main/aviary/wren/utils.py\n",
     "from aviary.wren.utils import get_aflow_label_aflow\n",
     "from matminer.datasets import load_dataset\n",
     "from tqdm import tqdm\n",
     "\n",
-    "from pymatviz import (\n",
-    "    annotate_bars,\n",
-    "    plot_structure_2d,\n",
-    "    ptable_heatmap_plotly,\n",
-    "    spacegroup_sunburst,\n",
-    ")\n",
-    "from pymatviz.plot_defaults import plt, px\n",
+    "from pymatviz import plot_structure_2d, ptable_heatmap_plotly, spacegroup_sunburst\n",
+    "from pymatviz.io import save_and_compress_svg\n",
+    "from pymatviz.powerups import annotate_bars\n",
     "from pymatviz.sankey import sankey_from_2_df_cols\n",
-    "from pymatviz.utils import get_crystal_sys, save_and_compress_svg\n",
+    "from pymatviz.utils import crystal_sys_from_spg_num\n",
     "\n",
     "\n",
     "__author__ = \"Janosh Riebesell\"\n",
@@ -79,7 +77,7 @@
     "\n",
     "df_perov[\"formula\"] = df_perov.structure.map(lambda cryst: cryst.formula)\n",
     "\n",
-    "df_perov[\"spglib_crys_sys\"] = df_perov.spglib_spg_num.map(get_crystal_sys)"
+    "df_perov[\"spglib_crys_sys\"] = df_perov.spglib_spg_num.map(crystal_sys_from_spg_num)"
    ]
   },
   {
@@ -291,7 +289,7 @@
    "outputs": [],
    "source": [
     "df_perov[\"aflow_spg_num\"] = df_perov.aflow_wyckoff.str.split(\"_\").str[2].astype(int)\n",
-    "df_perov[\"aflow_crys_sys\"] = df_perov.aflow_spg_num.map(get_crystal_sys)"
+    "df_perov[\"aflow_crys_sys\"] = df_perov.aflow_spg_num.map(crystal_sys_from_spg_num)"
    ]
   },
   {
@@ -332,11 +330,9 @@
    "source": [
     "fig = sankey_from_2_df_cols(df_perov, [\"spglib_spg_num\", \"aflow_spg_num\"])\n",
     "\n",
-    "fig.update_layout(\n",
-    "    title=\"Spacegroups as determined by Spglib vs Aflow<br>\"\n",
-    "    \"for the Matbench Perovskites dataset\",\n",
-    "    title_x=0.5,\n",
-    ")\n",
+    "title = \"Spglib vs Aflow Spacegroups<br>for the Matbench Perovskites dataset\"\n",
+    "\n",
+    "fig.layout.title.update(text=title, x=0.5)\n",
     "\n",
     "save_and_compress_svg(fig, \"sankey-spglib-vs-aflow-spacegroups\")\n",
     "\n",
@@ -361,11 +357,8 @@
    "source": [
     "fig = sankey_from_2_df_cols(df_perov, [\"spglib_crys_sys\", \"aflow_crys_sys\"])\n",
     "\n",
-    "fig.update_layout(\n",
-    "    title=\"Crystal systems as determined by Spglib vs Aflow<br>\"\n",
-    "    \"for the Matbench Perovskites dataset\",\n",
-    "    title_x=0.5,\n",
-    ")"
+    "title = \"Spglib vs Aflow Crystal systems<br>for the Matbench Perovskites dataset\"\n",
+    "fig.layout.title.update(text=title, x=0.5)"
    ]
   }
  ],

pymatviz/__init__.py

@@ -36,12 +36,16 @@
     plot_phonon_dos,
 )
 from pymatviz.ptable import (
+    ChildPlotters,
+    PTableProjector,
     count_elements,
     ptable_heatmap,
     ptable_heatmap_plotly,
     ptable_heatmap_ratio,
+    ptable_heatmap_splits,
     ptable_hists,
-    ptable_plots,
+    ptable_lines,
+    ptable_scatters,
 )
 from pymatviz.relevance import precision_recall_curve, roc_curve
 from pymatviz.sankey import sankey_from_2_df_cols

pymatviz/correlation.py

@@ -1,3 +1,5 @@
+"""Plot distributions of correlation matrix eigenvalues."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING

pymatviz/cumulative.py

@@ -1,3 +1,5 @@
+"""Plot the cumulative distribution of residuals and absolute errors."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any

pymatviz/histograms.py

@@ -1,3 +1,5 @@
+"""Histograms and bar charts."""
+
 from __future__ import annotations
 
 from collections.abc import Sequence

pymatviz/io.py

@@ -1,3 +1,5 @@
+"""I/O utilities for saving figures and dataframes to various image formats."""
+
 from __future__ import annotations
 
 import copy

pymatviz/parity.py

@@ -1,3 +1,5 @@
+"""Parity, residual and density plots."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any

pymatviz/phonons.py

@@ -1,3 +1,5 @@
+"""Plotting functions for pymatgen phonon band structures and density of states."""
+
 from __future__ import annotations
 
 import sys
@@ -384,8 +386,8 @@ def plot_phonon_dos(
     fig.layout.font.size = 16 * (fig.layout.width or 800) / 800
     fig.layout.legend.update(x=0.005, y=0.99, orientation="h", yanchor="top")
 
-    qual_colors = px.colors.qualitative.Plotly
     if last_peak_anno:
+        qual_colors = px.colors.qualitative.Plotly
         for idx, (key, dos) in enumerate(doses.items()):
             last_peak = dos.get_last_peak()
             color = (

pymatviz/powerups.py

@@ -1,3 +1,7 @@
+"""Powerups/enhancements such as parity lines, annotations and marginals for matplotlib
+and plotly figures.
+"""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any
@@ -196,12 +200,10 @@ def annotate_metrics(
         "MSE": lambda x, y: ((x - y) ** 2).mean(),
         "MAPE": mape,
         "R2": r2_score,
-        # TODO: check this for correctness
         "R2_adj": lambda x, y: 1 - (1 - r2_score(x, y)) * (len(x) - 1) / (len(x) - 2),
     }
     for key in set(metrics) - set(funcs):
-        func = getattr(sklearn.metrics, key, None)
-        if func:
+        if func := getattr(sklearn.metrics, key, None):
             funcs[key] = func
     if bad_keys := set(metrics) - set(funcs):
         raise ValueError(f"Unrecognized metrics: {bad_keys}")