Saving work

abipy/dynamics/hist.py

@@ -374,7 +374,7 @@ def plot_ax(self, ax, what, fontsize=8, **kwargs) -> None:
             ax.set_ylabel('F stats (eV/A)')
 
         else:
-            raise ValueError("Invalid value for what: `%s`" % str(what))
+            raise ValueError(f"Invalid value for {what=}")
 
         ax.set_xlabel('Step')
         ax.grid(True)

abipy/ml/aseml.py

@@ -50,7 +50,7 @@
 from abipy.tools.serialization import HasPickleIO
 from abipy.abio.enums import StrEnum, EnumMixin
 from abipy.tools.plotting import (set_axlims, add_fig_kwargs, get_ax_fig_plt, get_axarray_fig_plt, set_grid_legend,
-    set_visible, set_ax_xylabels)
+    set_visible, set_ax_xylabels, linear_fit_ax)
 
 
 ###################
@@ -331,20 +331,6 @@ def diff_stats(xs, ys):
     )
 
 
-def linear_fit_ax(ax, xs, ys, fontsize, with_label=True, with_ideal_line=False) -> tuple[float]:
-    """
-    """
-    from scipy.stats import linregress
-    result = linregress(xs, ys)
-    label = r"Linear fit $\alpha={:.2f}$, $r^2$={:.2f}".format(result.slope, result.rvalue**2)
-    ax.plot(xs, result.slope*xs + result.intercept, 'r', label=label if with_label else None)
-    if with_ideal_line:
-        # Plot y = x line
-        ax.plot([xs[0], xs[-1]], [ys[0], ys[-1]], color='k', linestyle='-',
-                linewidth=1, label='Ideal' if with_label else None)
-    return result
-
-
 def make_square_axes(ax_mat):
     """
     Make an axes square in screen units.

abipy/tools/context_managers.py

@@ -17,7 +17,7 @@ class Timer:
         with Timer("Timing code section"):
             do_stuff()
 
-    or 
+    or
 
         with Timer(header=f"Begin ABINIT", footer="ABINIT GS") as timer:
             do_stuff()
@@ -40,7 +40,7 @@ def __exit__(self, type, value, traceback):
         self.time = perf_counter() - self.time
         self.readout = f'Time: {self.time:.3f} seconds'
         if self.footer is not None:
-            msg = f'{self.footer} completed in {self.time:.3f} seconds.'
+            msg = f'{self.footer} completed in {self.time:.3f} seconds'.lstrip()
             print(msg, file=self.file)
 
 

abipy/tools/iotools.py

@@ -25,6 +25,18 @@ def make_executable(filepath: str) -> None:
     os.chmod(filepath, mode)
 
 
+def try_files(filepaths: list) -> Path:
+    """
+    Return the first existent file in filepaths
+    or raise RuntimeError.
+    """
+    for path in filepaths:
+        path = Path(str(path))
+        if path.exists(): return path
+
+    raise RuntimeError("Cannot find {filepaths=}")
+
+
 def yaml_safe_load(string: str) -> Any:
     """Load Yaml string"""
     return yaml.YAML(typ='safe', pure=True).load(string)

abipy/tools/plotting.py

@@ -464,6 +464,25 @@ def plot_xy_with_hue(data: pd.DataFrame, x: str, y: str, hue: str,
     return fig
 
 
+def linear_fit_ax(ax, xs, ys, fontsize, with_label=True, with_ideal_line=False, **kwargs) -> tuple[float]:
+    """
+    Calculate a linear least-squares regression for two sets of measurements.
+    kwargs are passed to ax.plot.
+    """
+    from scipy.stats import linregress
+    fit = linregress(xs, ys)
+    label = r"Linear fit $\alpha={:.2f}$, $r^2$={:.2f}".format(fit.slope, fit.rvalue**2)
+    if "color" not in kwargs:
+        kwargs["color"] = "r"
+
+    ax.plot(xs, fit.slope*xs + fit.intercept, label=label if with_label else None, **kwargs)
+    if with_ideal_line:
+        # Plot y = x line
+        ax.plot([xs[0], xs[-1]], [ys[0], ys[-1]], color='k', linestyle='-',
+                linewidth=1, label='Ideal' if with_label else None)
+    return fit
+
+
 @add_fig_kwargs
 def plot_array(array, color_map=None, cplx_mode="abs", **kwargs) -> Figure:
     """

abipy/tools/serialization.py

@@ -15,6 +15,7 @@
 from pathlib import Path
 from monty.json import MontyDecoder, MontyEncoder
 from pymatgen.core.periodic_table import Element
+from abipy.tools.context_managers import Timer
 
 
 def pmg_serialize(method):
@@ -145,15 +146,16 @@ class HasPickleIO:
 
     @classmethod
     def pickle_load(cls, workdir, basename=None):
-        """Reconstruct the object from a pickle file located in workdir."""
-        #if workdir is None: workdir = os.getcwd()
+        """
+        Reconstruct the object from a pickle file located in workdir.
+        """
+
         filepath = Path(workdir) / f"{cls.__name__}.pickle" if basename is None else Path(workdir) / basename
-        with open(filepath, "rb") as fh:
+        with open(filepath, "rb") as fh, Timer(header=f"Reconstructing {cls.__name__} instance from file: {str(filepath)}", footer="") as timer:
             return pickle.load(fh)
 
     def pickle_dump(self, workdir, basename=None) -> None:
         """Write pickle file."""
-        #if workdir is None: workdir = os.getcwd()
         filepath = Path(workdir) / f"{self.__class__.__name__}.pickle" if basename is None else Path(workdir) / basename
-        with open(filepath, "wb") as fh:
+        with open(filepath, "wb") as fh, Timer(header=f"Saving {self.__class__.__name__} instance to file: {str(filepath)}", footer="") as timer:
             pickle.dump(self, fh)