fix tests

Author: janosh

assets/scripts/ptable_plotly/ptable_heatmap_splits_plotly.py

@@ -74,6 +74,7 @@ def elem_color_scale(element: str, _val: float, split_idx: int) -> str:
 )
 fig.layout.title.update(text=title, x=0.4, y=0.8)
 fig.show()
+pmv.io.save_and_compress_svg(fig, "ptable-heatmap-splits-plotly-3-color-schemes")
 
 
 # %% Visualize multiple element color schemes on a split periodic table heatmap
@@ -89,3 +90,4 @@ def elem_color_scale(element: str, _val: float, split_idx: int) -> str:
 title = "<b>Element color schemes</b><br>left: VESTA, right: ALLOY"
 fig.layout.title.update(text=title, x=0.4, y=0.8)
 fig.show()
+pmv.io.save_and_compress_svg(fig, "ptable-heatmap-splits-plotly-2-color-schemes")

pymatviz/coordination/plotly.py

@@ -14,7 +14,7 @@
 from plotly.subplots import make_subplots
 from pymatgen.analysis.local_env import NearNeighbors
 
-from pymatviz.colors import ELEM_COLORS_JMOL, ELEM_COLORS_VESTA
+from pymatviz.colors import ELEM_COLORS_JMOL
 from pymatviz.coordination.helpers import (
     CnSplitMode,
     calculate_average_cn,
@@ -145,14 +145,10 @@ def coordination_hist(
     if isinstance(element_color_scheme, dict):
         # Merge custom colors with default Jmol colors to get a complete color scheme
         element_colors = ELEM_COLORS_JMOL | element_color_scheme
-    elif element_color_scheme == ElemColorScheme.jmol:
-        element_colors = ELEM_COLORS_JMOL
-    elif element_color_scheme == ElemColorScheme.vesta:
-        element_colors = ELEM_COLORS_VESTA
-    elif isinstance(element_color_scheme, dict):
-        element_colors = element_color_scheme
+    elif isinstance(element_color_scheme, ElemColorScheme):
+        element_colors = element_color_scheme.color_map
     else:
-        raise ValueError(
+        raise TypeError(
             f"Invalid {element_color_scheme=}. Must be {', '.join(ElemColorScheme)} "
             f"or a custom dict."
         )
@@ -403,12 +399,10 @@ def coordination_vs_cutoff_line(
 
     if isinstance(element_color_scheme, dict):
         element_colors = ELEM_COLORS_JMOL | element_color_scheme
-    elif element_color_scheme == ElemColorScheme.jmol:
-        element_colors = ELEM_COLORS_JMOL
-    elif element_color_scheme == ElemColorScheme.vesta:
-        element_colors = ELEM_COLORS_VESTA
+    elif isinstance(element_color_scheme, ElemColorScheme):
+        element_colors = element_color_scheme.color_map
     else:
-        raise ValueError(
+        raise TypeError(
             f"Invalid {element_color_scheme=}. Must be {', '.join(ElemColorScheme)} "
             "or a custom dict."
         )

pymatviz/enums.py

@@ -14,6 +14,8 @@
 
     from typing_extensions import Self
 
+    from pymatviz.typing import RgbColorType
+
 # TODO: remove following definition of StrEnum once Python 3.11+
 if sys.version_info >= (3, 11):
     from enum import StrEnum
@@ -798,6 +800,13 @@ class ElemColorScheme(LabelEnum):
     # custom made for pymatviz
     alloy = "alloy", "Alloy", "High-contrast color scheme optimized for metal alloys"
 
+    @property
+    def color_map(self) -> dict[str, RgbColorType]:
+        """Return map from element symbol to color."""
+        import pymatviz.colors as pmv_colors
+
+        return getattr(pmv_colors, f"ELEM_COLORS_{self.value.upper()}")
+
 
 @unique
 class SiteCoords(LabelEnum):

pymatviz/structure_viz/mpl.py

@@ -20,7 +20,6 @@
 from pymatgen.core import Structure
 from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
 
-from pymatviz.colors import ELEM_COLORS_JMOL, ELEM_COLORS_VESTA
 from pymatviz.enums import ElemColorScheme, Key
 from pymatviz.structure_viz.helpers import (
     NO_SYM_MSG,
@@ -202,13 +201,11 @@ class used to plot chemical bonds. Allowed are edgecolor, facecolor, color,
                 warnings.warn(NO_SYM_MSG, UserWarning, stacklevel=2)
 
         # Get default colors
-        if str(elem_colors) == str(ElemColorScheme.jmol):
-            elem_colors = ELEM_COLORS_JMOL
-        elif str(elem_colors) == str(ElemColorScheme.vesta):
-            elem_colors = ELEM_COLORS_VESTA
+        if isinstance(elem_colors, ElemColorScheme):
+            elem_colors = elem_colors.color_map
         elif not isinstance(elem_colors, dict):
             valid_color_schemes = "', '".join(ElemColorScheme)
-            raise ValueError(
+            raise TypeError(
                 f"colors must be a dict or one of ('{valid_color_schemes}')"
             )
 

tests/coordination/test_plotly.py

@@ -305,7 +305,7 @@ def test_coordination_hist_color_schemes(structures: Sequence[Structure]) -> Non
 
 def test_coordination_hist_invalid_elem_colors(structures: Sequence[Structure]) -> None:
     """Test invalid color scheme handling."""
-    with pytest.raises(ValueError, match="Invalid.*element_color_scheme"):
+    with pytest.raises(TypeError, match="Invalid.*element_color_scheme"):
         coordination_hist(structures[0], element_color_scheme="invalid")  # type: ignore[arg-type]
 
 

tests/structure_viz/test_structure_viz_mpl.py

@@ -10,7 +10,6 @@
 from pymatgen.core import Structure
 
 import pymatviz as pmv
-from pymatviz.colors import ELEM_COLORS_JMOL, ELEM_COLORS_VESTA
 from pymatviz.enums import ElemColorScheme, Key
 
 
@@ -163,23 +162,16 @@ def subplot_title(struct: Structure, key: str | int) -> str:
 
 def test_structure_2d_color_warning() -> None:
     # Copernicium is not in the default color scheme
-    elem_symbol = "Cn"
+    elem_symbol = "Fl"
     struct = Structure(np.eye(3) * 5, [elem_symbol] * 2, coords=COORDS)
     fallback_color = "gray"
 
     for elem_colors in ElemColorScheme:
-        if elem_colors == ElemColorScheme.jmol:
-            elem_color_symbols = ", ".join(ELEM_COLORS_JMOL)
-        elif elem_colors == ElemColorScheme.vesta:
-            elem_color_symbols = ", ".join(ELEM_COLORS_VESTA)
-        else:
-            raise ValueError(f"Unexpected {elem_colors=}")
-
         with pytest.warns(
             UserWarning,
             match=f"{elem_symbol=} not in elem_colors, using "
             f"{fallback_color=}\nelement color palette specifies the "
-            f"following elements: {elem_color_symbols}",
+            f"following elements: {', '.join(elem_colors.color_map)}",
         ):
             pmv.structure_2d(struct, elem_colors=elem_colors)