Added Loggers for optimizer, routine and line search removed tqdm

Author: MPMPMPMPMPMPMP

varipeps/__init__.py

@@ -20,11 +20,4 @@
 
 jax_config.update("jax_enable_x64", True)
 
-from tqdm_loggable.tqdm_logging import tqdm_logging
-import datetime
-
-tqdm_logging.set_log_rate(datetime.timedelta(seconds=60))
-
-del datetime
-del tqdm_logging
 del jax_config

varipeps/config.py

@@ -1,11 +1,12 @@
 from dataclasses import dataclass
 from enum import Enum, IntEnum, auto, unique
+from typing import TypeVar, Tuple, Any, Type, NoReturn
+import logging
 
 import numpy as np
 
 from jax.tree_util import register_pytree_node_class
 
-from typing import TypeVar, Tuple, Any, Type, NoReturn
 
 T_VariPEPS_Config = TypeVar("T_VariPEPS_Config", bound="VariPEPS_Config")
 
@@ -54,6 +55,15 @@ class Slurm_Restart_Mode(IntEnum):
     AUTOMATIC_RESTART = auto()  #: Write restart script and start new slurm job with it
 
 
+@unique
+class LogLevel(IntEnum):
+    OFF = 0
+    ERROR = logging.ERROR
+    WARNING = logging.WARNING
+    INFO = logging.INFO
+    DEBUG = logging.DEBUG
+
+
 @dataclass
 @register_pytree_node_class
 class VariPEPS_Config:
@@ -310,6 +320,16 @@ class VariPEPS_Config:
     # Slurm
     slurm_restart_mode: Slurm_Restart_Mode = Slurm_Restart_Mode.WRITE_NEED_RESTART_FILE
 
+    # Logging configuration
+    log_level_global: LogLevel = LogLevel.INFO
+    log_level_optimizer: LogLevel = LogLevel.INFO
+    log_level_ctmrg: LogLevel = LogLevel.INFO
+    log_level_line_search: LogLevel = LogLevel.INFO
+    log_to_console: bool = True
+    log_to_file: bool = False
+    log_file: str = "varipeps.log"
+    log_step_summary_every_n: int = 1
+
     def update(self, name: str, value: Any) -> NoReturn:
         self.__setattr__(name, value)
 
@@ -395,6 +415,7 @@ class ConfigModuleWrapper:
         "Projector_Method",
         "Wavevector_Type",
         "Slurm_Restart_Mode",
+        "LogLevel",
         "VariPEPS_Config",
         "config",
     }

varipeps/ctmrg/routine.py

@@ -5,6 +5,10 @@
 from jax import jit, custom_vjp, vjp, tree_util
 from jax.lax import cond, while_loop
 import jax.debug as jdebug
+import logging
+import time
+
+logger = logging.getLogger("varipeps.ctmrg")
 
 from varipeps import varipeps_config, varipeps_global_state
 from varipeps.peps import PEPS_Tensor, PEPS_Tensor_Split_Transfer, PEPS_Unit_Cell
@@ -515,9 +519,8 @@ def corner_svd_func(old, new, old_corner, conv_eps, config):
         eps,
         config,
     )
-
-    if config.ctmrg_print_steps:
-        debug_print("CTMRG: {}: {}", count, measure)
+    if logger.isEnabledFor(logging.DEBUG):
+        jax.debug.callback(lambda cnt, msr: logger.debug(f"CTMRG: Step {cnt}: {msr}"), count, measure, ordered=True)
         if config.ctmrg_verbose_output:
             jax.debug.callback(print_verbose, verbose_data, ordered=True)
 
@@ -620,7 +623,7 @@ def calc_ctmrg_env(
     best_norm_smallest_S = None
     best_truncation_eps = None
     have_been_increased = False
-
+    t0 = time.perf_counter()
     while True:
         tmp_count = 0
         corner_singular_vals = None
@@ -720,6 +723,11 @@ def calc_ctmrg_env(
         else:
             converged = False
             end_count = tmp_count
+        if logger.isEnabledFor(logging.INFO):
+            if logger.isEnabledFor(logging.WARNING) and not converged:
+                logger.warning("CTMRG: ❌ did not converge, took %.2f seconds. (Steps: %d, Smallest SVD Norm: %.3e)", time.perf_counter() - t0, end_count, norm_smallest_S)
+            else:
+                logger.info("CTMRG: ✅ converged, took %.2f seconds. (Steps: %d, Smallest SVD Norm: %.3e)", time.perf_counter() - t0, end_count, norm_smallest_S)
 
         if converged and (
             working_unitcell[0, 0][0][0].chi > best_chi or best_result is None
@@ -751,9 +759,9 @@ def calc_ctmrg_env(
                 working_unitcell = working_unitcell.change_chi(new_chi)
                 initial_unitcell = initial_unitcell.change_chi(new_chi)
 
-                if varipeps_config.ctmrg_print_steps:
-                    debug_print(
-                        "CTMRG: Increasing chi to {} since smallest SVD Norm was {}.",
+                if logger.isEnabledFor(logging.INFO):
+                    logger.info(
+                        "Increasing chi to {} since smallest SVD Norm was {}.",
                         new_chi,
                         norm_smallest_S,
                     )
@@ -785,9 +793,9 @@ def calc_ctmrg_env(
             if not new_chi in already_tried_chi:
                 working_unitcell = working_unitcell.change_chi(new_chi)
 
-                if varipeps_config.ctmrg_print_steps:
-                    debug_print(
-                        "CTMRG: Decreasing chi to {} since smallest SVD Norm was {} or routine did not converge.",
+                if logger.isEnabledFor(logging.INFO):
+                    logger.info(
+                        "Decreasing chi to {} since smallest SVD Norm was {} or routine did not converge.",
                         new_chi,
                         norm_smallest_S,
                     )
@@ -809,9 +817,9 @@ def calc_ctmrg_env(
                 new_truncation_eps
                 <= varipeps_config.ctmrg_increase_truncation_eps_max_value
             ):
-                if varipeps_config.ctmrg_print_steps:
-                    debug_print(
-                        "CTMRG: Increasing SVD truncation eps to {}.",
+                if logger.isEnabledFor(logging.INFO):
+                    logger.info(
+                        "Increasing SVD truncation eps to {}.",
                         new_truncation_eps,
                     )
                 varipeps_global_state.ctmrg_effective_truncation_eps = (
@@ -937,8 +945,8 @@ def _ctmrg_rev_while_body(carry):
 
     count += 1
 
-    if config.ad_custom_print_steps:
-        debug_print("Custom VJP: {}: {}", count, measure)
+    if logger.isEnabledFor(logging.DEBUG):
+        jax.debug.callback(lambda cnt, msr: logger.debug(f"Custom VJP: Step {cnt}, Measure {msr}"), count, measure, ordered=True)
         if config.ad_custom_verbose_output:
             jax.debug.callback(print_verbose, verbose_data, ordered=True, ad=True)
 
@@ -1014,12 +1022,14 @@ def calc_ctmrg_env_rev(
 
     varipeps_global_state.ctmrg_effective_truncation_eps = last_truncation_eps
 
+    if logger.isEnabledFor(logging.INFO):
+        t0 = time.perf_counter()
     t_bar, converged, end_count = _ctmrg_rev_workhorse(
         peps_tensors, new_unitcell, unitcell_bar, varipeps_config, varipeps_global_state
     )
 
     varipeps_global_state.ctmrg_effective_truncation_eps = None
-
+    debug_print("Custom VJP: Converged: {}, Steps: {}", converged, end_count)
     if end_count == varipeps_config.ad_custom_max_steps and not converged:
         raise CTMRGGradientNotConvergedError
 

varipeps/optimization/line_search.py

@@ -1,6 +1,5 @@
 import enum
 
-from tqdm_loggable.auto import tqdm
 
 import jax
 import jax.numpy as jnp
@@ -14,6 +13,9 @@
 from varipeps.expectation import Expectation_Model
 from varipeps.mapping import Map_To_PEPS_Model
 from varipeps.utils.debug_print import debug_print
+import logging
+
+logger = logging.getLogger("varipeps.line_search")
 
 from .inner_function import (
     calc_ctmrg_expectation,
@@ -443,6 +445,7 @@ def line_search(
                     additional_input,
                     enforce_elementwise_convergence=enforce_elementwise_convergence,
                 )
+                logger.info("🔎 Line search step %d, E=%.6f, alpha=%.3e", count + 1, new_value, alpha)
 
                 if new_unitcell[0, 0][0][0].chi > unitcell[0, 0][0][0].chi:
                     tmp_value = current_value
@@ -463,10 +466,11 @@ def line_search(
                     else:
                         unitcell = unitcell.change_chi(new_unitcell[0, 0][0][0].chi)
 
-                        debug_print(
-                            "Line search: Recalculate original unitcell with higher chi {}.",
-                            new_unitcell[0, 0][0][0].chi,
-                        )
+                        if logger.isEnabledFor(logging.DEBUG):
+                            logger.debug(
+                                "Line search: Recalculate original unitcell with higher chi %s.",
+                                new_unitcell[0, 0][0][0].chi,
+                            )
 
                         if varipeps_config.ad_use_custom_vjp:
                             (
@@ -534,6 +538,7 @@ def line_search(
                         additional_input,
                         calc_preconverged=True,
                     )
+                tqdm.write(f"Line search step {count+1}, E={new_value:<.6f}, alpha={alpha:.4f}")
                 new_gradient = [elem.conj() for elem in new_gradient_seq]
 
                 if new_unitcell[0, 0][0][0].chi > unitcell[0, 0][0][0].chi:
@@ -554,11 +559,11 @@ def line_search(
                         ) = cache_original_unitcell[new_unitcell[0, 0][0][0].chi]
                     else:
                         unitcell = unitcell.change_chi(new_unitcell[0, 0][0][0].chi)
-
-                        debug_print(
-                            "Line search: Recalculate original unitcell with higher chi {}.",
-                            new_unitcell[0, 0][0][0].chi,
-                        )
+                        if logger.isEnabledFor(logging.DEBUG):
+                            logger.debug(
+                                "Line search: Recalculate original unitcell with higher chi %s.",
+                                new_unitcell[0, 0][0][0].chi,
+                            )
 
                         if varipeps_config.ad_use_custom_vjp:
                             (
@@ -1002,7 +1007,7 @@ def line_search(
                 )
 
                 if alpha <= 0:
-                    tqdm.write("Found negative alpha in secant operation!")
+                    logger.warning("Found negative alpha in secant operation!")
 
                 hz_secant_alpha = alpha
 

varipeps/optimization/optimizer.py

@@ -10,8 +10,6 @@
 
 from scipy.optimize import OptimizeResult
 
-from tqdm_loggable.auto import tqdm
-
 import h5py
 
 import numpy as np
@@ -22,6 +20,10 @@
 from jax.lax import scan
 from jax.flatten_util import ravel_pytree
 
+import logging
+
+logger = logging.getLogger("varipeps.optimizer")
+
 from varipeps import varipeps_config, varipeps_global_state
 from varipeps.config import Optimizing_Methods, Slurm_Restart_Mode
 from varipeps.peps import PEPS_Unit_Cell
@@ -603,7 +605,7 @@ def random_noise(a):
     slurm_restart_written = False
     slurm_new_job_id = None
 
-    with tqdm(desc="Optimizing PEPS state", initial=count) as pbar:
+    if True:
         while count < varipeps_config.optimizer_max_steps:
             runtime_start = time.perf_counter()
 
@@ -697,9 +699,6 @@ def random_noise(a):
                     max_trunc_error_list[random_noise_retries] = []
                     step_runtime[random_noise_retries] = []
 
-                    pbar.reset()
-                    pbar.refresh()
-
                     continue
 
             if working_unitcell[0, 0][0][0].chi != chi_before_ctmrg:
@@ -762,7 +761,7 @@ def random_noise(a):
             signal_reset_descent_dir = False
 
             if _scalar_descent_grad(descent_dir, working_gradient) > 0:
-                tqdm.write("Found bad descent dir. Reset to negative gradient!")
+                logger.warning("Found bad descent dir. Reset to negative gradient!")
                 descent_dir = [-elem for elem in working_gradient]
 
             conv = jnp.linalg.norm(ravel_pytree(working_gradient)[0])
@@ -811,7 +810,7 @@ def random_noise(a):
                             is Projector_Method.HALF
                         )
                     ):
-                        tqdm.write(
+                        logger.warning(
                             "Convergence is not sufficient. Retry with some random noise on best result."
                         )
 
@@ -1017,21 +1016,15 @@ def random_noise(a):
 
             count += 1
 
-            pbar.update()
-            pbar.set_postfix(
-                {
-                    "Energy": f"{working_value:0.10f}",
-                    "Retries": random_noise_retries,
-                    "Convergence": f"{conv:0.8f}",
-                    "Line search step": (
-                        f"{linesearch_step:0.8f}"
-                        if linesearch_step is not None
-                        else "0"
-                    ),
-                    "Max. trunc. err.": f"{max_trunc_error:0.8g}",
-                }
+            logger.info(
+                "📉 Step %d | Energy: %0.10f | Retries: %d | Conv: %0.8f | Line search step: %s | Max. trunc. err.: %0.8g",
+                count,
+                working_value,
+                random_noise_retries,
+                conv,
+                f"{float(linesearch_step):0.8f}" if linesearch_step is not None else "None",
+                max_trunc_error,
             )
-            pbar.refresh()
 
             if count % varipeps_config.optimizer_autosave_step_count == 0:
                 _autosave_wrapper(
@@ -1187,7 +1180,7 @@ def random_noise(a):
                             slurm_data["WorkDir"],
                         )
                         if slurm_new_job_id is None:
-                            tqdm.write(
+                            logger.error(
                                 "Failed to start new Slurm job or parse its job id."
                             )
                         break

varipeps/utils/__init__.py

@@ -4,3 +4,4 @@
 from . import projector_dict
 from . import slurm
 from . import svd
+from . import logging_config

varipeps/utils/debug_print.py

@@ -1,22 +1,23 @@
 import functools
+import logging
 
-from tqdm_loggable.auto import tqdm
 
 import jax.debug as jdebug
 from jax._src.debugging import formatter
 
-# Adapting function from jax.debug to work with tqdm
+logger = logging.getLogger("varipeps.ctmrg")
 
 
 def _format_print_callback(fmt: str, *args, **kwargs):
-    tqdm.write(fmt.format(*args, **kwargs))
+    # Send to logger (respects per-module levels/handlers)
+    logger.debug(fmt.format(*args, **kwargs))
 
 
 def debug_print(fmt: str, *args, ordered: bool = True, **kwargs) -> None:
     """
     Prints values and works in staged out JAX functions.
 
-    Function adapted from :obj:`jax.debug.print` to work with tqdm. See there
+    Function adapted from :obj:`jax.debug.print` to work with logger. See there
     for original authors and function.
 
     Args: