Prefer pymatviz interactive plotly version of periodic table heatmap if available (#3180)

* try from pymatviz import ptable_heatmap_plotly in periodic_table_heatmap()

* fix tests

* snake_case

* add keyword pymatviz: bool = True to disable periodic_table_heatmap deprecation warning

Author: janosh

pymatgen/symmetry/kpath.py

@@ -1824,7 +1824,7 @@ def _get_IRBZ(self, recip_point_group, W, key_points, face_center_inds, atol):
 
         g = np.dot(W.T, W)  # just using change of basis matrix rather than
         # Lattice.get_cartesian_coordinates for conciseness
-        ginv = np.linalg.inv(g)
+        g_inv = np.linalg.inv(g)
         D = np.linalg.det(W)
 
         primary_orientation = secondary_orientation = tertiary_orientation = None
@@ -1871,7 +1871,7 @@ def _get_IRBZ(self, recip_point_group, W, key_points, face_center_inds, atol):
                 face_center_found = False
                 for point in IRBZ_points:
                     if point[0] in face_center_inds:
-                        cross = D * np.dot(ginv, np.cross(ax, point[1]))
+                        cross = D * np.dot(g_inv, np.cross(ax, point[1]))
                         if not np.allclose(cross, 0, atol=atol):
                             rot_boundaries = [cross, -1 * np.dot(op, cross)]
                             face_center_found = True
@@ -1880,7 +1880,7 @@ def _get_IRBZ(self, recip_point_group, W, key_points, face_center_inds, atol):
                 if not face_center_found:
                     print("face center not found")
                     for point in IRBZ_points:
-                        cross = D * np.dot(ginv, np.cross(ax, point[1]))
+                        cross = D * np.dot(g_inv, np.cross(ax, point[1]))
                         if not np.allclose(cross, 0, atol=atol):
                             rot_boundaries = [cross, -1 * np.dot(op, cross)]
                             used_axes.append(ax)
@@ -1896,7 +1896,7 @@ def _get_IRBZ(self, recip_point_group, W, key_points, face_center_inds, atol):
                 face_center_found = False
                 for point in IRBZ_points:
                     if point[0] in face_center_inds:
-                        cross = D * np.dot(ginv, np.cross(ax, point[1]))
+                        cross = D * np.dot(g_inv, np.cross(ax, point[1]))
                         if not np.allclose(cross, 0, atol=atol):
                             rot_boundaries = [cross, np.dot(op, cross)]
                             face_center_found = True
@@ -1905,7 +1905,7 @@ def _get_IRBZ(self, recip_point_group, W, key_points, face_center_inds, atol):
                 if not face_center_found:
                     print("face center not found")
                     for point in IRBZ_points:
-                        cross = D * np.dot(ginv, np.cross(ax, point[1]))
+                        cross = D * np.dot(g_inv, np.cross(ax, point[1]))
                         if not np.allclose(cross, 0, atol=atol):
                             rot_boundaries = [cross, -1 * np.dot(op, cross)]
                             used_axes.append(ax)
@@ -1921,7 +1921,7 @@ def _get_IRBZ(self, recip_point_group, W, key_points, face_center_inds, atol):
                 face_center_found = False
                 for point in IRBZ_points:
                     if point[0] in face_center_inds:
-                        cross = D * np.dot(ginv, np.cross(ax, point[1]))
+                        cross = D * np.dot(g_inv, np.cross(ax, point[1]))
                         if not np.allclose(cross, 0, atol=atol):
                             rot_boundaries = [cross, -1 * np.dot(op, cross)]
                             face_center_found = True
@@ -1930,7 +1930,7 @@ def _get_IRBZ(self, recip_point_group, W, key_points, face_center_inds, atol):
                 if not face_center_found:
                     print("face center not found")
                     for point in IRBZ_points:
-                        cross = D * np.dot(ginv, np.cross(ax, point[1]))
+                        cross = D * np.dot(g_inv, np.cross(ax, point[1]))
                         if not np.allclose(cross, 0, atol=atol):
                             rot_boundaries = [cross, -1 * np.dot(op, cross)]
                             used_axes.append(ax)
@@ -1946,7 +1946,7 @@ def _get_IRBZ(self, recip_point_group, W, key_points, face_center_inds, atol):
                 face_center_found = False
                 for point in IRBZ_points:
                     if point[0] in face_center_inds:
-                        cross = D * np.dot(ginv, np.cross(ax, point[1]))
+                        cross = D * np.dot(g_inv, np.cross(ax, point[1]))
                         if not np.allclose(cross, 0, atol=atol):
                             rot_boundaries = [
                                 cross,
@@ -1958,7 +1958,7 @@ def _get_IRBZ(self, recip_point_group, W, key_points, face_center_inds, atol):
                 if not face_center_found:
                     print("face center not found")
                     for point in IRBZ_points:
-                        cross = D * np.dot(ginv, np.cross(ax, point[1]))
+                        cross = D * np.dot(g_inv, np.cross(ax, point[1]))
                         if not np.allclose(cross, 0, atol=atol):
                             rot_boundaries = [cross, -1 * np.dot(op, cross)]
                             used_axes.append(ax)
@@ -1974,7 +1974,7 @@ def _get_IRBZ(self, recip_point_group, W, key_points, face_center_inds, atol):
                 face_center_found = False
                 for point in IRBZ_points:
                     if point[0] in face_center_inds:
-                        cross = D * np.dot(ginv, np.cross(ax, point[1]))
+                        cross = D * np.dot(g_inv, np.cross(ax, point[1]))
                         if not np.allclose(cross, 0, atol=atol):
                             rot_boundaries = [cross, -1 * np.dot(op, cross)]
                             face_center_found = True
@@ -1983,7 +1983,7 @@ def _get_IRBZ(self, recip_point_group, W, key_points, face_center_inds, atol):
                 if not face_center_found:
                     print("face center not found")
                     for point in IRBZ_points:
-                        cross = D * np.dot(ginv, np.cross(ax, point[1]))
+                        cross = D * np.dot(g_inv, np.cross(ax, point[1]))
                         if not np.allclose(cross, 0, atol=atol):
                             rot_boundaries = [cross, -1 * np.dot(op, cross)]
                             used_axes.append(ax)
@@ -1999,7 +1999,7 @@ def _get_IRBZ(self, recip_point_group, W, key_points, face_center_inds, atol):
                 face_center_found = False
                 for point in IRBZ_points:
                     if point[0] in face_center_inds:
-                        cross = D * np.dot(ginv, np.cross(ax, point[1]))
+                        cross = D * np.dot(g_inv, np.cross(ax, point[1]))
                         if not np.allclose(cross, 0, atol=atol):
                             rot_boundaries = [cross, -1 * np.dot(op, cross)]
                             face_center_found = True
@@ -2008,7 +2008,7 @@ def _get_IRBZ(self, recip_point_group, W, key_points, face_center_inds, atol):
                 if not face_center_found:
                     print("face center not found")
                     for point in IRBZ_points:
-                        cross = D * np.dot(ginv, np.cross(ax, point[1]))
+                        cross = D * np.dot(g_inv, np.cross(ax, point[1]))
                         if not np.allclose(cross, 0, atol=atol):
                             rot_boundaries = [cross, -1 * np.dot(op, cross)]
                             used_axes.append(ax)

pymatgen/util/plotting.py

@@ -180,7 +180,7 @@ def _decide_fontcolor(rgba: tuple) -> Literal["black", "white"]:
 
 
 def periodic_table_heatmap(
-    elemental_data,
+    elemental_data=None,
     cbar_label="",
     cbar_label_size=14,
     show_plot=False,
@@ -193,39 +193,90 @@ def periodic_table_heatmap(
     symbol_fontsize=14,
     max_row=9,
     readable_fontcolor=False,
+    pymatviz: bool = True,
+    **kwargs,
 ):
     """
     A static method that generates a heat map overlaid on a periodic table.
 
     Args:
-         elemental_data (dict): A dictionary with the element as a key and a
+        elemental_data (dict): A dictionary with the element as a key and a
             value assigned to it, e.g. surface energy and frequency, etc.
             Elements missing in the elemental_data will be grey by default
             in the final table elemental_data={"Fe": 4.2, "O": 5.0}.
-         cbar_label (str): Label of the color bar. Default is "".
-         cbar_label_size (float): Font size for the color bar label. Default is 14.
-         cmap_range (tuple): Minimum and maximum value of the color map scale.
+        cbar_label (str): Label of the color bar. Default is "".
+        cbar_label_size (float): Font size for the color bar label. Default is 14.
+        cmap_range (tuple): Minimum and maximum value of the color map scale.
             If None, the color map will automatically scale to the range of the
             data.
-         show_plot (bool): Whether to show the heatmap. Default is False.
-         value_format (str): Formatting string to show values. If None, no value
+        show_plot (bool): Whether to show the heatmap. Default is False.
+        value_format (str): Formatting string to show values. If None, no value
             is shown. Example: "%.4f" shows float to four decimals.
-         value_fontsize (float): Font size for values. Default is 10.
-         symbol_fontsize (float): Font size for element symbols. Default is 14.
-         cmap (str): Color scheme of the heatmap. Default is 'YlOrRd'.
+        value_fontsize (float): Font size for values. Default is 10.
+        symbol_fontsize (float): Font size for element symbols. Default is 14.
+        cmap (str): Color scheme of the heatmap. Default is 'YlOrRd'.
             Refer to the matplotlib documentation for other options.
-         blank_color (str): Color assigned for the missing elements in
+        blank_color (str): Color assigned for the missing elements in
             elemental_data. Default is "grey".
-         edge_color (str): Color assigned for the edge of elements in the
+        edge_color (str): Color assigned for the edge of elements in the
             periodic table. Default is "white".
-         max_row (int): Maximum number of rows of the periodic table to be
+        max_row (int): Maximum number of rows of the periodic table to be
             shown. Default is 9, which means the periodic table heat map covers
             the standard 7 rows of the periodic table + 2 rows for the lanthanides
             and actinides. Use a value of max_row = 7 to exclude the lanthanides and
             actinides.
-         readable_fontcolor (bool): Whether to use readable font color depending
+        readable_fontcolor (bool): Whether to use readable font color depending
             on background color. Default is False.
+        pymatviz (bool): Whether to use pymatviz to generate the heatmap. Defaults to True.
+            See https://github.com/janosh/pymatviz.
+        kwargs: Passed to pymatviz.ptable_heatmap_plotly
     """
+    if pymatviz:
+        try:
+            from pymatviz import ptable_heatmap_plotly
+
+            if elemental_data:
+                kwargs.setdefault("elem_values", elemental_data)
+                print('elemental_data is deprecated, use elem_values={"Fe": 4.2, "O": 5.0} instead')
+            if cbar_label:
+                kwargs.setdefault("color_bar", {}).setdefault("title", cbar_label)
+                print('cbar_label is deprecated, use color_bar={"title": cbar_label} instead')
+            if cbar_label_size != 14:
+                kwargs.setdefault("color_bar", {}).setdefault("titlefont", {}).setdefault("size", cbar_label_size)
+                print('cbar_label_size is deprecated, use color_bar={"titlefont": {"size": cbar_label_size}} instead')
+            if cmap:
+                kwargs.setdefault("colorscale", cmap)
+                print("cmap is deprecated, use colorscale=cmap instead")
+            if cmap_range:
+                kwargs.setdefault("cscale_range", cmap_range)
+                print("cmap_range is deprecated, use cscale_range instead")
+            if value_format:
+                kwargs.setdefault("precision", value_format)
+                print("value_format is deprecated, use precision instead")
+            if blank_color != "grey":
+                print("blank_color is deprecated")
+            if edge_color != "white":
+                print("edge_color is deprecated")
+            if symbol_fontsize != 14:
+                print("symbol_fontsize is deprecated, use font_size instead")
+                kwargs.setdefault("font_size", symbol_fontsize)
+            if value_fontsize != 10:
+                print("value_fontsize is deprecated, use font_size instead")
+                kwargs.setdefault("font_size", value_fontsize)
+            if max_row != 9:
+                print("max_row is deprecated, use max_row instead")
+            if readable_fontcolor:
+                print("readable_fontcolor is deprecated, use font_colors instead, e.g. ('black', 'white')")
+
+            return ptable_heatmap_plotly(**kwargs)
+        except ImportError:
+            print(
+                "You're using a deprecated version of periodic_table_heatmap(). Consider `pip install pymatviz` which "
+                "offers an interactive plotly periodic table heatmap. You can keep calling this same function from "
+                "pymatgen. Some of the arguments have changed which you'll be warned about. "
+                "To disable this warning, pass pymatviz=False."
+            )
+
     # Convert primitive_elemental data in the form of numpy array for plotting.
     if cmap_range is not None:
         max_val = cmap_range[1]
@@ -286,7 +337,7 @@ def periodic_table_heatmap(
     ax.axis("off")
     ax.invert_yaxis()
 
-    # Set the scalermap for fontcolor
+    # Set the scalarmap for fontcolor
     norm = colors.Normalize(vmin=min_val, vmax=max_val)
     scalar_cmap = cm.ScalarMappable(norm=norm, cmap=cmap)
 
@@ -335,20 +386,20 @@ def format_formula(formula):
     """
     formatted_formula = ""
     number_format = ""
-    for i, s in enumerate(formula):
-        if s.isdigit():
+    for idx, char in enumerate(formula):
+        if char.isdigit():
             if not number_format:
                 number_format = "_{"
-            number_format += s
-            if i == len(formula) - 1:
+            number_format += char
+            if idx == len(formula) - 1:
                 number_format += "}"
                 formatted_formula += number_format
         else:
             if number_format:
                 number_format += "}"
                 formatted_formula += number_format
                 number_format = ""
-            formatted_formula += s
+            formatted_formula += char
 
     return f"${formatted_formula}$"
 

pymatgen/util/tests/test_plotting.py

@@ -5,12 +5,21 @@
 from pymatgen.util.plotting import periodic_table_heatmap, van_arkel_triangle
 from pymatgen.util.testing import PymatgenTest
 
+try:
+    import pymatviz
+    from plotly.graph_objects import Figure
+except ImportError:
+    pymatviz = None
+
 
 class FuncTestCase(PymatgenTest):
     def test_plot_periodic_heatmap(self):
         random_data = {"Te": 0.11083, "Au": 0.75756, "Th": 1.24758, "Ni": -2.0354}
         ret_val = periodic_table_heatmap(random_data)
-        assert ret_val is plt
+        if pymatviz:
+            assert isinstance(ret_val, Figure)
+        else:
+            assert ret_val is plt
 
         # Test all keywords
         periodic_table_heatmap(
@@ -21,7 +30,7 @@ def test_plot_periodic_heatmap(self):
             cmap_range=[0, 1],
             cbar_label="Hello World",
             blank_color="white",
-            value_format="%.4f",
+            value_format=".4f",
             edge_color="black",
             value_fontsize=12,
             symbol_fontsize=18,