cleanup

Author: dwhswenson

paths_cli/commands/md.py

@@ -1,8 +1,6 @@
-import functools
-import operator
-
 import click
 
+import paths_cli.utils
 from paths_cli.parameters import (INPUT_FILE, OUTPUT_FILE, ENGINE,
                                   MULTI_ENSEMBLE, INIT_SNAP)
 
@@ -39,7 +37,6 @@ def __init__(self, ensembles):
         self.satisfied = {ens: False for ens in ensembles}
 
     def _check_previous_frame(self, trajectory, start, unsatisfied):
-        # TODO: add some debug logging in here
         if -start > len(trajectory):
             # we've done the whole traj; don't keep going
             return False
@@ -91,10 +88,8 @@ def md_main(output_storage, engine, ensembles, nsteps, initial_frame):
         continue_cond = paths.LengthEnsemble(nsteps).can_append
 
     trajectory = engine.generate(initial_frame, running=continue_cond)
-    if output_storage is not None:
-        output_storage.save(trajectory)
-        output_storage.tags['final_conditions'] = trajectory
-
+    paths_cli.utils.tag_final_result(trajectory, output_storage,
+                                     'final_conditions')
     return trajectory, None
 
 CLI = md

paths_cli/commands/visit_all.py

@@ -1,5 +1,6 @@
 import click
 
+import paths_cli.utils
 from paths_cli.parameters import (INPUT_FILE, OUTPUT_FILE, ENGINE, STATES,
                                   INIT_SNAP)
 
@@ -34,9 +35,8 @@ def visit_all_main(output_storage, states, engine, initial_frame):
     timestep = getattr(engine, 'timestep', None)
     visit_all_ens = paths.VisitAllStatesEnsemble(states, timestep=timestep)
     trajectory = engine.generate(initial_frame, [visit_all_ens.can_append])
-    if output_storage is not None:
-        output_storage.save(trajectory)
-        output_storage.tags['final_conditions'] = trajectory
+    paths_cli.utils.tag_final_result(trajectory, output_storage,
+                                     'final_conditions')
 
     return trajectory, None  # no simulation object to return here
 

paths_cli/utils.py

@@ -0,0 +1,4 @@
+def tag_final_result(result, storage, tag='final_conditions'):
+    if storage:
+        storage.save(result)
+        storage.tags[tag] = result