Convert output of INIT_CONDS parameter to trajs

openpathsampling · Aug 25, 2024 · d1c7539 · d1c7539
1 parent de53902
commit d1c7539
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 4 deletions.
diff --git a/paths_cli/parameters.py b/paths_cli/parameters.py
@@ -18,7 +18,31 @@
     store='schemes',
 )
 
-INIT_CONDS = OPSStorageLoadMultiple(
+class InitCondsLoader(OPSStorageLoadMultiple):
+    def _extract_trajectories(self, obj):
+        import openpathsampling as paths
+        if isinstance(obj, paths.SampleSet):
+            yield from (s.trajectory for s in obj)
+        elif isinstance(obj, paths.Sample):
+            yield obj.trajectory
+        elif isinstance(obj, paths.Trajectory):
+            yield obj
+        elif isinstance(obj, paths.BaseSnapshot):
+            yield paths.Trajectory([obj])
+        elif isinstance(obj, list):
+            for o in obj:
+                yield from self._extract_trajectories(o)
+        else:
+            raise RuntimeError("Unknown initial conditions type: "
+                               f"{obj} (type: {type(obj)}")
+
+    def get(self, storage, names):
+        results = super().get(storage, names)
+        final_results = list(self._extract_trajectories(results))
+        return final_results
+
+
+INIT_CONDS = InitCondsLoader(
     param=Option('-t', '--init-conds', multiple=True,
                  help=("identifier for initial conditions "
                        + "(sample set or trajectory)" + HELP_MULTIPLE)),

diff --git a/paths_cli/tests/test_parameters.py b/paths_cli/tests/test_parameters.py
@@ -243,8 +243,8 @@ def test_get(self, getter):
         storage = paths.Storage(filename, mode='r')
         get_type, getter_style = self._parse_getter(getter)
         expected = {
-            'sset': self.sample_set,
-            'traj': self.traj
+            'sset': [s.trajectory for s in self.sample_set],
+            'traj': [self.traj]
         }[get_type]
         get_arg = {
             'name': 'traj',
@@ -277,7 +277,14 @@ def test_get_none(self, num_in_file):
 
         st = paths.Storage(filename, mode='r')
         obj = INIT_CONDS.get(st, None)
-        assert obj == stored_things[num_in_file - 1]
+        # TODO: fix this for all being trajectories
+        expected = [
+            [self.traj],
+            [s.trajectory for s in self.sample_set],
+            [s.trajectory for s in self.other_sample_set],
+            [s.trajectory for s in self.other_sample_set],
+        ]
+        assert obj == expected[num_in_file - 1]
 
     def test_get_multiple(self):
         filename = self.create_file('number-traj')