xarpes 0.2.4__py3-none-any.whl → 0.6.0__py3-none-any.whl

xarpes/plotting.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2024 xARPES Developers
+# Copyright (C) 2025 xARPES Developers
 # This program is free software under the terms of the GNU GPLv3 license.
 
 # get_ax_fig_plt and add_fig_kwargs originate from pymatgen/util/plotting.py.
@@ -11,21 +11,7 @@
 
 """Functions related to plotting."""
 
-from functools import wraps
 import matplotlib.pyplot as plt
-import matplotlib as mpl
-
-def plot_settings(name='default'):
-    mpl.rc('xtick', labelsize=10, direction='in')
-    mpl.rc('ytick', labelsize=10, direction='in')
-    lw = dict(default=2.0, large=4.0)[name]
-    mpl.rcParams['lines.linewidth'] = lw
-    mpl.rcParams['lines.markersize'] = 3
-    mpl.rcParams['xtick.major.size'] = 4
-    mpl.rcParams['xtick.minor.size'] = 2
-    mpl.rcParams['xtick.major.width'] = 0.8
-    mpl.rcParams.update({'font.size': 16})
-    mpl.use('Qt5Agg') # Backend for showing plots in terminal
 
 def get_ax_fig_plt(ax=None, **kwargs):
     r"""Helper function used in plot functions supporting an optional `Axes`
@@ -58,13 +44,16 @@ def get_ax_fig_plt(ax=None, **kwargs):
 
     return ax, fig, plt
 
+
 def add_fig_kwargs(func):
     """Decorator that adds keyword arguments for functions returning matplotlib
     figures.
 
-    The function should return either a matplotlib figure or None to signal
-    some sort of error/unexpected event.
+    The function should return either a matplotlib figure or a tuple, where the
+    first element is a matplotlib figure, or None to signal some sort of
+    error/unexpected event.
     """
+    from functools import wraps
     @wraps(func)
     def wrapper(*args, **kwargs):
         # pop the kwds used by the decorator.
@@ -75,14 +64,26 @@ def add_fig_kwargs(func):
         tight_layout = kwargs.pop('tight_layout', False)
         ax_grid = kwargs.pop('ax_grid', None)
         ax_annotate = kwargs.pop('ax_annotate', None)
-        fig_close = kwargs.pop('fig_close', False)
-
-        # Call func and return immediately if None is returned.
-        fig = func(*args, **kwargs)
+        fig_close = kwargs.pop('fig_close', True)
+        
+        import string
+
+        # Call the original function
+        result = func(*args, **kwargs)
+
+        # Determine if result is a figure or tuple with first element as figure
+        if isinstance(result, tuple):
+            fig = result[0]
+            rest = result[1:]
+        else:
+            fig = result
+            rest = None
+
+        # Return immediately if no figure is returned
         if fig is None:
-            return fig
+            return result
 
-        # Operate on matplotlib figure.
+        # Operate on the matplotlib figure
         if title is not None:
             fig.suptitle(title)
 
@@ -95,9 +96,10 @@ def add_fig_kwargs(func):
                 ax.grid(bool(ax_grid))
 
         if ax_annotate:
-            tags = ascii_letters
+            tags = string.ascii_letters
             if len(fig.axes) > len(tags):
-                tags = (1 + len(ascii_letters) // len(fig.axes)) * ascii_letters
+                tags = (1 + len(string.ascii_letters) // len(fig.axes)) * \
+                string.ascii_letters
             for ax, tag in zip(fig.axes, tags):
                 ax.annotate(f'({tag})', xy=(0.05, 0.95),
                             xycoords='axes fraction')
@@ -106,10 +108,8 @@ def add_fig_kwargs(func):
             try:
                 fig.tight_layout()
             except Exception as exc:
-                # For some unknown reason, this problem shows up only on travis.
-                # https://stackoverflow.com/questions/22708888/valueerror-when-using-matplotlib-tight-layout
-                print('Ignoring Exception raised by fig.tight_layout\n',
-                      str(exc))
+                print('Ignoring Exception raised by fig.tight_layout ' +
+                      '\n', str(exc))
 
         if savefig:
             fig.savefig(savefig)
@@ -119,12 +119,16 @@ def add_fig_kwargs(func):
         if fig_close:
             plt.close(fig=fig)
 
-        return fig
+        # Reassemble the tuple if necessary and return
+        if rest is not None:
+            return (fig, *rest)
+        else:
+            return fig
 
     # Add docstring to the decorated method.
     doc_str = """\n\n
 
-        notes
+        Notes
         -----
 
         Keyword arguments controlling the display of the figure:
@@ -141,18 +145,16 @@ def add_fig_kwargs(func):
         ax_grid           True (False) to add (remove) grid from all axes in
                           fig.
                           Default: None i.e. fig is left unchanged.
-        ax_annotate       Add labels to  subplots e.g. (a), (b).
+        ax_annotate       Add labels to subplots e.g. (a), (b).
                           Default: False
-        fig_close         Close figure. Default: False.
+        fig_close         Close figure. Default: True.
         ================  ====================================================
 
 """
 
     if wrapper.__doc__ is not None:
-        # Add s at the end of the docstring.
         wrapper.__doc__ += f'\n{doc_str}'
     else:
-        # Use s
         wrapper.__doc__ = doc_str
 
     return wrapper