Use fastprogress instead of tqdm progressbar

pymc3/parallel_sampling.py

@@ -17,28 +17,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 +240,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 +348,7 @@ def __init__(
         start_chain_num=0,
         progressbar=True,
     ):
-        if progressbar:
-            from tqdm import tqdm
+        from fastprogress import progress_bar
 
         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()}
 
 

pymc3/smc/smc.py

@@ -1,7 +1,7 @@
 from collections import OrderedDict
 
 import numpy as np
-from tqdm import tqdm
+from fastprogress import progress_bar
 import multiprocessing as mp
 import warnings
 from theano import function as theano_function
@@ -91,7 +91,9 @@ def initialize_population(self):
         var_info = OrderedDict()
         if self.start is None:
             init_rnd = sample_prior_predictive(
-                self.draws, var_names=[v.name for v in self.model.unobserved_RVs], model=self.model
+                self.draws,
+                var_names=[v.name for v in self.model.unobserved_RVs],
+                model=self.model,
             )
         else:
             init_rnd = self.start
@@ -103,7 +105,9 @@ def initialize_population(self):
 
         for i in range(self.draws):
 
-            point = Point({v.name: init_rnd[v.name][i] for v in self.variables}, model=self.model)
+            point = Point(
+                {v.name: init_rnd[v.name][i] for v in self.variables}, model=self.model
+            )
             population.append(self.model.dict_to_array(point))
 
         self.posterior = np.array(floatX(population))
@@ -119,7 +123,9 @@ def setup_kernel(self):
         if self.kernel.lower() == "abc":
             warnings.warn(EXPERIMENTAL_WARNING)
             if len(self.model.observed_RVs) != 1:
-                warnings.warn("SMC-ABC only works properly with models with one observed variable")
+                warnings.warn(
+                    "SMC-ABC only works properly with models with one observed variable"
+                )
             simulator = self.model.observed_RVs[0]
             self.likelihood_logp = PseudoLikelihood(
                 self.epsilon,
@@ -131,15 +137,19 @@ def setup_kernel(self):
                 self.sum_stat,
             )
         elif self.kernel.lower() == "metropolis":
-            self.likelihood_logp = logp_forw([self.model.datalogpt], self.variables, shared)
+            self.likelihood_logp = logp_forw(
+                [self.model.datalogpt], self.variables, shared
+            )
 
     def initialize_logp(self):
         """
         initialize the prior and likelihood log probabilities
         """
         if self.parallel and self.cores > 1:
             self.pool = mp.Pool(processes=self.cores)
-            priors = self.pool.starmap(self.prior_logp, [(sample,) for sample in self.posterior])
+            priors = self.pool.starmap(
+                self.prior_logp, [(sample,) for sample in self.posterior]
+            )
             likelihoods = self.pool.starmap(
                 self.likelihood_logp, [(sample,) for sample in self.posterior]
             )
@@ -204,22 +214,27 @@ def update_proposal(self):
         cov = np.atleast_2d(cov)
         cov += 1e-6 * np.eye(cov.shape[0])
         if np.isnan(cov).any() or np.isinf(cov).any():
-            raise ValueError('Sample covariances not valid! Likely "draws" is too small!')
+            raise ValueError(
+                'Sample covariances not valid! Likely "draws" is too small!'
+            )
         self.proposal = MultivariateNormalProposal(cov)
 
     def tune(self):
         """
         Tune scaling and n_steps based on the acceptance rate.
         """
-        ave_scaling = np.exp(np.log(self.scalings.mean()) + (self.acc_per_chain.mean() - 0.234))
+        ave_scaling = np.exp(
+            np.log(self.scalings.mean()) + (self.acc_per_chain.mean() - 0.234)
+        )
         self.scalings = 0.5 * (
             ave_scaling + np.exp(np.log(self.scalings) + (self.acc_per_chain - 0.234))
         )
 
         if self.tune_steps:
             acc_rate = max(1.0 / self.proposed, self.acc_rate)
             self.n_steps = min(
-                self.max_steps, max(2, int(np.log(1 - self.p_acc_rate) / np.log(1 - acc_rate)))
+                self.max_steps,
+                max(2, int(np.log(1 - self.p_acc_rate) / np.log(1 - acc_rate))),
             )
 
         self.proposed = self.draws * self.n_steps
@@ -264,7 +279,7 @@ def mutate(self):
                     draw,
                     *parameters
                 )
-                for draw in tqdm(range(self.draws), disable=not self.progressbar)
+                for draw in progress_bar(range(self.draws), display=self.progressbar)
             ]
         posterior, acc_list, priors, likelihoods = zip(*results)
         self.posterior = np.array(posterior)
@@ -335,7 +350,9 @@ def metrop_kernel(
 
         new_tempered_logp = pl + ll * beta
 
-        q_old, accept = metrop_select(new_tempered_logp - old_tempered_logp, q_new, q_old)
+        q_old, accept = metrop_select(
+            new_tempered_logp - old_tempered_logp, q_new, q_old
+        )
 
         if accept:
             accepted += 1
@@ -369,7 +386,9 @@ class PseudoLikelihood:
     Pseudo Likelihood
     """
 
-    def __init__(self, epsilon, observations, function, model, var_info, distance, sum_stat):
+    def __init__(
+        self, epsilon, observations, function, model, var_info, distance, sum_stat
+    ):
         """
         epsilon : float
             Standard deviation of the gaussian pseudo likelihood.
@@ -419,7 +438,9 @@ def posterior_to_function(self, posterior):
 
     def gauss_kernel(self, value):
         epsilon = self.epsilon
-        return (-(value ** 2) / epsilon ** 2 + np.log(1 / (2 * np.pi * epsilon ** 2))) / 2.0
+        return (
+            -(value ** 2) / epsilon ** 2 + np.log(1 / (2 * np.pi * epsilon ** 2))
+        ) / 2.0
 
     def absolute_error(self, a, b):
         if self.sum_stat: