Merge branch 'main' into cmd-bootstrap-init

Author: dwhswenson

paths_cli/commands/load_trajectory.py

@@ -0,0 +1,51 @@
+import click
+
+import paths_cli.utils
+from paths_cli.parameters import APPEND_FILE, ENGINE, MULTI_TAG
+from paths_cli import OPSCommandPlugin
+
+
+@click.command(
+    "load-trajectory",
+    short_help="Load an external trajectory file",
+)
+@click.argument('traj_file')
+@click.option(
+    '--top',
+    help=(
+        "Topology file (typically PDB). Only for required "
+        "formats."
+    ),
+    default=None,
+)
+@APPEND_FILE.clicked(required=True)
+@MULTI_TAG.clicked()
+def load_trajectory(traj_file, top, append_file, tag):
+    """Load a trajectory from a file.
+
+    This uses MDTraj under the hood, and can load any file format that
+    MDTraj can. NB: This stores in a format based on OpenMM snapshots.
+    Trajectories loaded this way will work with engines compatible with
+    that input (e.g., GROMACS).
+    """
+    from openpathsampling.engines.openmm.tools import ops_load_trajectory
+    if top:
+        traj = ops_load_trajectory(traj_file, top=top)
+    else:
+        traj = ops_load_trajectory(traj_file)
+
+    storage = APPEND_FILE.get(append_file)
+    storage.save(traj)
+    for tag_name in tag:
+        storage.tags[tag_name] = traj
+
+
+PLUGIN = OPSCommandPlugin(
+    command=load_trajectory,
+    section="Miscellaneous",
+    requires_ops=(1, 6),
+    requires_cli=(0, 4),
+)
+
+if __name__ == "__main__":  # -no-cov-
+    load_trajectory()

paths_cli/parameters.py

@@ -18,7 +18,29 @@
     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, 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)),

paths_cli/tests/commands/test_equilibrate.py

@@ -11,7 +11,7 @@
 def print_test(output_storage, scheme, init_conds, multiplier, extra_steps):
     print(isinstance(output_storage, paths.Storage))
     print(scheme.__uuid__)
-    print(init_conds.__uuid__)
+    print([o.__uuid__ for o in init_conds])
     print(multiplier, extra_steps)
 
 
@@ -31,8 +31,10 @@ def test_equilibrate(tps_fixture):
             ["setup.nc", "-o", "foo.nc"]
         )
         out_str = "True\n{schemeid}\n{condsid}\n1 0\n"
-        expected_output = out_str.format(schemeid=scheme.__uuid__,
-                                         condsid=init_conds.__uuid__)
+        expected_output = out_str.format(
+            schemeid=scheme.__uuid__,
+            condsid=[o.trajectory.__uuid__ for o in init_conds],
+        )
         assert results.exit_code == 0
         assert results.output == expected_output
 

paths_cli/tests/commands/test_load_trajectory.py

@@ -0,0 +1,89 @@
+from click.testing import CliRunner
+from contextlib import contextmanager
+import pytest
+from importlib import resources
+from openpathsampling.tests.test_helpers import data_filename
+import openpathsampling as paths
+
+from paths_cli.commands.load_trajectory import *
+
+
+@contextmanager
+def run_load_trajectory(args):
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        storage = paths.Storage("setup.nc", 'w')
+        storage.close()
+        results = runner.invoke(
+            load_trajectory,
+            args
+        )
+        assert results.exit_code == 0
+        st = paths.Storage("setup.nc", mode='r')
+        assert len(st.trajectories) == 1
+        yield st
+
+
+@pytest.mark.parametrize("with_top", [True, False])
+@pytest.mark.parametrize("with_tag", [True, False])
+def test_load_trajectory_pdb(with_top, with_tag):
+    # test that we can load a PDB file with or without topology; also tests
+    # that the taging works correctly
+    pytest.importorskip("openmm")
+    pytest.importorskip("mdtraj")
+    pdb_path = data_filename("ala_small_traj.pdb")
+    out_file = "setup.nc"
+    args = [
+        pdb_path,
+        '--append-file', out_file,
+    ]
+    if with_top:
+        args.extend(['--top', pdb_path])
+
+    if with_tag:
+        args.extend(['--tag', 'init_snap'])
+
+    with run_load_trajectory(args) as st:
+        traj = st.trajectories[0]
+        assert len(traj) == 10
+        if with_tag:
+            tagged = st.tags['init_snap']
+            assert tagged == traj
+
+def test_load_trajectory_trr():
+    pytest.importorskip("openmm")
+    pytest.importorskip("mdtraj")
+    trr = data_filename("gromacs_engine/project_trr/0000000.trr")
+    gro = data_filename("gromacs_engine/conf.gro")
+    out_file = "setup.nc"
+    args = [
+        trr,
+        '--append-file', out_file,
+        '--top', gro,
+    ]
+    with run_load_trajectory(args) as st:
+        traj = st.trajectories[0]
+        assert len(traj) == 4
+
+def test_load_trajectory_bad_topology():
+    pytest.importorskip("openmm")
+    pytest.importorskip("mdtraj")
+    trr = data_filename("gromacs_engine/project_trr/0000000.trr")
+    pdb = data_filename("tip4p_water.pdb")
+    out_file = "setup.nc"
+    args = [
+        trr,
+        '--append-file', out_file,
+        '--top', pdb,
+    ]
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        storage = paths.Storage("setup.nc", 'w')
+        storage.close()
+        result = runner.invoke(
+            load_trajectory,
+            args
+        )
+        assert result.exit_code == 1
+        assert "topology" in str(result.exception)
+        assert "same atoms" in str(result.exception)

paths_cli/tests/commands/test_pathsampling.py

@@ -11,7 +11,7 @@
 def print_test(output_storage, scheme, init_conds, n_steps):
     print(isinstance(output_storage, paths.Storage))
     print(scheme.__uuid__)
-    print(init_conds.__uuid__)
+    print([traj.__uuid__ for traj in init_conds])
     print(n_steps)
 
 @patch('paths_cli.commands.pathsampling.pathsampling_main', print_test)
@@ -26,7 +26,8 @@ def test_pathsampling(tps_fixture):
 
         results = runner.invoke(pathsampling, ['setup.nc', '-o', 'foo.nc',
                                                '-n', '1000'])
-        expected_output = (f"True\n{scheme.__uuid__}\n{init_conds.__uuid__}"
+        initcondsid = [samp.trajectory.__uuid__ for samp in init_conds]
+        expected_output = (f"True\n{scheme.__uuid__}\n{initcondsid}"
                            "\n1000\n")
 
         assert results.output == expected_output

paths_cli/tests/test_parameters.py

@@ -214,8 +214,12 @@ def create_file(self, getter):
         get_type, getter_style = self._parse_getter(getter)
         main, other = {
             'traj': (self.traj, self.other_traj),
-            'sset': (self.sample_set, self.other_sample_set)
+            'sset': (self.sample_set, self.other_sample_set),
+            'samp': (self.sample_set[0], self.other_sample_set[0]),
         }[get_type]
+        if get_type == 'samp':
+            storage.save(main)
+            storage.save(other)
         if get_type == 'sset':
             storage.save(self.sample_set)
             storage.save(self.other_sample_set)
@@ -231,20 +235,23 @@ def create_file(self, getter):
 
         if other_tag:
             storage.tags[other_tag] = other
+
         storage.close()
         return filename
 
     @pytest.mark.parametrize("getter", [
         'name-traj', 'number-traj', 'tag-final-traj', 'tag-initial-traj',
-        'name-sset', 'number-sset', 'tag-final-sset', 'tag-initial-sset'
+        'name-sset', 'number-sset', 'tag-final-sset', 'tag-initial-sset',
+        'name-samp', 'number-samp',
     ])
     def test_get(self, getter):
         filename = self.create_file(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],
+            'samp': [self.sample_set[0].trajectory],
         }[get_type]
         get_arg = {
             'name': 'traj',
@@ -277,7 +284,13 @@ 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]
+        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')
@@ -297,6 +310,18 @@ def test_cannot_guess(self):
         with pytest.raises(RuntimeError):
             self.PARAMETER.get(storage, None)
 
+    def test_get_bad_name(self):
+        filename = self._filename("bad_tag")
+        storage = paths.Storage(filename, 'w')
+        storage.save(self.traj)
+        storage.save(self.other_traj)
+        storage.tags['bad_tag'] = "foo"
+        storage.close()
+
+        storage = paths.Storage(filename, 'r')
+        with pytest.raises(RuntimeError, match="initial conditions type"):
+            self.PARAMETER.get(storage, "bad_tag")
+
 
 class TestINIT_SNAP(ParamInstanceTest):
     PARAMETER = INIT_SNAP