Revert to using bilby print func, but implement tqdm

bilby/core/sampler/dynesty.py

@@ -6,6 +6,7 @@ import sys
 import pickle
 import signal
 
+import tqdm
 import matplotlib.pyplot as plt
 import numpy as np
 from pandas import DataFrame
@@ -159,11 +160,12 @@ class Dynesty(NestedSampler):
             self.kwargs['walks'] = self.ndim * 10
         if not self.kwargs['update_interval']:
             self.kwargs['update_interval'] = int(0.6 * self.kwargs['nlive'])
-        if self.kwargs['print_func'] == "bilby-classic":
+        if self.kwargs['print_func'] is None:
             self.kwargs['print_func'] = self._print_func
+            self.pbar = tqdm.tqdm(file=sys.stdout)
         Sampler._verify_kwargs_against_default_kwargs(self)
 
-    def _print_func(self, pbar, results, niter, ncall=None, dlogz=None, *args, **kwargs):
+    def _print_func(self, results, niter, ncall=None, dlogz=None, *args, **kwargs):
         """ Replacing status update for dynesty.result.print_func """
 
         # Extract results at the current iteration.
@@ -189,13 +191,15 @@ class Dynesty(NestedSampler):
             key = 'logz'
 
         # Constructing output.
-        raw_string = "\r {}| {}={:6.3f} +/- {:6.3f} | dlogz: {:6.3f} > {:6.3f}"
-        print_str = raw_string.format(
-            niter, key, logz, logzerr, delta_logz, dlogz)
+        string = []
+        string.append("bound:{:d}".format(bounditer))
+        string.append("ncall: {:d}".format(ncall))
+        string.append("eff: {:0.1f}".format(eff))
+        string.append("{}={:0.2f}+/-{:0.2f}".format(key, logz, logzerr))
+        string.append("dlogz:{:0.3f} > {:0.3f}".format(delta_logz, dlogz))
 
-        # Printing.
-        sys.stdout.write(print_str)
-        sys.stdout.flush()
+        self.pbar.set_postfix_str(" | ".join(string), refresh=True)
+        self.pbar.update(niter - self.pbar.n)
 
     def _apply_dynesty_boundaries(self):
         self._periodic = list()
@@ -232,6 +236,7 @@ class Dynesty(NestedSampler):
 
         # Flushes the output to force a line break
         if self.kwargs["verbose"]:
+            self.pbar.close()
             print("")
 
         check_directory_exists_and_if_not_mkdir(self.outdir)