Update PPPlot slice plotting to use logging_utils.

Author: NikoOinonen

ppafm/PPPlot.py

@@ -1,11 +1,11 @@
 #!/usr/bin/python
 
-import sys
-
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.colors import LinearSegmentedColormap
 
+from .logging_utils import ProgressLogger
+
 # =========== defaults
 
 default_figsize = (8, 8)
@@ -58,11 +58,6 @@ def colorize_XY2RG(Xs, Ys):
     return c, vmax
 
 
-def write_plotting_slice(i):
-    sys.stdout.write(f"\r plotting slice # {i}")
-    sys.stdout.flush()
-
-
 # =========== plotting functions
 
 
@@ -85,8 +80,9 @@ def plotImages(
     V0=0.0,
     cbar_label=None,
 ):
+    progress_logger = ProgressLogger("plot", pre_message="Plotting slice # ")
     for ii, i in enumerate(slices):
-        write_plotting_slice(i)
+        progress_logger.print_message(i, is_last=ii == (len(slices) - 1))
         if symmetric_map:
             limit = max(abs(np.min(F[i] - V0)), abs(np.max(F[i] - V0)))
             vmin = -limit + V0
@@ -115,8 +111,9 @@ def plotImages(
 def plotVecFieldRG(
     prefix, dXs, dYs, slices, extent=None, zs=None, figsize=default_figsize, interpolation=default_interpolation, atoms=None, bonds=None, atomSize=default_atom_size
 ):
+    progress_logger = ProgressLogger("plot", pre_message="Plotting slice # ")
     for ii, i in enumerate(slices):
-        write_plotting_slice(i)
+        progress_logger.print_message(i, is_last=ii == (len(slices) - 1))
         plt.figure(figsize=(10, 10))
         HSBs, vmax = colorize_XY2RG(dXs[i], dYs[i])
         plt.imshow(HSBs, extent=extent, origin="lower", interpolation=interpolation)
@@ -151,8 +148,9 @@ def plotDistortions(
     bonds=None,
     atomSize=default_atom_size,
 ):
+    progress_logger = ProgressLogger("plot", pre_message="Plotting slice # ")
     for ii, i in enumerate(slices):
-        write_plotting_slice(i)
+        progress_logger.print_message(i, is_last=ii == (len(slices) - 1))
         plt.figure(figsize=figsize)
         plt.plot(X[i, ::by, ::by].flat, Y[i, ::by, ::by].flat, "r.", markersize=markersize)
         if BG is not None:
@@ -193,8 +191,9 @@ def plotArrows(
     bonds=None,
     atomSize=default_atom_size,
 ):
+    progress_logger = ProgressLogger("plot", pre_message="Plotting slice # ")
     for ii, i in enumerate(slices):
-        write_plotting_slice(i)
+        progress_logger.print_message(i, is_last=ii == (len(slices) - 1))
         plt.figure(figsize=figsize)
         plt.quiver(Xs[::by, ::by], Ys[::by, ::by], dX[::by, ::by], dY[::by, ::by], color="k", headlength=10, headwidth=10, scale=15)
         if BG is not None:

ppafm/data.py

@@ -47,9 +47,9 @@ def _common_parent(paths):
 def _extract_members(archive_handle, members, target_dir):
     progress_logger = ProgressLogger(logger_name=_progress_logger_name, pre_message=f"Extracting dataset to `{target_dir}`: ")
     for i, m in enumerate(members):
-        progress_logger.print_progress(i, 1, len(members))
+        progress_logger.print_percent(i, 1, len(members))
         archive_handle.extract(m, target_dir)
-    progress_logger.print_progress(len(members), 1, len(members))
+    progress_logger.print_percent(len(members), 1, len(members))
 
 
 def _extract_targz(archive_path, target_dir):
@@ -119,7 +119,7 @@ def download_dataset(name: str, target_dir: Union[Path, str]):
     with TemporaryDirectory() as temp_dir:
         temp_file = Path(temp_dir) / f"dataset_{name}"
         progress_logger = ProgressLogger(logger_name=_progress_logger_name, pre_message=f"Downloading dataset `{name}`: ")
-        _, response = urlretrieve(dataset_url, temp_file, reporthook=progress_logger.print_progress)
+        _, response = urlretrieve(dataset_url, temp_file, reporthook=progress_logger.print_percent)
         original_file_name = response.get_filename()
         target_dir.mkdir(exist_ok=True, parents=True)
         if original_file_name.endswith(".tar.gz"):

ppafm/logging_utils.py

@@ -1,7 +1,7 @@
 import logging
 import os
 import sys
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
 # Setup the root logger for the ppafm package. All other loggers will be derived from this one.
 _root_logger = logging.getLogger("ppafm")
@@ -113,7 +113,7 @@ def perf_log_enabled() -> bool:
 
 
 class ProgressLogger:
-    """Print gradual progress as a percent number."""
+    """Print gradual progress messages."""
 
     def __init__(self, logger_name: str = "progress", pre_message: str = ""):
 
@@ -143,7 +143,7 @@ def __init__(self, logger_name: str = "progress", pre_message: str = ""):
         self.pre_message = pre_message
         self._previous_percent = -self._percent_increment
 
-    def print_progress(self, block_num: int, block_size: int, total_size: int):
+    def print_percent(self, block_num: int, block_size: int, total_size: int):
         if total_size == -1:
             return
         current_size = block_num * block_size
@@ -158,3 +158,9 @@ def print_progress(self, block_num: int, block_size: int, total_size: int):
             if self._print_to_terminal:
                 msg += "\n"
             self.logger.info(msg)
+
+    def print_message(self, message: Any, is_last: bool):
+        message = f"{self.pre_message}{message}"
+        if is_last and self._print_to_terminal:
+            message += "\n"
+        self.logger.info(message)