Fixed import error of plotting function from pymatgen for py3.9+.

turbomoleio/core/control.py

@@ -36,7 +36,7 @@
 import numpy as np
 from monty.json import MSONable
 from monty.os import cd, makedirs_p
-from pymatgen.util.plotting import add_fig_kwargs, get_ax_fig_plt
+from pymatgen.util.plotting import add_fig_kwargs
 
 from turbomoleio.core.base import get_mol_and_indices_frozen
 from turbomoleio.core.datagroups import DataGroups
@@ -157,7 +157,14 @@ def plot(self, ax=None, **kwargs):
         Returns:
             A matplotlib Figure
         """
-        ax, fig, plt = get_ax_fig_plt(ax=ax)
+        try:
+            from pymatgen.util.plotting import get_ax_fig_plt
+
+            ax, fig, plt = get_ax_fig_plt(ax=ax)
+        except ImportError:
+            from pymatgen.util.plotting import get_ax_fig
+
+            ax, fig = get_ax_fig(ax=ax)
 
         ax.plot(range(self.n_steps), self.total, **kwargs)
         ax.set_xlabel("Steps")

turbomoleio/output/data.py

@@ -34,7 +34,7 @@
 from monty.json import MSONable
 from pymatgen.core.structure import Molecule
 from pymatgen.core.units import bohr_to_ang
-from pymatgen.util.plotting import add_fig_kwargs, get_ax_fig_plt
+from pymatgen.util.plotting import add_fig_kwargs
 
 from turbomoleio.output.parser import Parser
 
@@ -655,7 +655,14 @@ def plot_energies(self, ax=None, **kwargs):
         Returns:
             A matplotlib Figure.
         """
-        ax, fig, plt = get_ax_fig_plt(ax=ax)
+        try:
+            from pymatgen.util.plotting import get_ax_fig_plt
+
+            ax, fig, plt = get_ax_fig_plt(ax=ax)
+        except ImportError:
+            from pymatgen.util.plotting import get_ax_fig
+
+            ax, fig = get_ax_fig(ax=ax)
 
         ax.plot(range(len(self.energies)), self.energies, **kwargs)
         ax.set_xlabel("Steps")