Refactor progress bars (#1272)

tqdm progress bars are used in a couple of places. They don't play so well in
non-interactive jobs, testing, ... . They also cause trouble with nbspinx
(#1246, #1271).
Progress bars can be disabled for specific tasks, but not globally (or at
least not very conveniently).
Since recently, tqdm can be controlled via environment variables (e.g.,
disabling all progress bars or changing update frequency). However,
this works by changing the argument defaults, so it only works if we don't
pass explicit `disable=...`. Therefore, this PR introduces some wrapper that
checks whether the user explicitly enabled/disabled progress bars. If not,
we go with the tqdm default, which means showing all progress bars unless
globally disabled. An additional `enabled` argument is added for convenience.

---------

Co-authored-by: Dilan Pathirana <[email protected]>

doc/conf.py

@@ -29,6 +29,11 @@
 # Note: Debugging will proceed. Set PYDEVD_DISABLE_FILE_VALIDATION=1 to disable this validation.
 os.environ["PYDEVD_DISABLE_FILE_VALIDATION"] = "1"
 
+# TQDM and nbsphinx do not play well together. Therefore, disable TQDM
+# for the documentation build.
+# (`Content block expected for the "raw" directive; none found.`)
+os.environ["TQDM_DISABLE"] = "1"
+
 # -- General configuration ------------------------------------------------
 
 # If your documentation needs a minimal Sphinx version, state it here.

pypesto/engine/base.py

@@ -13,7 +13,7 @@ def __init__(self):
 
  @abc.abstractmethod
  def execute(
- self, tasks: list[Task], progress_bar: bool = True
+ self, tasks: list[Task], progress_bar: bool = None
  ) -> list[Any]:
  """Execute tasks.
 
@@ -22,6 +22,6 @@ def execute(
  tasks:
  List of tasks to execute.
  progress_bar:
- Whether to display a progress bar. Defaults to ``True``.
+ Whether to display a progress bar.
  """
  raise NotImplementedError("This engine is not intended to be called.")

pypesto/engine/mpi_pool.py

@@ -5,8 +5,8 @@
 import cloudpickle as pickle
 from mpi4py import MPI
 from mpi4py.futures import MPIPoolExecutor
-from tqdm import tqdm
 
+from ..util import tqdm
 from .base import Engine
 from .task import Task
 
@@ -32,7 +32,7 @@ def __init__(self):
  super().__init__()
 
  def execute(
- self, tasks: list[Task], progress_bar: bool = True
+ self, tasks: list[Task], progress_bar: bool = None
  ) -> list[Any]:
  """
  Pickle tasks and distribute work to workers.
@@ -42,7 +42,7 @@ def execute(
  tasks:
  List of :class:`pypesto.engine.Task` to execute.
  progress_bar:
- Whether to display a progress bar. Defaults to ``True``.
+ Whether to display a progress bar.
 
  Returns
  -------
@@ -55,6 +55,6 @@ def execute(
 
  with MPIPoolExecutor() as executor:
  results = executor.map(
- work, tqdm(pickled_tasks, disable=not progress_bar)
+ work, tqdm(pickled_tasks, enable=progress_bar)
  )
  return results

pypesto/engine/multi_process.py

@@ -5,8 +5,8 @@
 from typing import Any, Union
 
 import cloudpickle as pickle
-from tqdm import tqdm
 
+from ..util import tqdm
 from .base import Engine
 from .task import Task
 
@@ -52,7 +52,7 @@ def __init__(
  self.method: str = method
 
  def execute(
- self, tasks: list[Task], progress_bar: bool = True
+ self, tasks: list[Task], progress_bar: bool = None
  ) -> list[Any]:
  """Pickle tasks and distribute work over parallel processes.
 
@@ -61,7 +61,7 @@ def execute(
  tasks:
  List of :class:`pypesto.engine.Task` to execute.
  progress_bar:
- Whether to display a progress bar. Defaults to ``True``.
+ Whether to display a progress bar.
 
  Returns
  -------
@@ -81,7 +81,7 @@ def execute(
  tqdm(
  pool.imap(work, pickled_tasks),
  total=len(pickled_tasks),
- disable=not progress_bar,
+ enable=progress_bar,
  ),
  )
 

pypesto/engine/multi_thread.py

@@ -5,8 +5,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Union
 
-from tqdm import tqdm
-
+from ..util import tqdm
 from .base import Engine
 from .task import Task
 
@@ -43,7 +42,7 @@ def __init__(self, n_threads: Union[int, None] = None):
  self.n_threads: int = n_threads
 
  def execute(
- self, tasks: list[Task], progress_bar: bool = True
+ self, tasks: list[Task], progress_bar: bool = None
  ) -> list[Any]:
  """Deepcopy tasks and distribute work over parallel threads.
 
@@ -70,7 +69,7 @@ def execute(
  tqdm(
  pool.map(work, copied_tasks),
  total=len(copied_tasks),
- disable=not progress_bar,
+ enable=progress_bar,
  ),
  )
 

pypesto/engine/single_core.py

@@ -1,8 +1,7 @@
 """Engines without parallelization."""
 from typing import Any
 
-from tqdm import tqdm
-
+from ..util import tqdm
 from .base import Engine
 from .task import Task
 
@@ -18,7 +17,7 @@ def __init__(self):
  super().__init__()
 
  def execute(
- self, tasks: list[Task], progress_bar: bool = True
+ self, tasks: list[Task], progress_bar: bool = None
  ) -> list[Any]:
  """Execute all tasks in a simple for loop sequentially.
 
@@ -34,7 +33,10 @@ def execute(
  A list of results.
  """
  results = []
- for task in tqdm(tasks, disable=not progress_bar):
+ for task in tqdm(
+ tasks,
+ enable=progress_bar,
+ ):
  results.append(task.execute())
 
  return results

pypesto/ensemble/ensemble.py

@@ -877,7 +877,7 @@ def predict(
  include_llh_weights: bool = False,
  include_sigmay: bool = False,
  engine: Engine = None,
- progress_bar: bool = True,
+ progress_bar: bool = None,
  ) -> EnsemblePrediction:
  """
  Run predictions for a full ensemble.

pypesto/optimize/optimize.py

@@ -29,7 +29,7 @@ def minimize(
  startpoint_method: Union[StartpointMethod, Callable, bool] = None,
  result: Result = None,
  engine: Engine = None,
- progress_bar: bool = True,
+ progress_bar: bool = None,
  options: OptimizeOptions = None,
  history_options: HistoryOptions = None,
  filename: Union[str, Callable, None] = None,

pypesto/profile/profile.py

@@ -25,7 +25,7 @@ def parameter_profile(
  result_index: int = 0,
  next_guess_method: Union[Callable, str] = 'adaptive_step_regression',
  profile_options: ProfileOptions = None,
- progress_bar: bool = True,
+ progress_bar: bool = None,
  filename: Union[str, Callable, None] = None,
  overwrite: bool = False,
 ) -> Result:

pypesto/sample/adaptive_metropolis.py

@@ -86,7 +86,7 @@ def default_options(cls):
  # target acceptance rate
  'target_acceptance_rate': 0.234,
  # show progress
- 'show_progress': True,
+ 'show_progress': None,
  }
 
  def initialize(self, problem: Problem, x0: np.ndarray):

pypesto/sample/metropolis.py

@@ -1,12 +1,12 @@
 from typing import Dict, Sequence, Union
 
 import numpy as np
-from tqdm import tqdm
 
 from ..history import NoHistory
 from ..objective import NegLogPriors, ObjectiveBase
 from ..problem import Problem
 from ..result import McmcPtResult
+from ..util import tqdm
 from .sampler import InternalSample, InternalSampler
 
 
@@ -51,7 +51,7 @@ def default_options(cls):
  """Return the default options for the sampler."""
  return {
  'std': 1.0, # the proposal standard deviation
- 'show_progress': True, # whether to show the progress
+ 'show_progress': None, # whether to show the progress
  }
 
  def initialize(self, problem: Problem, x0: np.ndarray):
@@ -73,10 +73,10 @@ def sample(self, n_samples: int, beta: float = 1.0):
  lpost = -self.trace_neglogpost[-1]
  lprior = -self.trace_neglogprior[-1]
 
- show_progress = self.options['show_progress']
+ show_progress = self.options.get('show_progress', None)
 
  # loop over iterations
- for _ in tqdm(range(int(n_samples)), disable=not show_progress):
+ for _ in tqdm(range(int(n_samples)), enable=show_progress):
  # perform step
  x, lpost, lprior = self._perform_step(
  x=x, lpost=lpost, lprior=lprior, beta=beta

pypesto/sample/parallel_tempering.py

@@ -2,10 +2,10 @@
 from typing import Dict, List, Sequence, Union
 
 import numpy as np
-from tqdm import tqdm
 
 from ..problem import Problem
 from ..result import McmcPtResult
+from ..util import tqdm
 from .sampler import InternalSampler, Sampler
 
 
@@ -70,7 +70,7 @@ def default_options(cls) -> Dict:
  'max_temp': 5e4,
  'exponent': 4,
  'temper_log_posterior': False,
- 'show_progress': True,
+ 'show_progress': None,
  }
 
  def initialize(
@@ -89,9 +89,9 @@ def initialize(
 
  def sample(self, n_samples: int, beta: float = 1.0):
  """Sample and swap in between samplers."""
- show_progress = self.options['show_progress']
+ show_progress = self.options.get('show_progress', None)
  # loop over iterations
- for i_sample in tqdm(range(int(n_samples)), disable=not show_progress):
+ for i_sample in tqdm(range(int(n_samples)), enable=show_progress):
  # TODO test
  # sample
  for sampler, beta in zip(self.samplers, self.betas):

pypesto/util.py

@@ -11,6 +11,7 @@
 
 import numpy as np
 from scipy import cluster
+from tqdm import tqdm as _tqdm
 
 
 def _check_none(fun: Callable[..., Any]) -> Callable[..., Union[Any, None]]:
@@ -295,3 +296,40 @@ def delete_nan_inf(
  )
  )
  return x, fvals[finite_fvals]
+
+
+def tqdm(*args, enable: bool = None, **kwargs):
+ """
+ Create a progress bar using tqdm.
+
+ Parameters
+ ----------
+ args:
+ Arguments passed to tqdm.
+ enable:
+ Whether to enable the progress bar.
+ If None, use tqdm defaults.
+ Mutually exclusive with `disable`.
+ kwargs:
+ Keyword arguments passed to tqdm.
+
+ Returns
+ -------
+ progress_bar:
+ A progress bar.
+ """
+ # Drop the `disable` argument unless it is not-None.
+ # This way, we don't interfere with TQDM_DISABLE or other global
+ # tqdm settings.
+ disable = kwargs.pop("disable", None)
+
+ if enable is not None:
+ if disable is not None and enable != disable:
+ raise ValueError(
+ "Contradicting values for `enable` and `disable` passed."
+ )
+ disable = not enable
+
+ if disable is not None:
+ kwargs["disable"] = disable
+ return _tqdm(*args, **kwargs)