Merge pull request #6 from dwhswenson/join_one_way

[WIP] Add ability to use 1-way segments as input to pseudo-simulator

ops_piggybacker/mover_stubs.py

@@ -4,7 +4,30 @@ class NoEngine(paths.engines.DynamicsEngine):
     pass
 
 class ShootingStub(paths.pathmover.PathMover):
-    def __init__(self, ensemble, selector=None, engine=None):
+    """Stub to mimic a shooting move.
+    
+    Parameters
+    ----------
+    ensemble : paths.Ensemble
+        the ensemble for the shooting mover
+    selector : paths.ShootingPointSelector or None
+        the selector for the shooting point. Default None creates a
+        UniformSelector. Currently, only UniformSelector is supported.
+    engine : paths.engines.DynamicsEngine
+        the engine to report as the source of the dynamics
+    pre_joined : bool
+        whether the input trial trajectories are pre-joined into complete
+        trajectories, or take partial one-way segments which should by
+        dynamically joined. Currently defaults to pre_joined=True (likely to
+        change soon, though)
+
+    Attributes
+    ----------
+    mimic : paths.OneWayShootingMover
+        the mover that this stub mimics
+
+    """
+    def __init__(self, ensemble, selector=None, engine=None, pre_joined=True):
         super(ShootingStub, self).__init__()
         if engine is None:
             engine = NoEngine()
@@ -13,10 +36,47 @@ def __init__(self, ensemble, selector=None, engine=None):
         self.engine = engine
         self.selector = selector
         self.ensemble = ensemble
+        self.pre_joined = pre_joined
         self.mimic = paths.OneWayShootingMover(ensemble, selector, engine)
 
+    def join_one_way(self, input_trajectory, partial_trial,
+                     shooting_point, direction):
+        """Create a one-way trial trajectory
+
+        Parameters
+        ----------
+        input_trajectory : paths.Trajectory
+            the previous complete trajectory
+        partial_trial : paths.Trajectory
+            The partial (one-way) trial trajectory. Must *not* include the
+            shooting point.
+        shooting_point : paths.Snapshot
+            the snapshot for the shooting point -- must be a member of the
+            input trajectory
+        direction : +1 or -1
+            if positive, treat as forward shooting; if negative, treat as
+            backward shooting
+
+        Returns
+        -------
+        paths.Trajectory
+            the complete trial trajectory
+        """
+        shooting_idx = input_trajectory.index(shooting_point)
+        if direction > 0:
+            joined_trajectory = (input_trajectory[:shooting_idx+1] +
+                                 partial_trial)
+        elif direction < 0:
+            joined_trajectory = (partial_trial +
+                                 input_trajectory[shooting_idx:])
+        else: # pragma: no cover
+            raise RuntimeError("Bad direction for shooting: " +
+                               str(direction))
+        return joined_trajectory
+
 
-    def move(self, input_sample, trial_trajectory, shooting_point, accepted):
+    def move(self, input_sample, trial_trajectory, shooting_point, accepted,
+             direction=None):
         """Fake a move.
 
         Parameters
@@ -29,12 +89,23 @@ def move(self, input_sample, trial_trajectory, shooting_point, accepted):
             the shooting point snapshot for this trial
         accepted: bool
             whether the trial was accepted
+        direction: +1 or -1
+            direction of the shooting move (positive is forward, negative is
+            backward). Only relevant if self.pre_joined
+            is False.
         """
         initial_trajectory = input_sample.trajectory
         replica = input_sample.replica
         ensemble = input_sample.ensemble
 
-        # determine the direction
+        if not self.pre_joined:
+            trial_trajectory = self.join_one_way(initial_trajectory,
+                                                 trial_trajectory,
+                                                 shooting_point,
+                                                 direction)
+
+        # determine the direction based on trial trajectory (maybe check
+        # with given direction if given?)
         shared = trial_trajectory.shared_subtrajectory(initial_trajectory)
         if shared[0] == trial_trajectory[0]:
             choice = 0  # forward submover

ops_piggybacker/simulation_stubs.py

@@ -1,6 +1,19 @@
 import openpathsampling as paths
 
 class ShootingPseudoSimulator(paths.PathSimulator):
+    """Pseudo-simulator for shooting-only mimics.
+
+    Parameters
+    ----------
+    storage : openpathsampling.netcdfplus.Storage
+        file to store OPS-ready analysis
+    initial_conditions : openpathsampling.SampleSet
+        sample set giving the OPS version of the initial conditions
+    mover : ShootingStub
+        stub to mimic the shooting mover
+    network : openpathsampling.TransitionNetwork
+        transition network with information about this system
+    """
     def __init__(self, storage, initial_conditions, mover, network):
         super(ShootingPseudoSimulator, self).__init__(storage)
         self.scheme = paths.LockedMoveScheme(mover, network)
@@ -19,7 +32,9 @@ def run(self, step_info_list):
         Parameters
         ----------
         step_info_list : list of tuple
-            (replica, trial_trajectory, shooting_point_index, accepted)
+            (replica, trial_trajectory, shooting_point_index, accepted) or
+            (replica, one_way_trial_segment, shooting_point_index, accepted,
+            direction)
         """
         mcstep = None
 
@@ -30,17 +45,26 @@ def run(self, step_info_list):
 
         for step_info in step_info_list:
             self.step += 1
+            if len(step_info) == 4 and not self.mover.pre_joined: # pragma: no-cover
+                raise RuntimeError(
+                    "Shooting trial trajectories not pre-joined: " + 
+                    "step_info must be (replica, trial_segment, " + 
+                    "shooting_pt_idx, accepted, direction)")
 
             replica = step_info[0]
             trial_trajectory = step_info[1]
             shooting_point_index = step_info[2]
             accepted = step_info[3]
+            direction = None
+            if len(step_info) == 5:
+                direction = step_info[4]
 
             input_sample = self.globalstate[replica]
             shooting_point = input_sample.trajectory[shooting_point_index]
 
+
             subchange = self.mover.move(input_sample, trial_trajectory,
-                                        shooting_point, accepted)
+                                        shooting_point, accepted, direction)
 
             change = paths.PathSimulatorPathMoveChange(
                 subchange=subchange,

ops_piggybacker/tests/common_test_data.py

@@ -28,12 +28,13 @@ def shooting_move_info():
 
     # for traj in [t0, out1, out2, out3, out4]:
         # print [s.xyz[0][0] for s in traj]
-
+
+    # replica, full trial, shooting idx, one-way trial, direction
     moves = [
-        (0, out1, 4, True),
-        (0, out2, 2, True),
-        (0, out3, 5, False),
-        (0, out4, 4, True)
+        (0, out1, 4, True, t1, -1),
+        (0, out2, 2, True, t2, +1),
+        (0, out3, 5, False, t3, -1),
+        (0, out4, 4, True, t4, -1)
     ]
     initial_sample = paths.Sample(replica=0,
                                   trajectory=t0,

ops_piggybacker/tests/test_mover_stubs.py

@@ -14,6 +14,40 @@ class testShootingStub(object):
     def setup(self):
         # this instantiates an object, and ensures that instantiation works
         self.stub = oink.ShootingStub(common.tps_ensemble)
+        self.no_prejoin = oink.ShootingStub(common.tps_ensemble,
+                                            pre_joined=False)
+
+    def test_join_one_way(self):
+        moves = common.tps_shooting_moves
+        initial = common.initial_tps_sample.trajectory
+
+        out_1 = moves[0][1]
+        sp_1 = initial[moves[0][2]]
+        trial_1 = moves[0][4]
+        dir_1 = moves[0][5]
+        joined_1 = self.no_prejoin.join_one_way(initial, trial_1, sp_1, dir_1)
+        assert_equal(joined_1, out_1)
+
+        out_2 = moves[1][1]
+        sp_2 = joined_1[moves[1][2]]
+        trial_2 = moves[1][4]
+        dir_2 = moves[1][5]
+        joined_2 = self.no_prejoin.join_one_way(joined_1, trial_2, sp_2, dir_2)
+        assert_equal(joined_2, out_2)
+
+        out_3 = moves[2][1]
+        sp_3 = joined_2[moves[2][2]]
+        trial_3 = moves[2][4]
+        dir_3 = moves[2][5]
+        joined_3 = self.no_prejoin.join_one_way(joined_2, trial_3, sp_3, dir_3)
+        assert_equal(joined_3, out_3)
+
+        out_4 = moves[3][1]
+        sp_4 = joined_2[moves[3][2]]
+        trial_4 = moves[3][4]
+        dir_4 = moves[3][5]
+        joined_4 = self.no_prejoin.join_one_way(joined_2, trial_4, sp_4, dir_4)
+        assert_equal(joined_4, out_4)
 
     def test_backward_move(self):
         initial_sample = common.initial_tps_sample
@@ -37,6 +71,29 @@ def test_backward_move(self):
         assert_equal(change.accepted, True)
         assert_equal(type(change.canonical.mover), paths.BackwardShootMover)
 
+    def test_backward_move_not_pre_joined(self):
+        initial_sample = common.initial_tps_sample
+        move = common.tps_shooting_moves[0]
+        replica = move[0]
+        trial_trajectory = move[4]
+        shooting_point = initial_sample.trajectory[move[2]]
+        accepted = move[3]
+        direction = move[5]
+        change = self.no_prejoin.move(initial_sample, trial_trajectory,
+                                      shooting_point, accepted, direction)
+        # assertions about the shape of the decision history
+        assert_equal(change.mover, self.no_prejoin.mimic)
+        assert_equal(change.subchange, change.canonical)
+        assert_equal(change.subchange.subchange, None)
+        assert_equal(len(change.trials), 1)
+        # assertions true for any OneWayShooting
+        details0 = change.canonical.trials[0].details
+        assert_equal(details0.shooting_snapshot, shooting_point)
+        assert_equal(details0.initial_trajectory.index(shooting_point), move[2])
+        # assertions specific to this test
+        assert_equal(change.accepted, True)
+        assert_equal(type(change.canonical.mover), paths.BackwardShootMover)
+
     def test_forward_move(self):
         init_traj = common.tps_shooting_moves[0][1]
         initial_sample = paths.Sample(replica=0,
@@ -62,6 +119,32 @@ def test_forward_move(self):
         assert_equal(change.accepted, True)
         assert_equal(type(change.canonical.mover), paths.ForwardShootMover)
 
+    def test_forward_move_not_pre_joined(self):
+        init_traj = common.tps_shooting_moves[0][1]
+        initial_sample = paths.Sample(replica=0,
+                                      trajectory=init_traj,
+                                      ensemble=common.tps_ensemble)
+        move = common.tps_shooting_moves[1]
+        replica = move[0]
+        trial_trajectory = move[4]
+        shooting_point = initial_sample.trajectory[move[2]]
+        accepted = move[3]
+        direction = move[5]
+        change = self.no_prejoin.move(initial_sample, trial_trajectory,
+                                      shooting_point, accepted, direction)
+        # assertions about the shape of the decision history
+        assert_equal(change.mover, self.no_prejoin.mimic)
+        assert_equal(change.subchange, change.canonical)
+        assert_equal(change.subchange.subchange, None)
+        assert_equal(len(change.trials), 1)
+        # assertions true for any OneWayShooting
+        details0 = change.canonical.trials[0].details
+        assert_equal(details0.shooting_snapshot, shooting_point)
+        assert_equal(details0.initial_trajectory.index(shooting_point), move[2])
+        # assertions specific to this test
+        assert_equal(change.accepted, True)
+        assert_equal(type(change.canonical.mover), paths.ForwardShootMover)
+
     def test_rejected_move(self):
         initial_sample = common.initial_tps_sample
         move = common.tps_shooting_moves[0]
@@ -83,3 +166,26 @@ def test_rejected_move(self):
         # assertions specific to this test
         assert_equal(change.accepted, False)
         assert_equal(type(change.canonical.mover), paths.BackwardShootMover)
+
+    def test_rejected_move_not_pre_joined(self):
+        initial_sample = common.initial_tps_sample
+        move = common.tps_shooting_moves[0]
+        replica = move[0]
+        trial_trajectory = move[4]
+        shooting_point = initial_sample.trajectory[move[2]]
+        accepted = False
+        direction = move[5]
+        change = self.no_prejoin.move(initial_sample, trial_trajectory,
+                                      shooting_point, accepted, direction)
+        # assertions about the shape of the decision history
+        assert_equal(change.mover, self.no_prejoin.mimic)
+        assert_equal(change.subchange, change.canonical)
+        assert_equal(change.subchange.subchange, None)
+        assert_equal(len(change.trials), 1)
+        # assertions true for any OneWayShooting
+        details0 = change.canonical.trials[0].details
+        assert_equal(details0.shooting_snapshot, shooting_point)
+        assert_equal(details0.initial_trajectory.index(shooting_point), move[2])
+        # assertions specific to this test
+        assert_equal(change.accepted, False)
+        assert_equal(type(change.canonical.mover), paths.BackwardShootMover)

ops_piggybacker/tests/test_simulation_stubs.py

@@ -32,14 +32,59 @@ def setup(self):
             mover=shoot,
             network=network
         )
+
+        nojoin = oink.ShootingStub(tps_ensemble, pre_joined=False)
+        self.nojoin_pseudosim = oink.ShootingPseudoSimulator(
+            storage=self.storage,
+            initial_conditions=paths.SampleSet([initial_sample]),
+            mover=nojoin,
+            network=network
+        )
+
 
     def teardown(self):
         if os.path.isfile(data_filename(self.fname)):
             os.remove(data_filename(self.fname))
 
+    def test_noprejoin_run_and_analyze(self):
+        moves = [(move[0], move[4], move[2], move[3], move[5]) 
+                 for move in common.tps_shooting_moves]
+        self.nojoin_pseudosim.run(moves)
+        self.storage.close()
+        # open the file for analysis, check that its content is reasonable
+        analysis = paths.AnalysisStorage(data_filename(self.fname))
+        assert_equal(len(analysis.steps), 5) # initial + 4 steps
+        scheme = analysis.schemes[0]
+        assert_equal(scheme.movers.keys(), ['shooting'])
+        assert_equal(len(scheme.movers['shooting']), 1)
+        mover = scheme.movers['shooting'][0]
+
+        # use several OPS tools to analyze this file
+        ## scheme.move_summary
+        devnull = open(os.devnull, 'w')
+        scheme.move_summary(analysis, output=devnull) 
+        mover_keys = [k for k in scheme._mover_acceptance.keys()
+                      if k[0] == mover]
+        assert_equal(len(mover_keys), 1)
+        assert_equal(scheme._mover_acceptance[mover_keys[0]], [3,4])
+
+        ## move history tree
+        import openpathsampling.visualize as ops_vis
+        history = ops_vis.ReplicaHistoryTree(
+            storage=analysis,
+            steps=analysis.steps[0:],
+            replica=0
+        )
+        assert_equal(len(history.decorrelated_trajectories), 2)
+
+        ## path length histogram
+        path_lengths = [len(step.active[0].trajectory) 
+                        for step in analysis.steps]
+        assert_equal(path_lengths, [11, 9, 7, 7, 7])
 
     def test_run_and_analyze(self):
-        self.pseudosim.run(common.tps_shooting_moves)
+        moves = [tuple(move[0:4]) for move in common.tps_shooting_moves]
+        self.pseudosim.run(moves)
         self.storage.close()
         # open the file for analysis, check that its content is reasonable
         analysis = paths.AnalysisStorage(data_filename(self.fname))

ops_piggybacker/tests/tools.py

@@ -13,3 +13,8 @@
 def data_filename(fname, subdir='test_data'):
     return resource_filename('ops_piggybacker',
                              os.path.join('tests', subdir, fname))
+
+import logging
+logging.getLogger('openpathsampling.netcdfplus').setLevel(logging.CRITICAL)
+logging.getLogger('openpathsampling.ensemble').setLevel(logging.CRITICAL)
+logging.getLogger('openpathsampling.initialization').setLevel(logging.CRITICAL)