Improved logging (issue #152)

dingo/core/dataset.py

@@ -4,6 +4,7 @@
 import numpy as np
 import pandas as pd
 
+from dingo.core.utils.logging_utils import logger
 from dingo.core.utils.misc import get_version
 
 
@@ -136,7 +137,7 @@ def __init__(
             self.from_dictionary(dictionary)
 
     def to_file(self, file_name: str, mode: str = "w"):
-        print("Saving dataset to " + str(file_name))
+        logger.info("Saving dataset to " + str(file_name))
         save_dict = {
             k: v
             for k, v in vars(self).items()
@@ -150,9 +151,9 @@ def to_file(self, file_name: str, mode: str = "w"):
                 f.attrs["dataset_type"] = self.dataset_type
 
     def from_file(self, file_name: str):
-        print(f"Loading dataset from {str(file_name)}.")
+        logger.info(f"Loading dataset from {str(file_name)}.")
         if self._leave_on_disk_keys:
-            print(f"Omitting data keys {self._leave_on_disk_keys}.")
+            logger.info(f"Omitting data keys {self._leave_on_disk_keys}.")
 
         with h5py.File(file_name, "r") as f:
             loaded_dict = recursive_hdf5_load(

dingo/core/posterior_models/base_model.py

@@ -10,6 +10,7 @@
 
 import torch
 import dingo.core.utils as utils
+from dingo.core.utils.logging_utils import logger
 from torch.utils.data import Dataset
 import time
 import numpy as np
@@ -198,7 +199,7 @@ def network_to_device(self, device):
         #     raise NotImplementedError('This needs testing!')
         #     # dim = 0 [512, ...] -> [256, ...], [256, ...] on 2 GPUs
         #     self.network = torch.nn.DataParallel(self.network)
-        print(f"Putting posterior model to device {self.device}.")
+        logger.info(f"Putting posterior model to device {self.device}.")
         self.network.to(self.device)
 
     def initialize_optimizer_and_scheduler(self):
@@ -394,7 +395,7 @@ def train(
 
         if test_only:
             test_loss = test_epoch(self, test_loader)
-            print(f"test loss: {test_loss:.3f}")
+            logger.info(f"test loss: {test_loss:.3f}")
 
         else:
             while not runtime_limits.limits_exceeded(self.epoch):
@@ -403,24 +404,24 @@ def train(
                 # Training
                 lr = utils.get_lr(self.optimizer)
                 with threadpool_limits(limits=1, user_api="blas"):
-                    print(f"\nStart training epoch {self.epoch} with lr {lr}")
+                    logger.info(f"Start training epoch {self.epoch} with lr {lr}")
                     time_start = time.time()
                     train_loss = train_epoch(self, train_loader)
                     train_time = time.time() - time_start
 
-                    print(
+                    logger.info(
                         "Done. This took {:2.0f}:{:2.0f} min.".format(
                             *divmod(train_time, 60)
                         )
                     )
 
                     # Testing
-                    print(f"Start testing epoch {self.epoch}")
+                    logger.info(f"Start testing epoch {self.epoch}")
                     time_start = time.time()
                     test_loss = test_epoch(self, test_loader)
                     test_time = time.time() - time_start
 
-                    print(
+                    logger.info(
                         "Done. This took {:2.0f}:{:2.0f} min.".format(
                             *divmod(time.time() - time_start, 60)
                         )
@@ -449,7 +450,7 @@ def train(
                             }
                         )
                     except ImportError:
-                        print("wandb not installed. Skipping logging to wandb.")
+                        logger.warning("wandb not installed. Skipping logging to wandb.")
 
                 if early_stopping is not None:
                     # Whether to use train or test loss
@@ -464,9 +465,9 @@ def train(
                             join(train_dir, "best_model.pt"), save_training_info=False
                         )
                     if early_stopping.early_stop:
-                        print("Early stopping")
+                        logger.info("Early stopping")
                         break
-                print(f"Finished training epoch {self.epoch}.\n")
+                logger.info(f"Finished training epoch {self.epoch}.")
 
 
 def train_epoch(pm, dataloader):

dingo/core/result.py

@@ -16,6 +16,7 @@
 from bilby.core.prior import Constraint, DeltaFunction, PriorDict
 
 from dingo.core.dataset import DingoDataset
+from dingo.core.utils.logging_utils import logger
 from dingo.core.density import train_unconditional_density_estimator
 from dingo.core.utils.misc import recursive_check_dicts_are_equal
 from dingo.core.utils.plotting import get_latex_labels, plot_corner_multi
@@ -167,23 +168,24 @@ def reset_event(self, event_dataset):
         ):
             # This is really just for notification. Actions are only taken if the
             # event metadata differ.
-            print("\nNew event data differ from existing.")
+            logger.warning("New event data differ from existing.")
         self.context = context
 
         if self.event_metadata is not None and self.event_metadata != event_metadata:
-            print("Changes")
-            print("=======")
+            logger.warning("Changes")
+            logger.warning("=======")
             old_minus_new = dict(freeze(self.event_metadata) - freeze(event_metadata))
-            print("Old event metadata:")
+            logger.warning("Old event metadata:")
             for k in sorted(old_minus_new):
-                print(f"  {k}:  {self.event_metadata[k]}")
+                logger.warning(f"  {k}:  {self.event_metadata[k]}")
 
             new_minus_old = dict(freeze(event_metadata) - freeze(self.event_metadata))
-            print("New event metadata:")
+            logger.warning("New event metadata:")
             if self.importance_sampling_metadata.get("updates") is None:
                 self.importance_sampling_metadata["updates"] = {}
             for k in sorted(new_minus_old):
-                print(f"  {k}:  {event_metadata[k]}")
+                logger.warning(f"  {k}:  {event_metadata[k]}")
+                self.importance_sampling_metadata["updates"][k] = event_metadata[k]
                 self.importance_sampling_metadata["updates"][k] = event_metadata[k]
 
             self._rebuild_domain(verbose=True)
@@ -312,12 +314,12 @@ def importance_sample(self, num_processes: int = 1, **likelihood_kwargs):
         valid_samples = np.isfinite(log_prior + delta_log_prob_target)
         theta = theta.iloc[valid_samples]
 
-        print(f"Calculating {len(theta)} likelihoods.")
+        logger.info(f"Calculating {len(theta)} likelihoods.")
         t0 = time.time()
         log_likelihood = self.likelihood.log_likelihood_multi(
             theta, num_processes=num_processes
         )
-        print(f"Done. This took {time.time() - t0:.2f} seconds.")
+        logger.info(f"Done. This took {time.time() - t0:.2f} seconds.")
 
         self.log_noise_evidence = self.likelihood.log_Zn
         self.samples["log_prior"] = log_prior
@@ -494,7 +496,7 @@ def rejection_sample(
         rng = np.random.default_rng(random_state)
         weights = self.samples["weights"].to_numpy(dtype=float)
 
-        print(
+        logger.info(
             f"Rejection sampling: {self.num_samples} samples, "
             f"ESS = {self.effective_sample_size:.0f} "
             f"(efficiency = {100 * self.sample_efficiency:.1f}%)"
@@ -526,7 +528,7 @@ def rejection_sample(
         unweighted = unweighted.drop(
             columns=["weights", "log_prob", "delta_log_prob_target"], errors="ignore"
         )
-        print(f"Produced {len(unweighted)} unweighted samples.")
+        logger.info(f"Produced {len(unweighted)} unweighted samples.")
         return unweighted
 
     def parameter_subset(self, parameters):
@@ -619,12 +621,12 @@ def print_summary(self):
         Display the number of samples, and (if importance sampling is complete) the log
         evidence and number of effective samples.
         """
-        print("Number of samples:", len(self.samples))
+        logger.info("Number of samples: %d", len(self.samples))
         if self.log_evidence is not None:
-            print(
+            logger.info(
                 f"Log(evidence): {self.log_evidence:.3f} +- {self.log_evidence_std:.3f}"
             )
-            print(
+            logger.info(
                 f"Effective samples {self.n_eff:.1f}: "
                 f"(Sample efficiency = {100 * self.sample_efficiency:.2f}%)"
             )
@@ -815,7 +817,7 @@ def plot_log_probs(self, filename="log_probs.png"):
             plt.tight_layout()
             plt.savefig(filename)
         else:
-            print("Results not importance sampled. Cannot produce log_prob plot.")
+            logger.warning("Results not importance sampled. Cannot produce log_prob plot.")
 
     def plot_weights(self, filename="weights.png"):
         """Make a scatter plot of samples weights vs log proposal."""
@@ -844,7 +846,7 @@ def plot_weights(self, filename="weights.png"):
             plt.tight_layout()
             plt.savefig(filename)
         else:
-            print("Results not importance sampled. Cannot plot weights.")
+            logger.warning("Results not importance sampled. Cannot plot weights.")
 
     def get_all_injection_credible_levels(
         self, keys: list[str] = None, weighted: bool = False
@@ -1020,7 +1022,7 @@ def make_pp_plot(
 
     pvalues = []
     latex_labels = get_latex_labels(results[0].prior)
-    print("Key: KS-test p-value")
+    logger.info("Key: KS-test p-value")
     for ii, key in enumerate(credible_levels):
         pp = np.array(
             [
@@ -1030,7 +1032,7 @@ def make_pp_plot(
         )
         pvalue = scipy.stats.kstest(credible_levels[key], "uniform").pvalue
         pvalues.append(pvalue)
-        print("{}: {}".format(key, pvalue))
+        logger.info("{}: {}".format(key, pvalue))
         label = "{} ({:2.3f})".format(latex_labels[key], pvalue)
         plt.plot(x_values, pp, lines[ii], label=label, **kwargs)
 
@@ -1040,7 +1042,7 @@ def make_pp_plot(
         pvalues=pvalues,
         names=list(credible_levels.keys()),
     )
-    print("Combined p-value: {}".format(pvals.combined_pvalue))
+    logger.info("Combined p-value: {}".format(pvals.combined_pvalue))
 
     if title:
         ax.set_title(

dingo/core/samplers.py

@@ -11,6 +11,7 @@
 from torchvision.transforms import Compose
 
 from dingo.core.posterior_models import BasePosteriorModel
+from dingo.core.utils.logging_utils import logger
 from dingo.core.result import Result
 from dingo.core.result import DATA_KEYS as RESULT_DATA_KEYS
 from dingo.core.utils import torch_detach_to_cpu, IterationTracker
@@ -157,7 +158,7 @@ def _run_sampler(
             x = [x]
         else:
             if context is not None:
-                print("Unconditional model. Ignoring context.")
+                logger.warning("Unconditional model. Ignoring context.")
             x = []
 
         # For a normalizing flow, we get the log_prob for "free" when sampling,
@@ -205,7 +206,7 @@ def run_sampler(
         """
         self.samples = None
 
-        print(f"Running sampler to generate {num_samples} samples.")
+        logger.info(f"Running sampler to generate {num_samples} samples.")
         t0 = time.time()
         if not self.unconditional_model:
             if self.context is None:
@@ -229,7 +230,7 @@ def run_sampler(
         # correction for t_ref) and represent as DataFrame.
         self._post_process(samples)
         self.samples = pd.DataFrame(samples)
-        print(f"Done. This took {time.time() - t0:.1f} s.")
+        logger.info(f"Done. This took {time.time() - t0:.1f} s.")
         sys.stdout.flush()
 
     def log_prob(self, samples: pd.DataFrame | dict) -> np.ndarray:
@@ -430,8 +431,8 @@ def _run_sampler(
             init_samples = self.init_sampler._run_sampler(num_samples, context)
         else:
             if self.num_iterations == 1:
-                print(
-                    f"Warning: Removing initial outliers, but only carrying out "
+                logger.warning(
+                    f"Removing initial outliers, but only carrying out "
                     f"{self.num_iterations} GNPE iteration. This risks biasing "
                     f"results."
                 )
@@ -515,16 +516,15 @@ def _run_sampler(
                 p: x["extrinsic_parameters"][p] for p in self.gnpe_proxy_parameters
             }
 
-            print(
-                f"it {i}.\tmin pvalue: {self.iteration_tracker.pvalue_min:.3f}"
-                f"\tproxy mean: ",
-                *[f"{torch.mean(v).item():.5f}" for v in proxies.values()],
-                "\tproxy std:",
-                *[f"{torch.std(v).item():.5f}" for v in proxies.values()],
-                "\ttimes:",
-                time_sample_start - start_time,
-                time_sample_end - time_sample_start,
-                time.time() - time_sample_end,
+            proxy_means = " ".join(f"{torch.mean(v).item():.5f}" for v in proxies.values())
+            proxy_stds = " ".join(f"{torch.std(v).item():.5f}" for v in proxies.values())
+            logger.debug(
+                f"it {i}. min pvalue: {self.iteration_tracker.pvalue_min:.3f}"
+                f" proxy mean: {proxy_means}"
+                f" proxy std: {proxy_stds}"
+                f" times: {time_sample_start - start_time:.3f}"
+                f" {time_sample_end - time_sample_start:.3f}"
+                f" {time.time() - time_sample_end:.3f}"
             )
 
         #

dingo/core/utils/condor_utils.py

@@ -2,6 +2,8 @@
 from os.path import join
 import yaml
 
+from dingo.core.utils.logging_utils import logger
+
 
 def resubmit_condor_job(train_dir, train_settings, epoch):
     """
@@ -12,14 +14,14 @@ def resubmit_condor_job(train_dir, train_settings, epoch):
     :return:
     """
     if 'condor_settings' in train_settings:
-        print('Copying log files')
+        logger.info("Copying log files")
         copy_logfiles(train_dir, epoch=epoch)
 
         if epoch >= train_settings['train_settings']['runtime_limits'][
             'max_epochs_total']:
-            print('Training complete, job will not be resubmitted')
+            logger.info("Training complete, job will not be resubmitted")
         else:
-            print('Training incomplete, resubmitting job.')
+            logger.info("Training incomplete, resubmitting job.")
             create_submission_file_and_submit_job(train_dir)
 
 
@@ -76,7 +78,7 @@ def copy_logfiles(log_dir, epoch, name='info', suffixes=('.err','.log','.out')):
         try:
             copyfile(src, dest)
         except:
-            print('Could not copy ' + src)
+            logger.warning("Could not copy " + src)
 
 
 if __name__ == '__main__':