Merge pull request #1 from dwhswenson/visit-all

Author: dwhswenson

paths_cli/commands/visit_all.py

@@ -0,0 +1,46 @@
+import click
+
+from paths_cli.parameters import (INPUT_FILE, OUTPUT_FILE, ENGINE, STATES,
+                                  INIT_SNAP)
+
+@click.command(
+    "visit-all",
+    short_help="Run MD to generate initial trajectories",
+)
+@INPUT_FILE.clicked(required=True)
+@OUTPUT_FILE.clicked(required=True)
+@STATES.clicked(required=True)
+@ENGINE.clicked(required=False)
+@INIT_SNAP.clicked(required=False)
+def visit_all(input_file, output_file, state, engine, init_frame):
+    """Run until initial trajectory for TPS/MSTPS/MSTIS achieved.
+
+    This runs until all given states have been visited. That creates a long
+    trajectory, subtrajectories of which will work for the initial
+    trajectories in TPS, MSTPS, or MSTIS. Typically, you'll use a different
+    engine from the TPS production engine (often high temperature).
+    """
+    storage = INPUT_FILE.get(input_file)
+    visit_all_main(
+        output_storage=OUTPUT_FILE.get(output_file),
+        states=STATES.get(storage, state),
+        engine=ENGINE.get(storage, engine),
+        initial_frame=INIT_SNAP.get(storage, init_frame)
+    )
+
+
+def visit_all_main(output_storage, states, engine, initial_frame):
+    import openpathsampling as paths
+    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
+
+    return trajectory, None  # no simulation object to return here
+
+
+CLI = visit_all
+SECTION = "Simulation"
+REQUIRES_OPS = (1, 0)

paths_cli/tests/commands/__init__.py

@@ -0,0 +1,76 @@
+import os
+
+import pytest
+from unittest.mock import patch, MagicMock
+import tempfile
+from click.testing import CliRunner
+
+from paths_cli.commands.visit_all import *
+
+import openpathsampling as paths
+
+# patch with this for testing
+def print_test(output_storage, states, engine, initial_frame):
+    print(isinstance(output_storage, paths.Storage))
+    print(sorted([s.__uuid__ for s in states]))
+    print(engine.__uuid__)
+    print(initial_frame.__uuid__)
+
+
+@pytest.fixture()
+def tps_fixture(flat_engine, tps_network_and_traj):
+    network, traj = tps_network_and_traj
+    scheme = paths.OneWayShootingMoveScheme(network=network,
+                                            selector=paths.UniformSelector(),
+                                            engine=flat_engine)
+    init_conds = scheme.initial_conditions_from_trajectories(traj)
+    return (scheme, network, flat_engine, init_conds)
+
+@pytest.fixture()
+def visit_all_fixture(tps_fixture):
+    scheme, network, engine, init_conds = tps_fixture
+    states = sorted(network.all_states, key=lambda x: x.__uuid__)
+    init_frame = init_conds[0].trajectory[0]
+    return states, engine, init_frame
+
+
+@patch('paths_cli.commands.visit_all.visit_all_main', print_test)
+def test_visit_all(visit_all_fixture):
+    # this is an integration test; testing integration click & parameters
+    states, engine, init_frame = visit_all_fixture
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        storage = paths.Storage("setup.nc", 'w')
+        for obj in visit_all_fixture:
+            storage.save(obj)
+        storage.tags['initial_snapshot'] = init_frame
+        storage.close()
+
+        results = runner.invoke(
+            visit_all,
+            ["setup.nc", '-o', 'foo.nc', '-s', 'A', '-s', 'B',
+             '-e', 'flat', '-f', 'initial_snapshot']
+        )
+
+    expected_output = ("True\n[" + str(states[0].__uuid__) + ", "
+                       + str(states[1].__uuid__) + "]\n")
+    expected_output += "\n".join(str(obj.__uuid__)
+                                 for obj in [engine, init_frame]) + "\n"
+    assert results.exit_code == 0
+    assert results.output == expected_output
+
+def test_visit_all_main(visit_all_fixture):
+    # just a smoke test here
+    tempdir = tempfile.mkdtemp()
+    try:
+        store_name = os.path.join(tempdir, "visit_all.nc")
+        storage = paths.Storage(store_name, mode='w')
+        states, engine, init_frame = visit_all_fixture
+        traj, foo = visit_all_main(storage, states, engine, init_frame)
+        assert isinstance(traj, paths.Trajectory)
+        assert foo is None
+        assert len(storage.trajectories) == 1
+        storage.close()
+    finally:
+        os.remove(store_name)
+        os.rmdir(tempdir)

paths_cli/tests/commands/test_visit_all.py

@@ -0,0 +1,25 @@
+from openpathsampling.tests.test_helpers import make_1d_traj
+from openpathsampling.engines import toy as toys
+import openpathsampling as paths
+import pytest
+
+@pytest.fixture
+def flat_engine():
+    pes = toys.LinearSlope([0, 0, 0], 0)
+    topology = toys.Topology(n_spatial=3, masses=[1.0, 1.0, 1.0], pes=pes)
+    integ = toys.LeapfrogVerletIntegrator(dt=0.1)
+    options = {'integ': integ,
+               'n_frames_max': 1000,
+               'n_steps_per_frame': 1}
+    engine = toys.Engine(options=options, topology=topology).named("flat")
+    return engine
+
+@pytest.fixture
+def tps_network_and_traj():
+    cv = paths.CoordinateFunctionCV("x", lambda s: s.xyz[0][0])
+    state_A = paths.CVDefinedVolume(cv, float("-inf"), 0).named("A")
+    state_B = paths.CVDefinedVolume(cv, 1, float("inf")).named("B")
+    network = paths.TPSNetwork(state_A, state_B)
+    init_traj = make_1d_traj([-0.1, 0.1, 0.3, 0.5, 0.7, 0.9, 1.1],
+                             [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0])
+    return (network, init_traj)