draft of bootstrapping command

Author: dwhswenson

paths_cli/commands/bootstrap_init.py

@@ -1,35 +1,82 @@
+import click
+from paths_cli import OPSCommandPlugin
 from paths_cli.parameters import (
-    MULTI_INIT_SNAP, MULTI_NETWORK, OUTPUT_FILE, INPUT_FILE
+    INIT_SNAP, MULTI_NETWORK, OUTPUT_FILE, INPUT_FILE, ENGINE
 )
+from paths_cli.param_core import OPSStorageLoadSingle, Option
 
-@MULTI_INIT_SNAP.clicked()
+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()
 @MULTI_NETWORK.clicked()
+@ENGINE.clicked()
 @OUTPUT_FILE.clicked()
 @INPUT_FILE.clicked()
-def bootstrap_init(init_frame, network, output_file, input_file):
-    ...
-
-def _update_init_frames(transitions, init_frames):
-    transition_to_init_frame = {}
-    for transition in transitions:
-        # raise error if not TIS transition
-        transition_to_init_frame[transition] = None
-        for snap in init_frames:
-            if init_frame.initial_state(snap):
-                transition_to_init_frame[transition] = snap
-                break
-    return transition_to_init_frame
-
-
-def bootstrap_init_main(init_frames, networks, output_file):
-    all_transitions = {trans for trans in network.sampling_transitions
-                       for network in networks}
-    trans_to_init_frame = _update_init_frames(all_transitions, init_frames)
-    missing = {trans for trans, frame in trans_to_init_frame.items()
-               if frame is None}
-    # TODO warn about missing initial frames
-    found = set(trans_to_init_frame) - missing
-    completed = set()
-    for transition in found:
-        ...
+def bootstrap_init(init_state, final_state, network, engine, init_snap,
+                   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)
+    network = MULTI_NETWORK.get(storage, network)
+    init_state = INIT_STATE.get(storage, init_state)
+    final_state = FINAL_STATE.get(storage, final_stage)
+    transition = network.transitions[(init_state, final_state)]
+    bootstrap_init_main(
+        init_frame=INIT_SNAP.get(storage, init_snap),
+        network=network,
+        engine=engine,
+        transition=transition,
+        output_storage=OUTPUT_FILE.get(output_file)
+    )
 
+
+def bootstrap_init_main(init_frame, network, engine, transition,
+                        output_storage):
+    all_states = set(network.initial_states) | set(network.final_states)
+    allowed_states = {transition.initial_state, transition.final_state}
+    forbidden_states = list(all_states - allowed_states)
+    try:
+        extra_ensembles = network.ms_outers
+    except KeyError:
+        extra_ensembles = None
+
+    bootstrapper = paths.FullBootstrapping(
+        transition=transition,
+        snaphot=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),
+)