dwhswenson · Sep 18, 2018 · Sep 8, 2018 · Sep 8, 2018 · Sep 8, 2018 · Sep 8, 2018
diff --git a/openpathsampling/__init__.py b/openpathsampling/__init__.py
@@ -92,6 +92,10 @@
 
 from .high_level.ms_outer_interface import MSOuterTISInterface
 
+from .high_level.part_in_b_tps import (
+    PartInBFixedLengthTPSNetwork, PartInBFixedLengthTPSTransition
+)
+
 from .step_visualizer_2D import StepVisualizer2D
 
 from .movechange import (

diff --git a/openpathsampling/high_level/part_in_b_tps.py b/openpathsampling/high_level/part_in_b_tps.py
@@ -0,0 +1,56 @@
+from openpathsampling.high_level.network import FixedLengthTPSNetwork
+from openpathsampling.high_level.transition import FixedLengthTPSTransition
+import openpathsampling as paths
+
+class PartInBFixedLengthTPSTransition(FixedLengthTPSTransition):
+    """Fixed length TPS transition accepting any frame in the final state.
+
+    Transition that builds an ensemble used to facilitate the rate
+    calculation in fixed-length TPS. [1]_ Details in
+    :class:`.PartInBFixedLengthTPSNetwork`.
+
+    See also
+    --------
+    PartInBFixedLengthTPSNetwork
+
+    References
+    ----------
+
+    .. [1] C. Dellago, P.G. Bolhuis, and D. Chandler. J. Chem. Phys. 110,
+           6617 (1999). http://dx.doi.org/10.1063/1.478569
+    """
+    def _tps_ensemble(self, stateA, stateB):
+        return paths.SequentialEnsemble([
+            paths.LengthEnsemble(1) & paths.AllInXEnsemble(stateA),
+            paths.LengthEnsemble(self.length - 1) \
+                & paths.PartInXEnsemble(stateB)
+        ])
+
+class PartInBFixedLengthTPSNetwork(FixedLengthTPSNetwork):
+    """Network for fixed-length TPS accepting any frame in the final state
+
+    This network samples a single path ensemble where the paths must begin
+    in an initial state, run for a fixed total number of frames, and must
+    have at least one frame in a final state. This  was used to assist in
+    the flux part of the TPS rate calculation. [1]_ This version is
+    generalized to multiple states.
+
+    Parameters
+    ----------
+    intial_states : (list of) :class:`.Volume`
+        acceptable initial states
+    final_states : (list of) :class:`.Volume`
+        acceptable final states
+    length : int
+        length of paths in the path ensemble, in frames
+    allow_self_transitions : bool
+        whether self-transitions (A->A) are allowed; default is False. For
+        this network, A->B->A transitions are *always* allowed.
+
+    References
+    ----------
+
+    .. [1] C. Dellago, P.G. Bolhuis, and D. Chandler. J. Chem. Phys. 110,
+           6617 (1999). http://dx.doi.org/10.1063/1.478569
+    """
+    TransitionType = PartInBFixedLengthTPSTransition
diff --git a/openpathsampling/tests/test_network.py b/openpathsampling/tests/test_network.py
@@ -361,6 +361,9 @@ def test_storage(self):
 
 
 class TestTPSNetwork(object):
+    NetworkType = TPSNetwork
+    std_kwargs = {}
+
     def setup(self):
         from .test_helpers import CallIdentity
         xval = paths.FunctionCV("xval", lambda snap: snap.xyz[0][0])
@@ -376,40 +379,49 @@ def setup(self):
         self.traj['BC'] = make_1d_traj([0.01, 0.16, 0.25, 0.53])
         self.traj['CC'] = make_1d_traj([0.51, 0.35, 0.36, 0.55])
         self.traj['CA'] = make_1d_traj([0.52, 0.22, -0.22, -0.52])
+        self.traj['AB0'] = make_1d_traj([-0.51, -0.25, 0.0, -0.25])
+        self.traj['ABA'] = make_1d_traj([-0.51, -0.25, 0.0, -0.51])
 
-    # define all the test networks as properties: we can do something
-    # similar then for the fixed path length, and just need to override
-    # these properties to get all the same tests
+    # define all the test networks as properties
     @property
     def network2a(self):
-        return TPSNetwork(initial_states=[self.stateA],
-                          final_states=[self.stateB])
+        return self.NetworkType(initial_states=[self.stateA],
+                                final_states=[self.stateB],
+                                **self.std_kwargs)
 
     @property
     def network2b(self):
-        return TPSNetwork(initial_states=self.stateA,
-                          final_states=self.stateB)
+        return self.NetworkType(initial_states=self.stateA,
+                                final_states=self.stateB,
+                                **self.std_kwargs)
 
     @property
     def network2c(self):
-        return TPSNetwork.from_state_pairs([(self.stateA, self.stateB)])
+        return self.NetworkType.from_state_pairs(
+            [(self.stateA, self.stateB)],
+            **self.std_kwargs
+        )
 
     @property
     def network3a(self):
-        return TPSNetwork(initial_states=self.states,
-                          final_states=self.states)
+        return self.NetworkType(initial_states=self.states,
+                                final_states=self.states,
+                                **self.std_kwargs)
 
     @property
     def network3b(self):
-        return TPSNetwork.from_states_all_to_all(self.states)
+        return self.NetworkType.from_states_all_to_all(self.states,
+                                                       **self.std_kwargs)
 
     @property
     def network3c(self):
-        return TPSNetwork.from_state_pairs([
-            (self.stateA, self.stateB), (self.stateA, self.stateC),
-            (self.stateB, self.stateA), (self.stateB, self.stateC),
-            (self.stateC, self.stateA), (self.stateC, self.stateB)
-        ])
+        return self.NetworkType.from_state_pairs(
+            [(self.stateA, self.stateB), (self.stateA, self.stateC),
+             (self.stateB, self.stateA), (self.stateB, self.stateC),
+             (self.stateC, self.stateA), (self.stateC, self.stateB)
+            ],
+            **self.std_kwargs
+        )
 
     def test_initialization_2state(self):
         network2a = self.network2a
@@ -423,7 +435,7 @@ def test_initialization_2state(self):
         assert_equal(len(network2c.transitions), 1)
         assert_equal(set(network2a.all_states), set(network2b.all_states))
         assert_equal(set(network2b.all_states), set(network2c.all_states))
-        assert_equal(set(network2a.all_states), 
+        assert_equal(set(network2a.all_states),
                      set([self.stateA, self.stateB]))
 
 
@@ -447,11 +459,13 @@ def test_storage(self):
         topol = peng.Topology(n_spatial=1, masses=[1.0], pes=None)
         engine = peng.Engine({}, topol)
         self.template = peng.Snapshot(coordinates=np.array([[0.0]]),
-                                       velocities=np.array([[0.0]]),
-                                       engine=engine)
+                                      velocities=np.array([[0.0]]),
+                                      engine=engine)
 
         states = [self.stateA, self.stateB, self.stateC]
-        network_a = TPSNetwork(initial_states=states, final_states=states)
+        network_a = self.NetworkType(initial_states=states,
+                                     final_states=states,
+                                     **self.std_kwargs)
         assert_equal(len(network_a.sampling_transitions), 1)
         assert_equal(len(network_a.transitions), 6)
         storage_w = paths.storage.Storage(fname, "w")
@@ -468,8 +482,8 @@ def test_storage(self):
             os.remove(fname)
 
     def test_allow_self_transitions_false(self):
-        network = TPSNetwork.from_states_all_to_all(
-            self.states, allow_self_transitions=False
+        network = self.NetworkType.from_states_all_to_all(
+            self.states, allow_self_transitions=False, **self.std_kwargs
         )
         assert_equal(len(network.sampling_ensembles), 1)
         ensemble = network.sampling_ensembles[0]
@@ -482,8 +496,8 @@ def test_allow_self_transitions_false(self):
         assert_equal(ensemble(self.traj['CC']), False)
 
     def test_allow_self_transitions_true(self):
-        network = TPSNetwork.from_states_all_to_all(
-            self.states, allow_self_transitions=True
+        network = self.NetworkType.from_states_all_to_all(
+            self.states, allow_self_transitions=True, **self.std_kwargs
         )
         assert_equal(len(network.sampling_ensembles), 1)
         ensemble = network.sampling_ensembles[0]
@@ -495,77 +509,40 @@ def test_allow_self_transitions_true(self):
         assert_equal(ensemble(self.traj['BB']), True)
         assert_equal(ensemble(self.traj['CC']), True)
 
-class TestFixedLengthTPSNetwork(TestTPSNetwork):
-    @property
-    def network2a(self):
-        return FixedLengthTPSNetwork(initial_states=[self.stateA],
-                                     final_states=[self.stateB],
-                                     length=10)
-
-    @property
-    def network2b(self):
-        return FixedLengthTPSNetwork(initial_states=self.stateA,
-                                     final_states=self.stateB,
-                                     length=10)
-
-    @property
-    def network2c(self):
-        return FixedLengthTPSNetwork.from_state_pairs(
-            [(self.stateA, self.stateB)],
-            length=10
+    def test_allow_self_transitions_false_ABX(self):
+        network = self.NetworkType.from_states_all_to_all(
+            self.states, allow_self_transitions=False, **self.std_kwargs
         )
+        assert_equal(len(network.sampling_ensembles), 1)
+        ensemble = network.sampling_ensembles[0]
+        assert_equal(ensemble(self.traj['AB0']), False)
+        assert_equal(ensemble(self.traj['ABA']), False)
 
-    @property
-    def network3a(self):
-        return FixedLengthTPSNetwork(initial_states=self.states,
-                                     final_states=self.states,
-                                     length=10)
-
-    @property
-    def network3b(self):
-        return FixedLengthTPSNetwork.from_states_all_to_all(self.states,
-                                                            length=10)
-
-    @property
-    def network3c(self):
-        return FixedLengthTPSNetwork.from_state_pairs(
-            [(self.stateA, self.stateB), (self.stateA, self.stateC),
-             (self.stateB, self.stateA), (self.stateB, self.stateC),
-             (self.stateC, self.stateA), (self.stateC, self.stateB)],
-            length=10
+    def test_allow_self_transitions_true_ABX(self):
+        # special case that differs between fixed and flexible
+        network = self.NetworkType.from_states_all_to_all(
+            self.states, allow_self_transitions=True, **self.std_kwargs
         )
+        assert_equal(len(network.sampling_ensembles), 1)
+        ensemble = network.sampling_ensembles[0]
+        assert_equal(ensemble(self.traj['AB0']), False)
+        assert_equal(ensemble(self.traj['ABA']), False)
+
+class TestFixedLengthTPSNetwork(TestTPSNetwork):
+    NetworkType = FixedLengthTPSNetwork
+    std_kwargs = {'length': 4}
 
     def test_lengths(self):
         for network in [self.network2a, self.network2b, self.network2c,
                         self.network3a, self.network3b, self.network3c]:
-            assert_equal(network.sampling_transitions[0].length, 10)
-            assert_equal(list(network.transitions.values())[0].length, 10)
+            assert_equal(network.sampling_transitions[0].length, 4)
+            assert_equal(list(network.transitions.values())[0].length, 4)
 
-    def test_allow_self_transitions_false(self):
-        network = FixedLengthTPSNetwork.from_states_all_to_all(
-            self.states, allow_self_transitions=False, length=4
+    def test_allow_self_transitions_true_ABX(self):
+        network = self.NetworkType.from_states_all_to_all(
+            self.states, allow_self_transitions=True, **self.std_kwargs
         )
         assert_equal(len(network.sampling_ensembles), 1)
         ensemble = network.sampling_ensembles[0]
-        assert_equal(ensemble(self.traj['AA']), False)
-        assert_equal(ensemble(self.traj['AB']), True)
-        assert_equal(ensemble(self.traj['BA']), True)
-        assert_equal(ensemble(self.traj['BC']), True)
-        assert_equal(ensemble(self.traj['CA']), True)
-        assert_equal(ensemble(self.traj['BB']), False)
-        assert_equal(ensemble(self.traj['CC']), False)
-
-    def test_allow_self_transitions_true(self):
-        network = FixedLengthTPSNetwork.from_states_all_to_all(
-            self.states, allow_self_transitions=True, length=4
-        )
-        assert_equal(len(network.sampling_ensembles), 1)
-        ensemble = network.sampling_ensembles[0]
-        assert_equal(ensemble(self.traj['AA']), True)
-        assert_equal(ensemble(self.traj['AB']), True)
-        assert_equal(ensemble(self.traj['BA']), True)
-        assert_equal(ensemble(self.traj['BC']), True)
-        assert_equal(ensemble(self.traj['CA']), True)
-        assert_equal(ensemble(self.traj['BB']), True)
-        assert_equal(ensemble(self.traj['CC']), True)
-
+        assert_equal(ensemble(self.traj['AB0']), False)
+        assert_equal(ensemble(self.traj['ABA']), True)
diff --git a/openpathsampling/tests/test_part_in_b_tps.py b/openpathsampling/tests/test_part_in_b_tps.py
@@ -0,0 +1,32 @@
+from __future__ import absolute_import
+from nose.tools import (assert_equal, assert_not_equal, assert_almost_equal,
+                        raises)
+import logging
+from openpathsampling.high_level.part_in_b_tps import *
+from .test_network import TestFixedLengthTPSNetwork
+
+logging.getLogger('openpathsampling.initialization').setLevel(logging.CRITICAL)
+logging.getLogger('openpathsampling.ensemble').setLevel(logging.CRITICAL)
+logging.getLogger('openpathsampling.storage').setLevel(logging.CRITICAL)
+logging.getLogger('openpathsampling.netcdfplus').setLevel(logging.CRITICAL)
+
+class TestPartInBNetwork(TestFixedLengthTPSNetwork):
+    NetworkType = PartInBFixedLengthTPSNetwork
+
+    def test_allow_self_transitions_false_ABX(self):
+        network = self.NetworkType.from_states_all_to_all(
+            self.states, allow_self_transitions=False, **self.std_kwargs
+        )
+        assert_equal(len(network.sampling_ensembles), 1)
+        ensemble = network.sampling_ensembles[0]
+        assert_equal(ensemble(self.traj['AB0']), True)
+        assert_equal(ensemble(self.traj['ABA']), True)
+
+    def test_allow_self_transitions_true_ABX(self):
+        network = self.NetworkType.from_states_all_to_all(
+            self.states, allow_self_transitions=True, **self.std_kwargs
+        )
+        assert_equal(len(network.sampling_ensembles), 1)
+        ensemble = network.sampling_ensembles[0]
+        assert_equal(ensemble(self.traj['AB0']), True)
+        assert_equal(ensemble(self.traj['ABA']), True)