Merge pull request #88 from dwhswenson/cmd-bootstrap-init

dwhswenson · web-flow · commit 02b33746dbae · 2024-08-30T09:19:49.000+02:00
New command: `bootstrap-init`
diff --git a/paths_cli/commands/bootstrap_init.py b/paths_cli/commands/bootstrap_init.py
@@ -0,0 +1,85 @@
+import click
+from paths_cli import OPSCommandPlugin
+from paths_cli.parameters import (
+    INIT_SNAP, SCHEME, ENGINE, OUTPUT_FILE, INPUT_FILE
+)
+from paths_cli.param_core import OPSStorageLoadSingle, Option
+
+INIT_STATE = OPSStorageLoadSingle(
+    param=Option("--initial-state", help="initial state"),
+    store='volumes'
+)
+FINAL_STATE = OPSStorageLoadSingle(
+    param=Option("--final-state", help="final state"),
+    store='volumes'
+)
+
+@click.command(
+    "bootstrap-init",
+    short_help="TIS interface set initial trajectories from a snapshot",
+)
+@INIT_STATE.clicked(required=True)
+@FINAL_STATE.clicked(required=True)
+@INIT_SNAP.clicked()
+@SCHEME.clicked()
+@ENGINE.clicked()
+@OUTPUT_FILE.clicked()
+@INPUT_FILE.clicked()
+def bootstrap_init(initial_state, final_state, scheme, engine, init_frame,
+                   output_file, input_file):
+    """Use ``FullBootstrapping`` to create initial conditions for TIS.
+
+    This approach starts from a snapshot, runs MD to generate an initial
+    path for the innermost ensemble, and then performs one-way shooting
+    moves within each ensemble until the next ensemble as reached. This
+    continues until all ensembles have valid trajectories.
+
+    Note that intermediate sampling in this is not saved to disk.
+    """
+    storage = INPUT_FILE.get(input_file)
+    scheme = SCHEME.get(storage, scheme)
+    network = scheme.network
+    engine = ENGINE.get(storage, engine)
+    init_state = INIT_STATE.get(storage, initial_state)
+    final_state = FINAL_STATE.get(storage, final_state)
+    transition = network.transitions[(init_state, final_state)]
+    bootstrap_init_main(
+        init_frame=INIT_SNAP.get(storage, init_frame),
+        network=network,
+        engine=engine,
+        transition=transition,
+        output_storage=OUTPUT_FILE.get(output_file)
+    )
+
+
+def bootstrap_init_main(init_frame, network, engine, transition,
+                        output_storage):
+    import openpathsampling as paths
+    all_states = set(network.initial_states) | set(network.final_states)
+    allowed_states = {transition.stateA, transition.stateB}
+    forbidden_states = list(all_states - allowed_states)
+    try:
+        extra_ensembles = network.ms_outers
+    except KeyError:
+        extra_ensembles = None
+
+    bootstrapper = paths.FullBootstrapping(
+        transition=transition,
+        snapshot=init_frame,
+        engine=engine,
+        forbidden_states=forbidden_states,
+        extra_ensembles=extra_ensembles,
+    )
+    init_conds = bootstrapper.run()
+    if output_storage:
+        output_storage.tags['final_conditions'] = init_conds
+
+    return init_conds, bootstrapper
+
+
+PLUGIN = OPSCommandPlugin(
+    command=bootstrap_init,
+    section="Simulation",
+    requires_ops=(1, 0),
+    requires_cli=(0, 4),
+)
diff --git a/paths_cli/tests/commands/test_bootstrap_init.py b/paths_cli/tests/commands/test_bootstrap_init.py
@@ -0,0 +1,132 @@
+import pytest
+from click.testing import CliRunner
+from unittest.mock import patch
+import numpy as np
+
+from paths_cli.commands.bootstrap_init import *
+import openpathsampling as paths
+from openpathsampling.engines import toy
+from openpathsampling.tests.test_helpers import make_1d_traj
+
+@pytest.fixture
+def toy_2_state_engine():
+    pes = (
+        toy.OuterWalls([1.0, 1.0], [0.0, 0.0]) +
+        toy.Gaussian(-1.0, [12.0, 12.0], [-0.5, 0.0]) +
+        toy.Gaussian(-1.0, [12.0, 12.0], [0.5, 0.0])
+    )
+    topology=toy.Topology(
+        n_spatial = 2,
+        masses =[1.0, 1.0],
+        pes = pes
+    )
+    integ = toy.LangevinBAOABIntegrator(dt=0.02, temperature=0.1, gamma=2.5)
+    options = {
+        'integ': integ,
+        'n_frames_max': 5000,
+        'n_steps_per_frame': 1
+    }
+
+    engine = toy.Engine(
+        options=options,
+        topology=topology
+    )
+    return engine
+
+@pytest.fixture
+def toy_2_state_cv():
+    return paths.FunctionCV("x", lambda s: s.xyz[0][0])
+
+@pytest.fixture
+def toy_2_state_volumes(toy_2_state_cv):
+    state_A = paths.CVDefinedVolume(
+        toy_2_state_cv,
+        float("-inf"),
+        -0.3,
+    ).named("A")
+    state_B = paths.CVDefinedVolume(
+        toy_2_state_cv,
+        0.3,
+        float("inf"),
+    ).named("B")
+    return state_A, state_B
+
+@pytest.fixture
+def toy_2_state_tis(toy_2_state_cv, toy_2_state_volumes):
+    state_A, state_B = toy_2_state_volumes
+    interfaces = paths.VolumeInterfaceSet(
+        toy_2_state_cv,
+        float("-inf"),
+        [-0.3, -0.2, -0.1],
+    )
+    tis = paths.MISTISNetwork(
+        [(state_A, interfaces, state_B)],
+    )
+    return tis
+
+
+def print_test(init_frame, network, engine, transition, output_storage):
+    print(init_frame.__uuid__)
+    print(network.__uuid__)
+    print(engine.__uuid__)
+    # apparently transition UUID isn't preserved, but these are?
+    print(transition.stateA.__uuid__)
+    print(transition.stateB.__uuid__)
+    print([e.__uuid__ for e in transition.ensembles])
+    print(isinstance(output_storage, paths.Storage))
+
+@patch('paths_cli.commands.bootstrap_init.bootstrap_init_main', print_test)
+def test_bootstrap_init(tis_fixture):
+    scheme, network, engine, init_conds = tis_fixture
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        storage = paths.Storage("setup.nc", 'w')
+        storage.save(init_conds)
+        for obj in tis_fixture:
+            storage.save(obj)
+
+        storage.tags["init_snap"] = init_conds[0][0]
+        storage.close()
+
+        results = runner.invoke(bootstrap_init, [
+            'setup.nc',
+            '-o', 'foo.nc',
+            '--initial-state', "A",
+            '--final-state', "B",
+            '--init-frame', 'init_snap',
+        ])
+
+        transitions = list(network.transitions.values())
+        assert len(transitions) == 1
+        transition = transitions[0]
+        stateA = transition.stateA
+        stateB = transition.stateB
+        ensembles = transition.ensembles
+
+        expected_output = (
+            f"{init_conds[0][0].__uuid__}\n{network.__uuid__}\n"
+            f"{engine.__uuid__}\n"
+            f"{stateA.__uuid__}\n{stateB.__uuid__}\n"
+            f"{[e.__uuid__ for e in ensembles]}\n"
+            "True\n"
+        )
+        assert results.exit_code == 0
+        assert results.output == expected_output
+
+
+def test_bootstrap_init_main(toy_2_state_tis, toy_2_state_engine, tmp_path):
+    network = toy_2_state_tis
+    engine = toy_2_state_engine
+    scheme = paths.DefaultScheme(network, engine)
+    init_frame = toy.Snapshot(
+        coordinates=np.array([[-0.5, -0.5]]),
+        velocities=np.array([[0.0,0.0]]),
+        engine=engine
+    )
+    assert len(network.transitions) == 1
+    transition = list(network.transitions.values())[0]
+    output_storage = paths.Storage(tmp_path / "output.nc", mode='w')
+    init_conds, bootstrapper = bootstrap_init_main(init_frame, network,
+                                                   engine, transition,
+                                                   output_storage)
+    init_conds.sanity_check()
diff --git a/paths_cli/tests/conftest.py b/paths_cli/tests/conftest.py
@@ -58,3 +58,11 @@ def tis_network(cv_and_states):
                                           [0.0, 0.1, 0.2])
     network = paths.MISTISNetwork([(state_A, interfaces, state_B)])
     return network
+
+@pytest.fixture
+def tis_fixture(flat_engine, tis_network, transition_traj):
+    paths.InterfaceSet._reset()
+    scheme = paths.DefaultScheme(network=tis_network,
+                                 engine=flat_engine)
+    init_conds = scheme.initial_conditions_from_trajectories(transition_traj)
+    return (scheme, tis_network, flat_engine, init_conds)
