Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add progress bar keywords #4085

Merged
1 change: 1 addition & 0 deletions package/AUTHORS
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,7 @@ Chronological list of authors
- Domenico Marson
- Ahmed Salah Ghoneim
- Alexander Schlaich
- Egor Marin

External code
-------------
Expand Down
1 change: 1 addition & 0 deletions package/CHANGELOG
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ Enhancements
* AuxReaders are now pickle-able and copy-able (Issue #1785, PR #3887)
* Add pickling support for Atom, Residue, Segment, ResidueGroup
and SegmentGroup. (PR #3953)
* Add `progressbar_kwargs` parameter to `AnalysisBase.run` method, allowing to modify description, position etc of tqdm progressbars.

Changes
* As per NEP29 the minimum supported NumPy version has been raised to 1.21
Expand Down
13 changes: 11 additions & 2 deletions package/MDAnalysis/analysis/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -390,7 +390,7 @@ def _conclude(self):
pass # pylint: disable=unnecessary-pass

def run(self, start=None, stop=None, step=None, frames=None,
verbose=None):
verbose=None, *, progressbar_kwargs={}):
"""Perform the calculation

Parameters
Expand All @@ -411,12 +411,19 @@ def run(self, start=None, stop=None, step=None, frames=None,

verbose : bool, optional
Turn on verbosity

progressbar_kwargs : dict, optional
ProgressBar keywords with custom parameters regarding progress bar position, etc;
see :class:`MDAnalysis.lib.log.ProgressBar` for full list.


.. versionchanged:: 2.2.0
Added ability to analyze arbitrary frames by passing a list of
frame indices in the `frames` keyword argument.

.. versionchanged:: 2.5.0
Add `progressbar_kwargs` parameter,
allowing to modify description, position etc of tqdm progressbars
"""
logger.info("Choosing frames to analyze")
# if verbose unchanged, use class default
Expand All @@ -429,9 +436,11 @@ def run(self, start=None, stop=None, step=None, frames=None,
self._prepare()
logger.info("Starting analysis loop over %d trajectory frames",
self.n_frames)

for i, ts in enumerate(ProgressBar(
self._sliced_trajectory,
verbose=verbose)):
verbose=verbose,
**progressbar_kwargs)):
self._frame_index = i
self._ts = ts
self.frames[i] = ts.frame
Expand Down
6 changes: 6 additions & 0 deletions testsuite/MDAnalysisTests/analysis/test_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -274,6 +274,12 @@ def test_verbose_progressbar_run(u, capsys):
actual = err.strip().split('\r')[-1]
assert actual[:24] == expected[:24]

def test_verbose_progressbar_run_with_kwargs(u, capsys):
an = FrameAnalysis(u.trajectory).run(verbose=True, progressbar_kwargs={'desc':'custom'})
out, err = capsys.readouterr()
expected = u'custom: 100%|██████████| 98/98 [00:00<00:00, 8799.49it/s]'
actual = err.strip().split('\r')[-1]
assert actual[:30] == expected[:30]

def test_incomplete_defined_analysis(u):
with pytest.raises(NotImplementedError):
Expand Down