Add progress bar keywords to AnalysisBase.run (#4085)

* Add progress bar keywords * Update authors list for first isssue #4085 * Rename `pbar_kwargs` to `progressbar_kwargs` * Update changelog for PR #4085 * Remove unnecessary use of `rms` * Update testcase to include number of iterations * Update AUTHORS * Update CHANGELOG to make newest first --------- Co-authored-by: Hugo MacDermott-Opeskin <[email protected]>
MDAnalysis · Mar 29, 2023 · 5b01ed8 · 5b01ed8
1 parent 75f0fb9
commit 5b01ed8
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 2 deletions.
diff --git a/package/AUTHORS b/package/AUTHORS
@@ -214,6 +214,7 @@ Chronological list of authors
   - Alexander Schlaich
   - Josh Vermaas
   - Xiaoxu Ruan
+  - Egor Marin
 
 External code
 -------------

diff --git a/package/CHANGELOG b/package/CHANGELOG
@@ -37,6 +37,7 @@ Fixes
     (Issue #3336)
 
 Enhancements
+  * Add `progressbar_kwargs` parameter to `AnalysisBase.run` method, allowing to modify description, position etc of tqdm progressbars.
   * Add a nojump transformation, which unwraps trajectories so that particle
     paths are continuous. (Issue #3703, PR #4031)
   * Added AtomGroup TopologyAttr to calculate gyration moments (Issue #3904,

diff --git a/package/MDAnalysis/analysis/base.py b/package/MDAnalysis/analysis/base.py
@@ -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
@@ -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
@@ -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

diff --git a/testsuite/MDAnalysisTests/analysis/test_base.py b/testsuite/MDAnalysisTests/analysis/test_base.py
@@ -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):