more tests; docstrings

Author: dwhswenson

paths_cli/commands/md.py

@@ -39,25 +39,43 @@ def md(input_file, output_file, engine, ensemble, nsteps, init_frame):
     )
 
 class ProgressReporter(object):
+    """Generic class for a callable that reports progress.
+
+    Base class for ends-with-ensemble and fixed-length tricks.
+
+    Parameters
+    ----------
+    timestep : Any
+        timestep, optionally with units
+    update_freq : int
+        how often to report updates
+    """
     def __init__(self, timestep, update_freq):
         self.timestep = timestep
         self.update_freq = update_freq
 
     def steps_progress_string(self, n_steps):
+        """Return string for number of frames run and time elapsed
+
+        Not newline-terminated.
+        """
         report_str = "Ran {n_steps} frames"
         if self.timestep is not None:
-            report_str += " [{}]".format(str(n_steps * timestep))
+            report_str += " [{}]".format(str(n_steps * self.timestep))
         report_str += '.'
-        return report_str
+        return report_str.format(n_steps=n_steps)
 
     def progress_string(self, n_steps):
+        """Return the progress string. Subclasses may override.
+        """
         report_str = self.steps_progress_string(n_steps) + "\n"
         return report_str.format(n_steps=n_steps)
 
-
-    def report_progress(self, n_steps):
+    def report_progress(self, n_steps, force=False):
+        """Report the progress to the terminal.
+        """
         import openpathsampling as paths
-        if n_steps % self.update_freq == 0:
+        if (n_steps % self.update_freq == 0) or force:
             string = self.progress_string(n_steps)
             paths.tools.refresh_output(string)
 
@@ -77,6 +95,10 @@ class EnsembleSatisfiedContinueConditions(ProgressReporter):
     ----------
     ensembles: List[:class:`openpathsampling.Ensemble`]
         the ensembles to satisfy
+    timestep : Any
+        timestep, optionally with units
+    update_freq : int
+        how often to report updates
     """
     def __init__(self, ensembles, timestep=None, update_freq=10):
         super().__init__(timestep, update_freq)
@@ -136,7 +158,16 @@ def __call__(self, trajectory, trusted=False):
 
 
 class FixedLengthContinueCondition(ProgressReporter):
-    """
+    """Continuation condition for fixed-length runs.
+
+    Parameters
+    ----------
+    length : int
+        final length of the trajectory in frames
+    timestep : Any
+        timestep, optionally with units
+    update_freq : int
+        how often to report updates
     """
     def __init__(self, length, timestep=None, update_freq=10):
         super().__init__(timestep, update_freq)
@@ -161,6 +192,7 @@ def md_main(output_storage, engine, ensembles, nsteps, initial_frame):
         continue_cond = FixedLengthContinueCondition(nsteps)
 
     trajectory = engine.generate(initial_frame, running=continue_cond)
+    continue_cond.report_progress(len(trajectory) - 1, force=True)
     paths_cli.utils.tag_final_result(trajectory, output_storage,
                                      'final_conditions')
     return trajectory, None

paths_cli/tests/commands/test_md.py

@@ -11,6 +11,32 @@
 from openpathsampling.tests.test_helpers import \
         make_1d_traj, CalvinistDynamics
 
+class TestProgressReporter(object):
+    def setup(self):
+        self.progress = ProgressReporter(timestep=None, update_freq=5)
+
+    @pytest.mark.parametrize('timestep', [None, 0.1])
+    def test_progress_string(self, timestep):
+        progress = ProgressReporter(timestep, update_freq=5)
+        expected = "Ran 25 frames"
+        if timestep is not None:
+            expected += " [2.5]"
+        expected += '.\n'
+        assert progress.progress_string(25) == expected
+
+    @pytest.mark.parametrize('n_steps', [0, 5, 6])
+    @pytest.mark.parametrize('force', [True, False])
+    @patch('openpathsampling.tools.refresh_output',
+           lambda s: print(s, end=''))
+    def test_report_progress(self, n_steps, force, capsys):
+        self.progress.report_progress(n_steps, force)
+        expected = "Ran {n_steps} frames.\n".format(n_steps=n_steps)
+        out, err = capsys.readouterr()
+        if (n_steps in [0, 5]) or force:
+            assert out == expected
+        else:
+            assert out == ""
+
 class TestEnsembleSatisfiedContinueConditions(object):
     def setup(self):
         cv = paths.CoordinateFunctionCV('x', lambda x: x.xyz[0][0])