fixes for logger + added some other logging

Author: MPMPMPMPMPMPMP

varipeps/ctmrg/routine.py

@@ -898,6 +898,8 @@ def calc_ctmrg_env_fwd(
     Internal helper function of custom VJP to calculate the values in
     the forward sweep.
     """
+    if logger.isEnabledFor(logging.DEBUG):
+        logger.debug("Custom VJP: Starting forward CTMRG calculation.")
     new_unitcell, last_truncation_eps, norm_smallest_S = calc_ctmrg_env_custom_rule(
         peps_tensors, unitcell, _return_truncation_eps=True
     )
@@ -952,7 +954,7 @@ def _ctmrg_rev_while_body(carry):
     count += 1
 
     if logger.isEnabledFor(logging.DEBUG):
-        jax.debug.callback(lambda cnt, msr: logger.debug(f"Custom VJP: Step {cnt}, Measure {msr}"), count, measure, ordered=True)
+        jax.debug.callback(lambda cnt, msr: logger.debug(f"Custom VJP: Step {cnt}: {msr}"), count, measure, ordered=True)
         if config.ad_custom_verbose_output:
             jax.debug.callback(print_verbose, verbose_data, ordered=True, ad=True)
 
@@ -1023,6 +1025,8 @@ def calc_ctmrg_env_rev(
     Internal helper function of custom VJP to calculate the gradient in
     the backward sweep.
     """
+    if logger.isEnabledFor(logging.DEBUG):
+        logger.debug("Custom VJP: Starting reverse CTMRG calculation.")
     unitcell_bar, _ = input_bar
     peps_tensors, new_unitcell, input_unitcell, last_truncation_eps = res
 

varipeps/optimization/optimizer.py

@@ -204,10 +204,12 @@ def autosave_function(
     auxiliary_data: Optional[Dict[str, Any]] = None,
 ) -> None:
     if counter is not None:
+        logger.info(f"💾 Autosaving to {str(filename)}.{counter}")
         unitcell.save_to_file(
             f"{str(filename)}.{counter}", auxiliary_data=auxiliary_data
         )
     else:
+        logger.info(f"💾 Autosaving to {str(filename)}")
         unitcell.save_to_file(filename, auxiliary_data=auxiliary_data)
 
 

varipeps/utils/logging_config.py

@@ -34,8 +34,13 @@ def init_logging(cfg: Any | None = None) -> None:
     root.setLevel(_to_py_log_level(getattr(cfg, "log_level_global", logging.INFO)))
     root.propagate = False
 
+    # fmt = logging.Formatter(
+    #     fmt="%(asctime)s %(levelname)s %(name)s: %(message)s",
+    #     datefmt="%H:%M:%S",
+    # )
+
     fmt = logging.Formatter(
-        fmt="%(asctime)s %(levelname)s %(name)s: %(message)s",
+        fmt="%(asctime)s %(levelname)s %(message)s",
         datefmt="%H:%M:%S",
     )