Configure logging with support for verbose flags

pymbar/__init__.py

@@ -28,6 +28,10 @@
 __maintainer__ = "Levi N. Naden, Michael R. Shirts and John D. Chodera"
 __email__ = "[email protected],[email protected],[email protected]"
 
+from pymbar.logging_utils import setup_logging
+
+setup_logging()
+
 from pymbar import timeseries, testsystems, confidenceintervals
 from pymbar.mbar import MBAR
 from pymbar.other_estimators import bar, bar_zero, exp, exp_gauss

pymbar/logging_utils.py

@@ -0,0 +1,146 @@
+"""
+Utilities for our custom logging features
+"""
+from contextlib import contextmanager
+import logging
+
+
+class ForceEmitLogger(logging.Logger):
+    def log(self, level, msg, *args, **kwargs):
+        """
+        Do what ``logging.Logger`` does, but offer to override all
+        level thresholds if ``force_emit=True`` is passed to the
+        ``.log()`` call.
+        """
+        if not isinstance(level, int):
+            if logging.raiseExceptions:
+                raise TypeError("level must be an integer")
+            else:
+                return
+        if self.isEnabledFor(level):
+            self._log(level, msg, args, **kwargs)
+        elif kwargs.pop("force_emit", None):
+            with self.temporary_level(level):
+                self._log(level, msg, args, **kwargs)
+
+    def debug(self, msg, *args, **kwargs):
+        """
+        Do what ``logging.Logger`` does, but offer to override all
+        level thresholds if ``force_emit=True`` is passed to the
+        ``.log()`` call.
+        """
+        if self.isEnabledFor(logging.DEBUG):
+            self._log(logging.DEBUG, msg, args, **kwargs)
+        elif kwargs.pop("force_emit", None):
+            with self.temporary_level(logging.DEBUG):
+                self._log(logging.DEBUG, msg, args, **kwargs)
+
+    def info(self, msg, *args, **kwargs):
+        """
+        Do what ``logging.Logger`` does, but offer to override all
+        level thresholds if ``force_emit=True`` is passed to the
+        ``.log()`` call.
+        """
+        if self.isEnabledFor(logging.INFO):
+            self._log(logging.INFO, msg, args, **kwargs)
+        elif kwargs.pop("force_emit", None):
+            with self.temporary_level(logging.INFO):
+                self._log(logging.INFO, msg, args, **kwargs)
+
+    def warning(self, msg, *args, **kwargs):
+        """
+        Do what ``logging.Logger`` does, but offer to override all
+        level thresholds if ``force_emit=True`` is passed to the
+        ``.log()`` call.
+        """
+        if self.isEnabledFor(logging.WARNING):
+            self._log(logging.WARNING, msg, args, **kwargs)
+        elif kwargs.pop("force_emit", None):
+            with self.temporary_level(logging.INFO):
+                self._log(logging.INFO, msg, args, **kwargs)
+
+    def error(self, msg, *args, **kwargs):
+        """
+        Do what ``logging.Logger`` does, but offer to override all
+        level thresholds if ``force_emit=True`` is passed to the
+        ``.log()`` call.
+        """
+        if self.isEnabledFor(logging.ERROR):
+            self._log(logging.ERROR, msg, args, **kwargs)
+        elif kwargs.pop("force_emit", None):
+            with self.temporary_level(logging.ERROR):
+                self._log(logging.ERROR, msg, args, **kwargs)
+
+    def critical(self, msg, *args, **kwargs):
+        """
+        Do what ``logging.Logger`` does, but offer to override all
+        level thresholds if ``force_emit=True`` is passed to the
+        ``.log()`` call.
+        """
+        if self.isEnabledFor(logging.CRITICAL):
+            self._log(logging.CRITICAL, msg, args, **kwargs)
+        elif kwargs.pop("force_emit", None):
+            with self.temporary_level(logging.CRITICAL):
+                self._log(logging.CRITICAL, msg, args, **kwargs)
+
+    @contextmanager
+    def temporary_level(self, level=logging.INFO):
+        """
+        Context manager to change the logging level for a block.
+        """
+        logger_level = self.getEffectiveLevel()
+        handler_levels = [h.level for h in self.handlers]
+        self.setLevel(level)
+        try:
+            yield
+        finally:
+            self.setLevel(logger_level)
+            for handler, handler_level in zip(self.handlers, handler_levels):
+                handler.level = handler_level
+
+
+class PerLevelFormatter(logging.Formatter):
+    """
+    Adapted from https://stackoverflow.com/a/14859558
+    """
+
+    FORMATS = {
+        logging.ERROR: "ERROR! %(message)s",
+        logging.WARNING: "WARNING: %(message)s",
+        logging.INFO: "%(message)s",
+        logging.DEBUG: "Debug: %(message)s",
+    }
+
+    def __init__(self, fmt="%(levelName)d: %(message)s", datefmt=None, style="%", **kwargs):
+        super().__init__(fmt=fmt, datefmt=datefmt, style=style, **kwargs)
+
+    def format(self, record):
+
+        # Save the original format configured by the user
+        # when the logger formatter was instantiated
+        format_orig = self._style._fmt
+        self._style._fmt = self.FORMATS.get(record.levelno, self._style._fmt)
+        # Call the original formatter class to do the grunt work
+        result = super().format(record)
+        # Restore the original format configured by the user
+        self._style._fmt = format_orig
+
+        return result
+
+
+def setup_logging(level=logging.WARNING):
+    """
+    Basic configuration for the project, with a custom formatter
+    """
+    logging.setLoggerClass(ForceEmitLogger)
+    logger = logging.getLogger(__name__.split(".", 1)[0])
+    handler = logging.StreamHandler()
+    formatter = PerLevelFormatter()
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+    configure_logging_level(level)
+
+
+def configure_logging_level(level):
+    logger = logging.getLogger(__name__.split(".", 1)[0])
+    logger.setLevel(level)

pymbar/mbar.py

@@ -216,8 +216,9 @@ def __init__(
 
         K, N = np.shape(u_kn)
 
-        if verbose:
-            logger.info("K (total states) = {:d}, total samples = {:d}".format(K, N))
+        logger.info(
+            "K (total states) = {:d}, total samples = {:d}".format(K, N), force_emit=verbose
+        )
 
         if np.sum(self.N_k) != N:
             raise ParameterError(
@@ -285,18 +286,18 @@ def __init__(
                         logger.warning(dedent(msg[1:]))
 
         # Print number of samples from each state.
-        if self.verbose:
-            logger.info("N_k = ")
-            logger.info(self.N_k)
+        logger.info("N_k = ", force_emit=verbose)
+        logger.info(self.N_k, force_emit=verbose)
 
         # Determine list of k indices for which N_k != 0
         self.states_with_samples = np.where(self.N_k != 0)[0]
         self.states_with_samples = self.states_with_samples.astype(np.int64)
 
         # Number of states with samples.
         self.K_nonzero = self.states_with_samples.size
-        if verbose:
-            logger.info("There are {:d} states with samples.".format(self.K_nonzero))
+        logger.info(
+            "There are {:d} states with samples.".format(self.K_nonzero), force_emit=verbose
+        )
 
         # Initialize estimate of relative dimensionless free energy of each state to zero.
         # Note that f_k[0] will be constrained to be zero throughout.
@@ -306,8 +307,7 @@ def __init__(
         # If an initial guess of the relative dimensionless free energies is
         # specified, start with that.
         if initial_f_k is not None:
-            if self.verbose:
-                logger.info("Initializing f_k with provided initial guess.")
+            logger.info("Initializing f_k with provided initial guess.", force_emit=verbose)
             # Cast to np array.
             initial_f_k = np.array(initial_f_k, dtype=np.float64)
             # Check shape
@@ -317,20 +317,19 @@ def __init__(
                 )
             # Initialize f_k with provided guess.
             self.f_k = initial_f_k
-            if self.verbose:
-                logger.info(self.f_k)
+            logger.info(self.f_k, force_emit=verbose)
             # Shift all free energies such that f_0 = 0.
             self.f_k[:] = self.f_k[:] - self.f_k[0]
         else:
             # Initialize estimate of relative dimensionless free energies.
             self._initializeFreeEnergies(verbose, method=initialize)
 
-            if self.verbose:
-                logger.info(
-                    "Initial dimensionless free energies with method {:s}".format(initialize)
-                )
-                logger.info("f_k = ")
-                logger.info(self.f_k)
+            logger.info(
+                "Initial dimensionless free energies with method {:s}".format(initialize),
+                force_emit=verbose,
+            )
+            logger.info("f_k = ", force_emit=verbose)
+            logger.info(self.f_k, force_emit=verbose)
 
         if solver_protocol is None:
             solver_protocol = ({"method": None},)
@@ -348,13 +347,10 @@ def __init__(
         self.Log_W_nk = mbar_solvers.mbar_log_W_nk(self.u_kn, self.N_k, self.f_k)
 
         # Print final dimensionless free energies.
-        if self.verbose:
-            logger.info("Final dimensionless free energies")
-            logger.info("f_k = ")
-            logger.info(self.f_k)
-
-        if self.verbose:
-            logger.info("MBAR initialization complete.")
+        logger.info("Final dimensionless free energies", force_emit=verbose)
+        logger.info("f_k = ", force_emit=verbose)
+        logger.info(self.f_k, force_emit=verbose)
+        logger.info("MBAR initialization complete.", force_emit=verbose)
 
     @property
     def W_nk(self):
@@ -440,15 +436,16 @@ def compute_effective_sample_number(self, verbose=False):
         for k in range(self.K):
             w = np.exp(self.Log_W_nk[:, k])
             N_eff[k] = 1 / np.sum(w ** 2)
-            if verbose:
-                logger.info(
-                    "Effective number of sample in state {:d} is {:10.3f}".format(k, N_eff[k])
-                )
-                logger.info(
-                    "Efficiency for state {:d} is {:6f}/{:d} = {:10.4f}".format(
-                        k, N_eff[k], len(w), N_eff[k] / len(w)
-                    )
-                )
+            logger.info(
+                "Effective number of sample in state {:d} is {:10.3f}".format(k, N_eff[k]),
+                force_emit=verbose,
+            )
+            logger.info(
+                "Efficiency for state {:d} is {:6f}/{:d} = {:10.4f}".format(
+                    k, N_eff[k], len(w), N_eff[k] / len(w)
+                ),
+                force_emit=verbose,
+            )
 
         return N_eff
 
@@ -1359,8 +1356,7 @@ def compute_entropy_and_enthalpy(
         >>> results = mbar.compute_entropy_and_enthalpy()
 
         """
-        if verbose:
-            logger.info("Computing average energy and entropy by MBAR.")
+        logger.info("Computing average energy and entropy by MBAR.", force_emit=verbose)
 
         dims = len(np.shape(u_kn))
         if dims == 3:
@@ -1640,19 +1636,17 @@ def _initializeFreeEnergies(self, verbose=False, method="zeros"):
             * mean-reduced-potential: the mean reduced potential is used
 
         """
-
         if method == "zeros":
             # Use zeros for initial free energies.
-            if verbose:
-                logger.info("Initializing free energies to zero.")
+            logger.info("Initializing free energies to zero.", force_emit=verbose)
             self.f_k[:] = 0.0
         elif method == "mean-reduced-potential":
             # Compute initial guess at free energies from the mean reduced
             # potential from each state
-            if verbose:
-                logger.info(
-                    "Initializing free energies with mean reduced potential for each state."
-                )
+            logger.info(
+                "Initializing free energies with mean reduced potential for each state.",
+                force_emit=verbose,
+            )
             means = np.zeros([self.K], float)
             for k in self.states_with_samples:
                 means[k] = self.u_kn[k, 0 : self.N_k[k]].mean()