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 @@ -210,6 +210,7 @@ Chronological list of authors
- Moritz Schaeffler
- Xu Hong Chen
- Domenico Marson
- Egor Marin

External code
-------------
Expand Down
10 changes: 8 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, *, pbar_kwargs={}):
marinegor marked this conversation as resolved.
Show resolved Hide resolved
"""Perform the calculation

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

verbose : bool, optional
Turn on verbosity

pbar_kwargs : dict, optional
marinegor marked this conversation as resolved.
Show resolved Hide resolved
marinegor marked this conversation as resolved.
Show resolved Hide resolved
ProgressBar keywords with custom parameters regarding progress bar position, etc;
see :class:`MDAnalysis.lib.log.ProgressBar` for full list.


.. versionchanged:: 2.2.0
Expand All @@ -429,9 +433,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,
**pbar_kwargs)):
self._frame_index = i
self._ts = ts
self.frames[i] = ts.frame
Expand Down
8 changes: 7 additions & 1 deletion testsuite/MDAnalysisTests/analysis/test_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@
from numpy.testing import assert_equal, assert_almost_equal

import MDAnalysis as mda
from MDAnalysis.analysis import base
from MDAnalysis.analysis import base, rms
marinegor marked this conversation as resolved.
Show resolved Hide resolved

from MDAnalysisTests.datafiles import PSF, DCD, TPR, XTC
from MDAnalysisTests.util import no_deprecated_call
Expand Down 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, pbar_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[:24] == expected[:24]
marinegor marked this conversation as resolved.
Show resolved Hide resolved

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