Use fastprogress instead of tqdm progressbar

RELEASE-NOTES.md

@@ -1,6 +1,11 @@
 # Release Notes
 
-## PyMC3 3.8 (on deck)
+## PyMC3 3.9 (On deck)
+
+### New features
+- use [fastprogress](https://github.com/fastai/fastprogress) instead of tqdm [#3693](https://github.com/pymc-devs/pymc3/pull/3693)
+
+## PyMC3 3.8 (November 29 2019)
 
 ### New features
 - Implemented robust u turn check in NUTS (similar to stan-dev/stan#2800). See PR [#3605]

pymc3/parallel_sampling.py

@@ -9,6 +9,7 @@
 import errno
 
 import numpy as np
+from fastprogress import progress_bar
 
 from . import theanof
 
@@ -17,28 +18,31 @@
 
 def _get_broken_pipe_exception():
     import sys
-    if sys.platform == 'win32':
-        return RuntimeError("The communication pipe between the main process "
-                            "and its spawned children is broken.\n"
-                            "In Windows OS, this usually means that the child "
-                            "process raised an exception while it was being "
-                            "spawned, before it was setup to communicate to "
-                            "the main process.\n"
-                            "The exceptions raised by the child process while "
-                            "spawning cannot be caught or handled from the "
-                            "main process, and when running from an IPython or "
-                            "jupyter notebook interactive kernel, the child's "
-                            "exception and traceback appears to be lost.\n"
-                            "A known way to see the child's error, and try to "
-                            "fix or handle it, is to run the problematic code "
-                            "as a batch script from a system's Command Prompt. "
-                            "The child's exception will be printed to the "
-                            "Command Promt's stderr, and it should be visible "
-                            "above this error and traceback.\n"
-                            "Note that if running a jupyter notebook that was "
-                            "invoked from a Command Prompt, the child's "
-                            "exception should have been printed to the Command "
-                            "Prompt on which the notebook is running.")
+
+    if sys.platform == "win32":
+        return RuntimeError(
+            "The communication pipe between the main process "
+            "and its spawned children is broken.\n"
+            "In Windows OS, this usually means that the child "
+            "process raised an exception while it was being "
+            "spawned, before it was setup to communicate to "
+            "the main process.\n"
+            "The exceptions raised by the child process while "
+            "spawning cannot be caught or handled from the "
+            "main process, and when running from an IPython or "
+            "jupyter notebook interactive kernel, the child's "
+            "exception and traceback appears to be lost.\n"
+            "A known way to see the child's error, and try to "
+            "fix or handle it, is to run the problematic code "
+            "as a batch script from a system's Command Prompt. "
+            "The child's exception will be printed to the "
+            "Command Promt's stderr, and it should be visible "
+            "above this error and traceback.\n"
+            "Note that if running a jupyter notebook that was "
+            "invoked from a Command Prompt, the child's "
+            "exception should have been printed to the Command "
+            "Prompt on which the notebook is running."
+        )
     else:
         return None
 
@@ -237,7 +241,6 @@ def __init__(self, draws, tune, step_method, chain, seed, start):
             tune,
             seed,
         )
-        # We fork right away, so that the main process can start tqdm threads
         try:
             self._process.start()
         except IOError as e:
@@ -346,8 +349,6 @@ def __init__(
         start_chain_num=0,
         progressbar=True,
     ):
-        if progressbar:
-            from tqdm import tqdm
 
         if any(len(arg) != chains for arg in [seeds, start_points]):
             raise ValueError("Number of seeds and start_points must be %s." % chains)
@@ -369,14 +370,13 @@ def __init__(
 
         self._progress = None
         self._divergences = 0
+        self._total_draws = 0
         self._desc = "Sampling {0._chains:d} chains, {0._divergences:,d} divergences"
         self._chains = chains
-        if progressbar:
-            self._progress = tqdm(
-                total=chains * (draws + tune),
-                unit="draws",
-                desc=self._desc.format(self)
-            )
+        self._progress = progress_bar(
+            range(chains * (draws + tune)), display=progressbar, auto_update=False
+        )
+        self._progress.comment = self._desc.format(self)
 
     def _make_active(self):
         while self._inactive and len(self._active) < self._max_active:
@@ -393,11 +393,11 @@ def __iter__(self):
         while self._active:
             draw = ProcessAdapter.recv_draw(self._active)
             proc, is_last, draw, tuning, stats, warns = draw
-            if self._progress is not None:
-                if not tuning and stats and stats[0].get('diverging'):
-                    self._divergences += 1
-                    self._progress.set_description(self._desc.format(self))
-                self._progress.update()
+            self._total_draws += 1
+            if not tuning and stats and stats[0].get("diverging"):
+                self._divergences += 1
+                self._progress.comment = self._desc.format(self)
+            self._progress.update(self._total_draws)
 
             if is_last:
                 proc.join()
@@ -423,8 +423,7 @@ def __enter__(self):
 
     def __exit__(self, *args):
         ProcessAdapter.terminate_all(self._samplers)
-        if self._progress is not None:
-            self._progress.close()
+
 
 def _cpu_count():
     """Try to guess the number of CPUs in the system.

pymc3/sampling.py

@@ -40,7 +40,7 @@
 from .parallel_sampling import _cpu_count
 from pymc3.step_methods.hmc import quadpotential
 import pymc3 as pm
-from tqdm import tqdm
+from fastprogress import progress_bar
 
 
 import sys
@@ -568,11 +568,17 @@ def _sample_population(
     # create the generator that iterates all chains in parallel
     chains = [chain + c for c in range(chains)]
     sampling = _prepare_iter_population(
-        draws, chains, step, start, parallelize, tune=tune, model=model, random_seed=random_seed
+        draws,
+        chains,
+        step,
+        start,
+        parallelize,
+        tune=tune,
+        model=model,
+        random_seed=random_seed,
     )
 
-    if progressbar:
-        sampling = tqdm(sampling, total=draws)
+    sampling = progress_bar(sampling, total=draws, display=progressbar)
 
     latest_traces = None
     for it, traces in enumerate(sampling):
@@ -596,23 +602,20 @@ def _sample(
 
     sampling = _iter_sample(draws, step, start, trace, chain, tune, model, random_seed)
     _pbar_data = None
-    if progressbar:
-        _pbar_data = {"chain": chain, "divergences": 0}
-        _desc = "Sampling chain {chain:d}, {divergences:,d} divergences"
-        sampling = tqdm(sampling, total=draws, desc=_desc.format(**_pbar_data))
+    _pbar_data = {"chain": chain, "divergences": 0}
+    _desc = "Sampling chain {chain:d}, {divergences:,d} divergences"
+    sampling = progress_bar(sampling, total=draws, display=progressbar)
+    sampling.comment = _desc.format(**_pbar_data)
     try:
         strace = None
         for it, (strace, diverging) in enumerate(sampling):
             if it >= skip_first:
                 trace = MultiTrace([strace])
                 if diverging and _pbar_data is not None:
                     _pbar_data["divergences"] += 1
-                    sampling.set_description(_desc.format(**_pbar_data))
+                    sampling.comment = _desc.format(**_pbar_data)
     except KeyboardInterrupt:
         pass
-    finally:
-        if progressbar:
-            sampling.close()
     return strace
 
 
@@ -753,7 +756,7 @@ def __init__(self, steppers, parallelize):
                 )
                 import multiprocessing
 
-                for c, stepper in enumerate(tqdm(steppers)):
+                for c, stepper in enumerate(progress_bar(steppers)):
                     slave_end, master_end = multiprocessing.Pipe()
                     stepper_dumps = pickle.dumps(stepper, protocol=4)
                     process = multiprocessing.Process(
@@ -1235,9 +1238,13 @@ def sample_posterior_predictive(
         nchain = 1
 
     if keep_size and samples is not None:
-        raise IncorrectArgumentsError("Should not specify both keep_size and samples argukments")
+        raise IncorrectArgumentsError(
+            "Should not specify both keep_size and samples argukments"
+        )
     if keep_size and size is not None:
-        raise IncorrectArgumentsError("Should not specify both keep_size and size argukments")
+        raise IncorrectArgumentsError(
+            "Should not specify both keep_size and size argukments"
+        )
 
     if samples is None:
         samples = sum(len(v) for v in trace._straces.values())
@@ -1253,7 +1260,9 @@ def sample_posterior_predictive(
 
     if var_names is not None:
         if vars is not None:
-            raise IncorrectArgumentsError("Should not specify both vars and var_names arguments.")
+            raise IncorrectArgumentsError(
+                "Should not specify both vars and var_names arguments."
+            )
         else:
             vars = [model[x] for x in var_names]
     elif vars is not None:  # var_names is None, and vars is not.
@@ -1266,8 +1275,7 @@ def sample_posterior_predictive(
 
     indices = np.arange(samples)
 
-    if progressbar:
-        indices = tqdm(indices, total=samples)
+    indices = progress_bar(indices, total=samples, display=progressbar)
 
     ppc_trace_t = _DefaultTrace(samples)
     try:
@@ -1285,10 +1293,6 @@ def sample_posterior_predictive(
     except KeyboardInterrupt:
         pass
 
-    finally:
-        if progressbar:
-            indices.close()
-
     ppc_trace = ppc_trace_t.trace_dict
     if keep_size:
         for k, ary in ppc_trace.items():
@@ -1411,8 +1415,7 @@ def sample_posterior_predictive_w(
 
     indices = np.random.randint(0, len_trace, samples)
 
-    if progressbar:
-        indices = tqdm(indices, total=samples)
+    indices = progress_bar(indices, total=samples, display=progressbar)
 
     try:
         ppc = defaultdict(list)
@@ -1426,10 +1429,6 @@ def sample_posterior_predictive_w(
     except KeyboardInterrupt:
         pass
 
-    finally:
-        if progressbar:
-            indices.close()
-
     return {k: np.asarray(v) for k, v in ppc.items()}