make strip-snapshots stuff more reusable

Author: dwhswenson

paths_cli/commands/strip_snapshots.py

@@ -21,26 +21,37 @@ def strip_snapshots(input_file, output_file, cv, blocksize):
     By giving the --cv option (once for each CV), you can select to only
     save certain CVs. If you do not give that option, all CVs will be saved.
     """
+    input_storage = INPUT_FILE.get(input_file)
     return strip_snapshots_main(
-        input_storage=INPUT_FILE.get(input_file),
+        input_storage=input_storage,
         output_storage=OUTPUT_FILE.get(output_file),
-        cvs=CVS.get(cv),
+        cvs=CVS.get(input_storage, cv),
         blocksize=blocksize
     )
 
 
-def block_precompute_cvs(cvs, proxy_snaps, blocksize):
-    n_snaps = len(proxy_snaps)
-    n_blocks = (n_snaps // blocksize) +1
+def make_blocks(listlike, blocksize):
+    n_objs = len(listlike)
+    n_blocks = (n_objs // blocksize) + 1
+    minval = lambda i: i * blocksize
+    maxval = lambda i: min((i + 1) * blocksize, n_objs)
+    blocks = [listlike[minval(i):maxval(i)] for i in range(n_blocks)]
+    return blocks
 
-    blocks = [proxy_snaps[i*blocksize: min((i+1)*blocksize, n_snaps)]
-              for i in range(n_blocks)]
 
-    for block_num in tqdm(range(n_blocks), leave=False):
-        block = blocks[block_num]
-        for cv in cvs:
-            cv.enable_diskcache()
-            _ = cv(block)
+def precompute_cvs(cvs, block):
+    for cv in cvs:
+        cv.enable_diskcache()
+        _ = cv(block)
+
+
+def rewrite_file(stage_names, stage_mapping):
+    stages = tqdm(stage_names)
+    for stage in stages:
+        stages.set_description("%s" % stage)
+        store_func, inputs = stage_mapping[stage]
+        for obj in tqdm(inputs):
+            store_func(obj)
 
 
 def strip_snapshots_main(input_storage, output_storage, cvs, blocksize):
@@ -50,24 +61,35 @@ def strip_snapshots_main(input_storage, output_storage, cvs, blocksize):
     # save template
     output_storage.save(input_storage.snapshots[0])
 
+    precompute_func = lambda inps: precompute_cvs(cvs, inps)
     snapshot_proxies = input_storage.snapshots.all().as_proxies()
-    stages = tqdm(['precompute', 'cvs', 'snapshots', 'trajectories',
-                   'steps'])
-    for stage in stages:
-        stages.set_description("%s" % stage)
-        if stage == 'precompute':
-            block_precompute_cvs(cvs, snapshot_proxies, blocksize)
-        else:
-            store_func, inputs = {
-                'cvs': (output_storage.cvs.save, input_storage.cvs),
-                'snapshots': (output_storage.snapshots.mention,
-                              snapshot_proxies),
-                'trajectories': (output_storage.trajectories.mention,
-                                 input_storage.trajectories),
-                'steps': (output_storage.steps.save, input_storage.steps)
-            }[stage]
-            for obj in tqdm(inputs):
-                store_func(obj)
+    snapshot_blocks = make_blocks(snapshot_proxies, blocksize)
+    stages = ['precompute', 'cvs', 'snapshots', 'trajectories', 'steps']
+    stage_mapping = {
+        'precompute': (precompute_func, snapshot_blocks),
+        'cvs': (output_storage.cvs.save, input_storage.cvs),
+        'snapshots': (output_storage.snapshots.mention,
+                      snapshot_proxies),
+        'trajectories': (output_storage.trajectories.mention,
+                         input_storage.trajectories),
+        'steps': (output_storage.steps.save, input_storage.steps),
+    }
+    return rewrite_file(stages, stage_mapping)
+    # for stage in stages:
+        # stages.set_description("%s" % stage)
+        # if stage == 'precompute':
+            # block_precompute_cvs(cvs, snapshot_proxies, blocksize)
+        # else:
+            # store_func, inputs = {
+                # 'cvs': (output_storage.cvs.save, input_storage.cvs),
+                # 'snapshots': (output_storage.snapshots.mention,
+                              # snapshot_proxies),
+                # 'trajectories': (output_storage.trajectories.mention,
+                                 # input_storage.trajectories),
+                # 'steps': (output_storage.steps.save, input_storage.steps)
+            # }[stage]
+            # for obj in tqdm(inputs):
+                # store_func(obj)
 
 
 CLI = strip_snapshots