Repository: janosh/pymatviz

Issue:
[Feature] [Self-assigned] Add heatmap plotter
Hi @janosh, I happen to have some plotting scripts for my MPhil research which includes plotting heatmaps. Shall I proceed and try to merge them into your package?

For example:
![headmap1](https://github.com/janosh/pymatviz/assets/80093591/4c4ad8a7-c284-4ece-ab5d-767b6e32a36a)

And this (I didn't find a good way to fix these overlapping labels though):
![heatmap2](https://github.com/janosh/pymatviz/assets/80093591/5678252c-ea80-4a70-bbdf-ff00309d1ecf)


Discussion:
DanielYang59: @janosh Ready for review. Anything improvement needed? Thanks!
DanielYang59: I notice `matplotlib` already has the `interpolation` method built in to allow interpolation of the plot, as stated [here](https://matplotlib.org/stable/gallery/images_contours_and_fields/interpolation_methods.html), by `ax.imshow(grid, interpolation=interp_method, cmap='viridis')`.

In this case I think this entire code looks pointless, being a simple wrapper over `imshow`. Therefore I wish to remove this entire function altogether. However the heatmap itself looks good to me (it might not make sense to leave a standalone heatmap tutorial because `pymatviz` is a visualization package instead of a tutorial for `matplotlib`?)...Do you have any suggestions @janosh? Thanks.


DanielYang59: @janosh Can you please review this as a minor format tweak and I would open another PR for the periodic table scatter plotter? Thanks!
DanielYang59: Thanks @janosh! I would also add guides and assert generations for `ptable_hist` as well in the following PR.

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 334461f..84ac855 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -7,7 +7,7 @@ default_install_hook_types: [pre-commit, commit-msg]
 
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.1.14
+    rev: v0.2.0
     hooks:
       - id: ruff
         args: [--fix]
diff --git a/examples/_generate_assets.py b/examples/_generate_assets.py
index 7839429..b35d512 100644
--- a/examples/_generate_assets.py
+++ b/examples/_generate_assets.py
@@ -37,7 +37,7 @@ from pymatviz.uncertainty import error_decay_with_uncert, qq_gaussian
 from pymatviz.utils import TEST_FILES, df_ptable
 
 
-# %%
+# %% 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")
@@ -192,7 +192,7 @@ ax = elements_hist(df_expt_gap.composition, keep_top=15, v_offset=1)
 save_and_compress_svg(ax, "hist-elemental-prevalence")
 
 
-# %%
+# %% Spacegroup histograms
 for backend in ("plotly", "matplotlib"):
     fig = spacegroup_hist(df_phonons.spg_num, backend=backend)  # type: ignore[arg-type]
     save_and_compress_svg(fig, f"spg-num-hist-{backend}")
@@ -245,10 +245,10 @@ ax = marchenko_pastur(corr_mat_rank_deficient, gamma=n_cols / n_rows)
 save_and_compress_svg(ax, "marchenko-pastur-rank-deficient")
 
 
-# %%
+# %% Plot Matbench phonon structures
 n_rows, n_cols = 3, 4
 fig, axs = plt.subplots(n_rows, n_cols, figsize=(3 * n_cols, 3 * n_rows))
-title = f"{len(axs.flat)} Matbench Phonons Structures"
+title = f"{len(axs.flat)} Matbench phonon structures"
 fig.suptitle(title, fontweight="bold", fontsize=20)
 
 for row, ax in zip(df_phonons.itertuples(), axs.flat):
@@ -266,14 +266,12 @@ fig.show()
 save_and_compress_svg(fig, "matbench-phonons-structures-2d")
 
 
-# %% plot some disordered structures in 2d
+# %% Plot some disordered structures in 2D
 disordered_structs = {
     mp_id: MPRester().get_structure_by_material_id(mp_id, conventional_unit_cell=True)
     for mp_id in ["mp-19017", "mp-12712"]
 }
 
-
-# %%
 for mp_id, struct in disordered_structs.items():
     for site in struct:  # disorder structures in-place
         if "Fe" in site.species:
@@ -316,7 +314,7 @@ fig.add_annotation(code_anno)
 save_and_compress_svg(fig, "sankey-from-2-df-cols-randints")
 
 
-# %% plot phonon bands and DOS
+# %% Plot phonon bands and DOS
 with zopen(f"{TEST_FILES}/mp-2758-Sr4Se4-pbe.json.lzma") as file:
     dft_dct = json.loads(file.read())
 with zopen(f"{TEST_FILES}/mp-2758-Sr4Se4-mace-y7uhwpje.json.lzma") as file:
diff --git a/pymatviz/__init__.py b/pymatviz/__init__.py
index e0a528b..953766a 100644
--- a/pymatviz/__init__.py
+++ b/pymatviz/__init__.py
@@ -1,4 +1,13 @@
-"""pymatviz: A Python package for visualizing materials data."""
+"""pymatviz: A Python package for visualizing materials data.
+
+Importing this module has side-effects that apply sensible (often, not always) global
+defaults settings for plotly and matplotlib like increasing font size, prettier
+axis labels (plotly only) and higher figure resolution (matplotlib only).
+
+To use it, simply import this module before generating any plots:
+
+import pymatviz
+"""
 
 from __future__ import annotations
 
@@ -121,16 +130,6 @@ px.defaults.labels |= {
 # use pio.kaleido.scope.mathjax = None
 
 
-"""
-Importing this module has side-effects that apply sensible (often, not always) global
-defaults settings for plotly and matplotlib like increasing font size, prettier
-axis labels (plotly only) and higher figure resolution (matplotlib only).
-
-To use it, simply import this module before generating any plots:
-
-import pymatviz
-"""
-
 plt.rc("font", size=16)
 plt.rc("savefig", bbox="tight", dpi=200)
 plt.rc("figure", dpi=200, titlesize=18)
diff --git a/pymatviz/io.py b/pymatviz/io.py
index 179b316..b9972ac 100644
--- a/pymatviz/io.py
+++ b/pymatviz/io.py
@@ -109,13 +109,13 @@ def save_fig(
         if path.lower().endswith(".svelte"):
             # insert {...$$props} into top-level div to be able to post-process and
             # style plotly figures from within Svelte files
-            with open(path) as file:
+            with open(path, encoding="utf-8") as file:
                 text = file.read().replace("<div>", "<div {...$$props}>", 1)
-            with open(path, "w") as file:
+            with open(path, mode="w", encoding="utf-8") as file:
                 # add trailing newline for pre-commit end-of-file commit hook
                 file.write(text + "\n")
         if style:
-            with open(path, "r+") as file:
+            with open(path, mode="r+", encoding="utf-8") as file:
                 # replace first '<div ' with '<div {style=} '
                 file.write(file.read().replace("<div ", f"<div {style=} ", 1))
     else:
@@ -354,7 +354,7 @@ def df_to_html_table(
     if styles is not None:
         # insert styles at end of closing </style> tag so they override default styles
         html = html.replace("</style>", f"{styles}\n</style>")
-    with open(file_path, "w") as file:
+    with open(file_path, mode="w", encoding="utf-8") as file:
         file.write(html)
 
 
diff --git a/pymatviz/phonons.py b/pymatviz/phonons.py
index 9e437ac..ad92b9a 100644
--- a/pymatviz/phonons.py
+++ b/pymatviz/phonons.py
@@ -138,10 +138,7 @@ def plot_phonon_bands(
                 f"Only {PhononBands.__name__} objects supported, got {type_name}"
             )
         branches = {pretty_sym_point(branch["name"]) for branch in bs.branches}
-        if not common_branches:
-            common_branches = branches
-        else:
-            common_branches = common_branches & branches
+        common_branches = common_branches & branches if common_branches else branches
 
     if not common_branches:
         raise ValueError("No common branches found among the band structures.")
diff --git a/pymatviz/sunburst.py b/pymatviz/sunburst.py
index 10dc822..6344665 100644
--- a/pymatviz/sunburst.py
+++ b/pymatviz/sunburst.py
@@ -55,7 +55,7 @@ def spacegroup_sunburst(
         df_spg_counts["crystal_sys"] = [
             crystal_sys_from_spg_num(x) for x in df_spg_counts.spacegroup
         ]
-    except ValueError:  # column must be strings of space group symbols
+    except (ValueError, TypeError):  # column must be strings of space group symbols
         df_spg_counts["crystal_sys"] = [
             SpaceGroup(x).crystal_system for x in df_spg_counts.spacegroup
         ]
diff --git a/pymatviz/uncertainty.py b/pymatviz/uncertainty.py
index d6dab29..e48f6f3 100644
--- a/pymatviz/uncertainty.py
+++ b/pymatviz/uncertainty.py
@@ -185,9 +185,7 @@ def get_std_decay(y_true: ArrayLike, y_pred: ArrayLike, y_std: ArrayLike) -> Arr
     # increasing count of the number of samples in each element of cumsum()
     n_inc = range(1, len(abs_err) + 1)
 
-    decay_by_std = abs_err[y_std_sort].cumsum() / n_inc
-
-    return decay_by_std  # noqa: RET504
+    return abs_err[y_std_sort].cumsum() / n_inc
 
 
 def error_decay_with_uncert(
diff --git a/pymatviz/utils.py b/pymatviz/utils.py
index 5e3aac4..c42710f 100644
--- a/pymatviz/utils.py
+++ b/pymatviz/utils.py
@@ -4,7 +4,7 @@ import ast
 from contextlib import contextmanager
 from functools import partial
 from os.path import dirname
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, get_args
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -28,8 +28,8 @@ if TYPE_CHECKING:
 PKG_DIR = dirname(__file__)
 ROOT = dirname(PKG_DIR)
 TEST_FILES = f"{ROOT}/tests/files"
-VALID_BACKENDS = ("matplotlib", "plotly")
 Backend = Literal["matplotlib", "plotly"]
+VALID_BACKENDS = mpl_key, plotly_key = get_args(Backend)
 
 
 df_ptable = pd.read_csv(f"{ROOT}/pymatviz/elements.csv", comment="#").set_index(
@@ -169,20 +169,17 @@ def annotate_bars(
             ) from exc
 
 
-metric_labels_mpl = {"R2": "$R^2$", "R2_adj": "$R^2_{adj}$"}
-metric_labels_plotly = {"R2": "R<sup>2</sup>", "R2_adj": "R<sup>2</sup><sub>adj</sub>"}
-
-
-def pretty_metric_label(key: str, backend: Literal["matplotlib", "plotly"]) -> str:
-    """Map metric keys to their pretty-printed labels."""
+def pretty_metric_label(key: str, backend: Backend) -> str:
+    """Map metric keys to their pretty labels."""
     if backend not in VALID_BACKENDS:
         raise ValueError(f"Unexpected {backend=}, must be one of {VALID_BACKENDS}")
 
-    if backend == "plotly":
-        label = metric_labels_plotly.get(key, key)
-    else:
-        label = metric_labels_mpl.get(key, key)
-    return label
+    symbol_mapping = {
+        "R2": {mpl_key: "$R^2$", plotly_key: "R<sup>2</sup>"},
+        "R2_adj": {mpl_key: "$R^2_{adj}$", plotly_key: "R<sup>2</sup><sub>adj</sub>"},
+    }
+
+    return symbol_mapping.get(key, {}).get(backend, key)
 
 
 def annotate_metrics(
@@ -305,21 +302,24 @@ CrystalSystem = Literal[
 
 def crystal_sys_from_spg_num(spg: int) -> CrystalSystem:
     """Get the crystal system for an international space group number."""
-    # not using isinstance(n, int) to allow 0-decimal floats
-    if not (spg == int(spg) and 0 < spg < 231):
-        raise ValueError(f"Invalid space group {spg}")
+    # ensure integer or float with no decimal part
+    if not (isinstance(spg, (int, float)) and spg == int(spg)):
+        raise TypeError(f"Expect integer space group number, got {spg=}")
+
+    if not (1 <= spg <= 230):
+        raise ValueError(f"Invalid space group number {spg}, must be 1 <= num <= 230")
 
-    if 0 < spg < 3:
+    if 1 <= spg <= 2:
         return "triclinic"
-    if spg < 16:
+    if spg <= 15:
         return "monoclinic"
-    if spg < 75:
+    if spg <= 74:
         return "orthorhombic"
-    if spg < 143:
+    if spg <= 142:
         return "tetragonal"
-    if spg < 168:
+    if spg <= 167:
         return "trigonal"
-    if spg < 195:
+    if spg <= 194:
         return "hexagonal"
     return "cubic"
 
diff --git a/pyproject.toml b/pyproject.toml
index 5a1c485..cfe2b2e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -80,8 +80,8 @@ no_implicit_optional = false
 
 [tool.ruff]
 target-version = "py39"
-select = ["ALL"]
-ignore = [
+lint.select = ["ALL"]
+lint.ignore = [
     "ANN101",
     "ANN401",
     "B028",   # No explicit stacklevel keyword argument found
@@ -109,11 +109,11 @@ ignore = [
     "TD",
     "TRY003",
 ]
-pydocstyle.convention = "google"
-isort.lines-after-imports = 2
-isort.split-on-trailing-comma = false
+lint.pydocstyle.convention = "google"
+lint.isort.lines-after-imports = 2
+lint.isort.split-on-trailing-comma = false
 
-[tool.ruff.per-file-ignores]
+[tool.ruff.lint.per-file-ignores]
 "__init__.py" = ["F401"]
 "tests/*" = ["D103", "S101"]
 "examples/*" = ["INP001", "T201"] # T201: print found


Test patch:
diff --git a/tests/test_ptable.py b/tests/test_ptable.py
index d303efb..3ed6c99 100644
--- a/tests/test_ptable.py
+++ b/tests/test_ptable.py
@@ -80,7 +80,7 @@ def test_count_elements_by_atomic_nums() -> None:
 @pytest.mark.parametrize("range_limits", [(-1, 10), (100, 200)])
 def test_count_elements_bad_atomic_nums(range_limits: tuple[int, int]) -> None:
     with pytest.raises(ValueError, match="assumed to represent atomic numbers"):
-        count_elements({idx: 0 for idx in range(*range_limits)})
+        count_elements(dict.fromkeys(range(*range_limits), 0))
 
     with pytest.raises(ValueError, match="assumed to represent atomic numbers"):
         # string and integer keys for atomic numbers should be handled equally
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 1bd95b9..4803c59 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -127,7 +127,18 @@ def test_crystal_sys_from_spg_num(spg_num: int, crystal_sys: CrystalSystem) -> N
 
 @pytest.mark.parametrize("spg", [-1, 0, 231])
 def test_crystal_sys_from_spg_num_invalid(spg: int) -> None:
-    with pytest.raises(ValueError, match=f"Invalid space group {spg}"):
+    with pytest.raises(
+        ValueError,
+        match=f"Invalid space group number {spg}, must be 1 <= num <= 230",
+    ):
+        crystal_sys_from_spg_num(spg)
+
+
+@pytest.mark.parametrize("spg", [1.2, "3"])
+def test_crystal_sys_from_spg_num_typeerror(spg: int) -> None:
+    with pytest.raises(
+        TypeError, match=f"Expect integer space group number, got {spg=}"
+    ):
         crystal_sys_from_spg_num(spg)