Add INIT_SNAP parameter

Author: dwhswenson

paths_cli/parameters.py

@@ -126,6 +126,18 @@ def init_traj_fallback(parameter, storage, name):
             return storage.trajectories[0]
 
 
+def init_snap_fallback(parameter, storage, name):
+    # this is structured so that other things can be added to it later
+    result = None
+
+    if name is None:
+        result = storage.tags['initial_snapshot']
+        if result:
+            return result
+
+        if len(storage.snapshots) == 2:
+            # this is really only 1 snapshot; reversed copy gets saved
+            return storage.snapshots[0]
 
 ENGINE = OPSStorageLoadSingle(
     param=Option('-e', '--engine', help="identifer for the engine"),
@@ -142,12 +154,20 @@ def init_traj_fallback(parameter, storage, name):
 INIT_CONDS = OPSStorageLoadSingle(
     param=Option('-t', '--init-conds',
                  help=("identifier for initial conditions "
-                       + "(sample set or trajectory")),
+                       + "(sample set or trajectory)")),
     store='tags',
     num_store='samplesets',
     fallback=init_traj_fallback
 )
 
+INIT_SNAP = OPSStorageLoadSingle(
+    param=Option('-f', '--init-frame',
+                 help="identifier for initial snapshot"),
+    store='tags',
+    num_store='snapshots',
+    fallback=init_snap_fallback
+)
+
 CVS = OPSStorageLoadNames(
     param=Option('--cv', type=str, multiple=True,
                  help='name of CV; may select more than once'),

paths_cli/tests/test_parameters.py

@@ -233,8 +233,41 @@ def test_get_none(self, num_in_file):
         obj = INIT_CONDS.get(st, None)
         assert obj == stored_things[num_in_file - 1]
 
+class TestINIT_SNAP(ParamInstanceTest):
+    PARAMETER = INIT_SNAP
+    def setup(self):
+        super(TestINIT_SNAP, self).setup()
+        traj = make_1d_traj([1.0, 2.0])
+        self.other_snap = traj[0]
+        self.init_snap = traj[1]
+
+    def create_file(self, getter):
+        filename = self._filename(getter)
+        storage = paths.Storage(filename, 'w')
+        storage.save(self.other_snap)
+        if getter != 'none-num':
+            storage.save(self.init_snap)
+            storage.tags['initial_snapshot'] = self.init_snap
+        storage.close()
+        return filename
 
-        pass
+    @pytest.mark.parametrize('getter', ['none-num', 'tag', 'number',
+                                        'none-tag'])
+    def test_get(self, getter):
+        # get by number is 2 because of snapshot duplication in storage
+        get_arg = {'none-num': None,
+                   'none-tag': None,
+                   'tag': 'initial_snapshot',
+                   'number': 2}[getter]
+        expected = {'none-num': self.other_snap,
+                    'none-tag': self.init_snap,
+                    'tag': self.init_snap,
+                    'number': self.init_snap}[getter]
+        filename = self.create_file(getter)
+        storage = paths.Storage(filename, mode='r')
+
+        obj = self.PARAMETER.get(storage, get_arg)
+        assert obj == expected
 
 
 class MultiParamInstanceTest(ParamInstanceTest):