[Breaking] Refactor ptable heatmap plotter (#157)

* comment out heatmap plotter for refactor

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* group args for refactor

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove None as default

* remove commented out legacy code

* fix unsupported default None value for elem_colors

* add missing start_angle arg

* add heatmap child plotter

* remove Callable type from plot_kwargs

* set hide_f_block default to "AUTO"

* docstring tweak

* add values display functionality

* add docstring for heatmap plotter

* add fences for blocks in docstring

* add log scale preprocessor

* warning upon illegal log value

* tweak docstring

* skip ptable heatmap tests

* add more comments

* add data sum normalizer

* merge heat_mode and show_values to show_values_mode

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* change default symbol font to white and bold

* make default cbar title font 16 and bold

* fix missing % sign in percent mode

* rename show_values_mode to values_show_mode

* add tile_colors to store and overwrite tile colors

* add values f-string formatter

* add setting cbar range, need double check overwrite norm

* temp save: add tile_size docstring

* change border color to grey

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix value in percent mode

* remove percentage mode from normalizer

* fix merge main

* remove TODO tag

* enable cbar label format in percent mode

* WIP: separate label fmt generate as method

* WIP: add sci_notation arg

* add cbar label sci notation formatter

* fix sci notation in fraction mode

* add TODO tag for tests to migrate

* move get_cbar_label_formatter to utils

* turn off useMathText for consistency with values in tile

* remove percentage arg

* add tests for get_cbar_label_formatter

* track examples, would remove before merge

* overwrite add_child_plots method

* access colors from tile_colors property

* tiny docstring correction

* remove accidentally committed test file

* simplify sci notation of e+0

* SUBOPTIMAL: pass tile_size to subplots figsize

* WIP, not working: add f-block vertical offset

* fix wrong arg passing

* separate HMapPTableProjector

* preprocess_ptable_data and downstream funcs now record anomaly

* remove TODO tag

* overload preprocess_ptable_data type

* add overwrite_colors generator and unit test

* fix log_scale apply, and return type

* docstring tweaks

* In process: migration some of the heatmap tests

* remove TODO tag

* fix mypy errors

* migrate most unit tests

* migrate most unit tests

* add overwrite_colors arg

* add PTableData and unit test, not fully finished

* add drop_elements method

* add unit test for apply method

* move PTableData to ptable and relocate unit tests

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add placeholder wrap for from_formulas method

* fix nan replacement when nan and inf coexist

* fix log scale apply

* fix type errors

* simplify missing_strategy arg name to strategy

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* mark tile_size and f-block offset as WIP

* add text_colors to unify value/symbol colors and fix inconsistency

* fix typo in test and pick_bw_for_contrast allow color as str

* NEED CONFIRM: remove test for color consistency

* rename text_color to text_colors for consistency

* migrate asset maker for ptable-heatmap-atomic-mass

* migrate ptable-heatmap-percent

* fix colormap in percent mode

* fix log scale

* uncommit example figures

* add back exclude_elem

* minor clean up

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add tests to exclude_elems

* add overwrite_anomalies to handle nan/inf and tests

* remove finished TODO

* make exclude_elements in cline

* fix tile color and value when on_emtpy is show

* increate title font size

* remove cbar range to be consistent with previous setting

* WIP: migrate ptable_heatmap_ratio

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix arg in other examples

* make overwrite_anomalies in line

* fix color mapping in log mode

* update ratios plot

* optimize norm setter

* remove excluded elem to avoid metadata pollution and fix text color

* fix color mapping in percent/fraction mode

* break into submodules

* remove debug script

* Revert "break into submodules"

This reverts commit a43912e.

* fix merge main

* split ptable projector

* add ratio plot legends

* fix NaN display in ratio plotter

* migrate unit test for ratio plotter

* condense tile value, color and text_color generation

* WIP: merge tile symbol, color and value plot into add_heatmap_tiles

* first working version after tile refactor

* fix colorbar format

* rename args values_xx to value_xx

* add missing docstrings

* regenerate asset with new implementation

* finish ratio plotter migration

* fix wrong import of ptableprojector in unit test

* remove HMapPTableProjector into _projectors module

* add unit tests for generate_tile_value_colors

* fix mypy errors

* set NaN default color as lightgrey

* add deprecated warning for args

* remove inf_color to infty_color for backwards compatibility

* add more deprecated args

* WIP: migrate matbench_dielectric_eda.ipynb

* WIP: migrate mprester_ptable.ipynb

* set default on_empty to show and reset normalizer for values not present

* fix matbench_dielectric_eda ipynb

* ptable_heatmap now filter near zero values internally with a warning

* fix value type

* migrate mprester_ptable.ipynb

* migrate example boltztrap_mp

* migrate camd_2022

* set default value_fmt to .0f in log mode to be large-number friendly

* fix and migrate ricci_carrier_transport

* add debug tag

* migrate dielectric

* migrate expt_gap

* migrate other matbench

* simplify count_elements import

* fix unit test

* remove TODO tag

* drop not available structure

* add return_type for smoother migration to return figure

* reapply value filter for percent mode

* revert accidental revert during conflict merge

* fix return type in test and add test for return_type

* add debug tag

* make copy of df

* include ax example in asset

* remove ptabledata from_formula method

* change ptable_matplotlib.py doc str section format # -> ---

* rename Literal "AUTO"->"auto" for automatic text color and text format mode

* df_ptable.atomic_mass->df_ptable[Key.atomic_mass]

* rename HMapPTableProjector->HeatMapPTableProjector

* simplify assert isclose(a, b) -> assert a == b

* decrease filter_near_zero(tol=1e-4->1e-6)

* add isinstance(self.anomalies, dict) to remove need for typing.cast

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Janosh Riebesell <[email protected]>

Author: 3 people

.pre-commit-config.yaml

@@ -8,7 +8,7 @@ default_install_hook_types: [pre-commit, commit-msg]
 
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.5.3
+    rev: v0.5.4
     hooks:
       - id: ruff
         args: [--fix]
@@ -17,7 +17,7 @@ repos:
         types_or: [python, jupyter]
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.10.1
+    rev: v1.11.0
     hooks:
       - id: mypy
         additional_dependencies: [types-requests]

assets/ptable-heatmap-atomic-mass.svg

@@ -30,7 +30,7 @@
 import matplotlib.pyplot as plt
 from matminer.datasets import load_dataset
 
-from pymatviz import ptable_heatmap
+from pymatviz import count_elements, ptable_heatmap
 from pymatviz.enums import Key
 from pymatviz.io import save_fig
 
@@ -41,15 +41,20 @@
 
 
 # %%
-ax = ptable_heatmap(df_boltz[Key.formula], log=True)
-ax.set(title="Elements in BoltzTraP MP dataset")
-save_fig(ax, "boltztrap_mp-ptable-heatmap.pdf")
+fig = ptable_heatmap(
+    count_elements(df_boltz[Key.formula]), log=True, return_type="figure"
+)
+fig.suptitle("Elements in BoltzTraP MP dataset")
+save_fig(fig, "boltztrap_mp-ptable-heatmap.pdf")
 
 
 # %%
-ax = ptable_heatmap(df_boltz.sort_values("pf_n").tail(100)[Key.formula])
-ax.set(title="Elements of top 100 n-type powerfactors in BoltzTraP MP dataset")
-save_fig(ax, "boltztrap_mp-ptable-heatmap-top-100-nPF.pdf")
+fig = ptable_heatmap(
+    count_elements(df_boltz.sort_values("pf_n").tail(100)[Key.formula]),
+    return_type="figure",
+)
+fig.suptitle("Elements of top 100 n-type powerfactors in BoltzTraP MP dataset")
+save_fig(fig, "boltztrap_mp-ptable-heatmap-top-100-nPF.pdf")
 
 
 # %%

assets/ptable-heatmap-percent.svg

@@ -46,9 +46,9 @@
 
 # %%
 elem_counts = count_elements(df_camd.reduced_formula)
-ax = ptable_heatmap(elem_counts, log=True)
-ax.set(title="Elements in CAMD 2022 dataset")
-save_fig(ax, "camd-2022-ptable-heatmap.pdf")
+fig = ptable_heatmap(elem_counts, log=True, return_type="figure")
+fig.suptitle("Elements in CAMD 2022 dataset")
+save_fig(fig, "camd-2022-ptable-heatmap.pdf")
 
 
 # %%

assets/ptable-heatmap-ratio.svg

@@ -5,6 +5,7 @@
 from tqdm import tqdm
 
 from pymatviz import (
+    count_elements,
     crystal_sys_order,
     ptable_heatmap,
     ptable_heatmap_plotly,
@@ -51,9 +52,11 @@
 
 
 # %%
-ax = ptable_heatmap(df_diel[Key.formula], log=True)
-ax.set(title="Elemental prevalence in the Matbench dielectric dataset")
-save_fig(ax, "dielectric-ptable-heatmap.pdf")
+fig = ptable_heatmap(
+    count_elements(df_diel[Key.formula]), log=True, return_type="figure"
+)
+fig.suptitle("Elemental prevalence in the Matbench dielectric dataset")
+save_fig(fig, "dielectric-ptable-heatmap.pdf")
 
 
 # %%

assets/ptable-heatmap.svg

@@ -5,7 +5,7 @@
 from matminer.datasets import load_dataset
 from pymatgen.core import Composition
 
-from pymatviz import ptable_heatmap
+from pymatviz import count_elements, ptable_heatmap
 from pymatviz.enums import Key
 from pymatviz.io import save_fig
 
@@ -56,13 +56,13 @@ def mean_atomic_prop(comp: Composition, prop: str) -> float | None:
 
 
 # %%
-ax = ptable_heatmap(
-    df_gap.query("~composition.str.contains('Xe')")[Key.composition],
+fig = ptable_heatmap(
+    count_elements(df_gap.query("~composition.str.contains('Xe')")[Key.composition]),
     log=True,
-    text_color="black",
+    return_type="figure",
 )
-ax.set(title="Elements in Matbench experimental band gap dataset")
-save_fig(ax, "expt-gap-ptable-heatmap.pdf")
+fig.suptitle("Elements in Matbench experimental band gap dataset")
+save_fig(fig, "expt-gap-ptable-heatmap.pdf")
 
 
 # %%

examples/dataset_exploration/boltztrap_mp/boltztrap_mp-ptable-heatmap.pdf

@@ -2,7 +2,12 @@
 from matminer.datasets import load_dataset
 from tqdm import tqdm
 
-from pymatviz import ptable_heatmap, spacegroup_hist, spacegroup_sunburst
+from pymatviz import (
+    count_elements,
+    ptable_heatmap,
+    spacegroup_hist,
+    spacegroup_sunburst,
+)
 from pymatviz.enums import Key
 from pymatviz.io import save_fig
 
@@ -39,9 +44,9 @@
 df_2d[Key.volume] = [x.volume for x in df_2d[Key.structure]]
 df_2d[Key.formula] = [x.formula for x in df_2d[Key.structure]]
 
-ax = ptable_heatmap(df_2d[Key.formula], log=True)
-ax.set(title="Elemental prevalence in the Matbench Jarvis DFT 2D dataset")
-save_fig(ax, "jdft2d-ptable-heatmap.pdf")
+fig = ptable_heatmap(count_elements(df_2d[Key.formula]), log=True, return_type="figure")
+fig.suptitle("Elemental prevalence in the Matbench Jarvis DFT 2D dataset")
+save_fig(fig, "jdft2d-ptable-heatmap.pdf")
 
 
 # %%

examples/dataset_exploration/boltztrap_mp/explore_boltztrap_mp.py

@@ -20,6 +20,7 @@
 from tqdm import tqdm
 
 from pymatviz import (
+    count_elements,
     crystal_sys_order,
     ptable_heatmap,
     spacegroup_hist,
@@ -128,9 +129,11 @@ def has_isolated_atom(crystal: Structure, radius: float = 5) -> bool:
 # %%
 df_grvh[Key.formula] = df_grvh[Key.structure].map(lambda struct: struct.formula)
 
-ax = ptable_heatmap(df_grvh[Key.formula], log=True)
-ax.set(title="Elemental prevalence in the Matbench bulk/shear modulus datasets")
-save_fig(ax, "log_gvrh-ptable-heatmap.pdf")
+fig = ptable_heatmap(
+    count_elements(df_grvh[Key.formula]), log=True, return_type="figure"
+)
+fig.suptitle("Elemental prevalence in the Matbench bulk/shear modulus datasets")
+save_fig(fig, "log_gvrh-ptable-heatmap.pdf")
 
 
 # %%

examples/dataset_exploration/camd_2022/camd-2022-ptable-heatmap.pdf

@@ -1,7 +1,7 @@
 # %%
 from matminer.datasets import load_dataset
 
-from pymatviz import ptable_heatmap
+from pymatviz import count_elements, ptable_heatmap
 from pymatviz.enums import Key
 from pymatviz.io import save_fig
 
@@ -31,8 +31,11 @@
 # %%
 df_e_form[Key.formula] = df_e_form[Key.structure].map(lambda struct: struct.formula)
 
-
-# %%
-ax = ptable_heatmap(df_e_form[Key.formula], log=True)
-ax.set(title="Elemental prevalence in the Matbench formation energy dataset")
-save_fig(ax, "mp_e_form-ptable-heatmap.pdf")
+fig = ptable_heatmap(
+    count_elements(df_e_form[Key.formula]),
+    log=True,
+    value_kwargs={"fontsize": 10},
+    return_type="figure",
+)
+fig.suptitle("Elemental prevalence in the Matbench formation energy dataset")
+save_fig(fig, "mp_e_form-ptable-heatmap.pdf")

examples/dataset_exploration/camd_2022/explore_camd_2022.py

@@ -1,7 +1,7 @@
 # %%
 from matminer.datasets import load_dataset
 
-from pymatviz import ptable_heatmap
+from pymatviz import count_elements, ptable_heatmap
 from pymatviz.enums import Key
 from pymatviz.io import save_fig
 
@@ -32,11 +32,14 @@
 df_gap[Key.vol_per_atom] = df_gap[Key.volume] / df_gap[Key.n_sites]
 df_gap[Key.formula] = df_gap[Key.structure].map(lambda cryst: cryst.formula)
 
-
-# %%
-ax = ptable_heatmap(df_gap[Key.formula], log=True)
-ax.set(title="Elemental prevalence in the Matbench MP band gap dataset")
-save_fig(ax, "mp_gap-ptable-heatmap.pdf")
+fig = ptable_heatmap(
+    count_elements(df_gap[Key.formula]),
+    log=True,
+    value_kwargs={"fontsize": 10},
+    return_type="figure",
+)
+fig.suptitle("Elemental prevalence in the Matbench MP band gap dataset")
+save_fig(fig, "mp_gap-ptable-heatmap.pdf")
 
 
 # %%

examples/dataset_exploration/matbench/dielectric/dielectric-ptable-heatmap.pdf

@@ -2,7 +2,12 @@
 from matminer.datasets import load_dataset
 from tqdm import tqdm
 
-from pymatviz import plot_structure_2d, ptable_heatmap, spacegroup_sunburst
+from pymatviz import (
+    count_elements,
+    plot_structure_2d,
+    ptable_heatmap,
+    spacegroup_sunburst,
+)
 from pymatviz.enums import Key
 from pymatviz.io import save_fig
 from pymatviz.powerups import annotate_bars
@@ -50,9 +55,11 @@
 
 
 # %%
-ax = ptable_heatmap(df_perov[Key.formula], log=True)
-ax.set(title="Elements in Matbench Perovskites dataset")
-save_fig(ax, "perovskites-ptable-heatmap.pdf")
+fig = ptable_heatmap(
+    count_elements(df_perov[Key.formula]), log=True, return_type="figure"
+)
+fig.suptitle("Elements in Matbench Perovskites dataset")
+save_fig(fig, "perovskites-ptable-heatmap.pdf")
 
 
 # %%

examples/dataset_exploration/matbench/dielectric/explore_dielectric.py

@@ -19,7 +19,7 @@
 from matminer.datasets import load_dataset
 from tqdm import tqdm
 
-from pymatviz import ptable_heatmap, spacegroup_hist
+from pymatviz import count_elements, ptable_heatmap, spacegroup_hist
 from pymatviz.enums import Key
 from pymatviz.io import save_fig
 
@@ -42,9 +42,11 @@
 df_phonon[Key.formula] = df_phonon[Key.structure].map(lambda cryst: cryst.formula)
 df_phonon[Key.volume] = df_phonon[Key.structure].map(lambda cryst: cryst.volume)
 
-ax = ptable_heatmap(df_phonon[Key.formula], log=True)
-ax.set(title="Elemental prevalence in the Matbench phonons dataset")
-save_fig(ax, "phonons-ptable-heatmap.pdf")
+fig = ptable_heatmap(
+    count_elements(df_phonon[Key.formula]), log=True, return_type="figure"
+)
+fig.suptitle("Elemental prevalence in the Matbench phonons dataset")
+save_fig(fig, "phonons-ptable-heatmap.pdf")
 
 
 # %%

examples/dataset_exploration/matbench/expt_gap/explore_expt_gap.py

@@ -10,7 +10,7 @@
 # %%
 from matminer.datasets import load_dataset
 
-from pymatviz import ptable_heatmap
+from pymatviz import count_elements, ptable_heatmap
 from pymatviz.enums import Key
 from pymatviz.io import save_fig
 
@@ -25,6 +25,8 @@
 
 
 # %%
-ax = ptable_heatmap(df_steels[Key.composition], log=True)
-ax.set(title="Elemental prevalence in the Matbench steels dataset")
-save_fig(ax, "steels-ptable-heatmap.pdf")
+fig = ptable_heatmap(
+    count_elements(df_steels[Key.composition]), log=True, return_type="figure"
+)
+fig.suptitle("Elemental prevalence in the Matbench steels dataset")
+save_fig(fig, "steels-ptable-heatmap.pdf")

examples/dataset_exploration/matbench/expt_gap/expt-gap-ptable-heatmap.pdf

@@ -108,16 +108,16 @@
 r2scan_elem_counts = locals().get("r2scan_elem_counts")
 if r2scan_elem_counts is None:
     r2scan_elem_counts = count_elements(df_r2scan[Key.formula])
-ax = ptable_heatmap(r2scan_elem_counts)
+fig = ptable_heatmap(r2scan_elem_counts, return_type="figure")
 
-save_fig(ax, "r2scan-element-counts-ptable.pdf")
+save_fig(fig, "r2scan-element-counts-ptable.pdf")
 
 
 # %%
 pbe_elem_counts = locals().get("pbe_elem_counts")
 if pbe_elem_counts is None:
     pbe_elem_counts = count_elements(df_pbe[Key.formula])
-ax = ptable_heatmap(pbe_elem_counts)
+fig = ptable_heatmap(pbe_elem_counts, return_type="figure")
 
 
 # %% calculate per element energies
@@ -158,14 +158,15 @@
 
 
 # %% which elements have a higher share of missing r2scan data
-ax = ptable_heatmap(
+fig = ptable_heatmap(
     (pbe_elem_counts - r2scan_elem_counts) / pbe_elem_counts,
-    fmt=".1%",
-    cbar_fmt=".0%",
+    value_fmt=".1%",
+    cbar_label_fmt=".0%",
     cbar_title="Fraction of missing PBE calcs missing r2SCAN",
+    return_type="figure",
 )
 
-save_fig(ax, "ptable-has-pbe-but-no-r2scan.pdf")
+save_fig(fig, "ptable-has-pbe-but-no-r2scan.pdf")
 
 
 # %% per-elem mean abs magmoms
@@ -174,10 +175,13 @@
     for struct in df_r2scan[Key.structure]
 ).mean()
 
-ax = ptable_heatmap(
-    df_per_elem_magmoms, cbar_title=r"Mean |magmom| ($\mu_B$)", fmt=".1f"
+fig = ptable_heatmap(
+    df_per_elem_magmoms,
+    cbar_title=r"Mean |magmom| ($\mu_B$)",
+    value_fmt=".1f",
+    return_type="figure",
 )
-save_fig(ax, "magmoms-ptable.pdf")
+save_fig(fig, "magmoms-ptable.pdf")
 
 
 # %% spacegroup distribution

examples/dataset_exploration/matbench/jdft2d/explore_jdft2d.py

@@ -23,24 +23,27 @@
 from matminer.datasets import load_dataset
 from tqdm import tqdm
 
-from pymatviz import ptable_heatmap, spacegroup_hist
+from pymatviz import count_elements, ptable_heatmap, spacegroup_hist
 from pymatviz.enums import Key
 from pymatviz.io import save_fig
 
 
 # %%
 df_carrier = load_dataset("ricci_boltztrap_mp_tabular")
+df_carrier = df_carrier.dropna(subset=Key.structure)
 
-# getting space group symbols and numbers takes about 2 min
+# Getting space group symbols and numbers (take about 2 min)
 df_carrier[[Key.spg_symbol, Key.spg_num]] = [
     struct.get_space_group_info() for struct in tqdm(df_carrier[Key.structure])
 ]
 
 
 # %%
-ax = ptable_heatmap(df_carrier.pretty_formula.dropna(), log=True)
-ax.set(title="Elemental prevalence in the Ricci Carrier Transport dataset")
-save_fig(ax, "carrier-transport-ptable-heatmap.pdf")
+fig = ptable_heatmap(
+    count_elements(df_carrier.pretty_formula.dropna()), log=True, return_type="figure"
+)
+fig.suptitle("Elemental prevalence in the Ricci Carrier Transport dataset")
+save_fig(fig, "carrier-transport-ptable-heatmap.pdf")
 
 
 # %%