Merge pull request #287 from choderalab/testsystem_control_seed

Add the ability to control the random seed to testsystem

pymbar/testsystems/exponential_distributions.py

@@ -90,7 +90,7 @@ def analytical_entropies(self):
     def analytical_x_squared(self):
         return self.analytical_variances() + self.analytical_means() ** 2.
 
-    def sample(self, N_k=(10, 20, 30, 40, 50), mode='u_kln'):
+    def sample(self, N_k=(10, 20, 30, 40, 50), mode='u_kln', seed=None):
         """Draw samples from the distribution.
 
         Parameters
@@ -101,6 +101,7 @@ def sample(self, N_k=(10, 20, 30, 40, 50), mode='u_kln'):
         mode : str, optional, default='u_kln'
             If 'u_kln', return K x K x N_max matrix where u_kln[k,l,n] is reduced potential of sample n from state k evaluated at state l.
             If 'u_kn', return K x N_tot matrix where u_kn[k,n] is reduced potential of sample n (in concatenated indexing) evaluated at state k.
+        seed: int, optional, default=None.  Provides control over the random seed for replicability.
 
         Returns
         -------
@@ -123,6 +124,9 @@ def sample(self, N_k=(10, 20, 30, 40, 50), mode='u_kln'):
            N_k[k] is the number of samples generated from state k
 
         """
+
+        np.random.seed(seed)
+
         N_k = np.array(N_k, np.int32)
         if len(N_k) != self.n_states:
             raise Exception("N_k has %d states while self.n_states has %d states." % (len(N_k), self.n_states))

pymbar/testsystems/harmonic_oscillators.py

@@ -97,7 +97,7 @@ def analytical_free_energies(self, subtract_component=0):
     def analytical_entropies(self, subtract_component = 0):
         return self.analytical_observable(observable = 'potential energy') - self.analytical_free_energies(subtract_component)
 
-    def sample(self, N_k=[10, 20, 30, 40, 50], mode='u_kn'):
+    def sample(self, N_k=[10, 20, 30, 40, 50], mode='u_kn', seed = None):
         """Draw samples from the distribution.
 
         Parameters
@@ -109,6 +109,8 @@ def sample(self, N_k=[10, 20, 30, 40, 50], mode='u_kn'):
             If 'u_kln', return K x K x N_max matrix where u_kln[k,l,n] is reduced potential of sample n from state k evaluated at state l.
             If 'u_kn', return K x N_tot matrix where u_kn[k,n] is reduced potential of sample n (in concatenated indexing) evaluated at state k.
 
+        seed: int, optional, default=None.  Provides control over the random seed for replicability.
+
         Returns
         -------
         if mode == 'u_kn':
@@ -130,6 +132,9 @@ def sample(self, N_k=[10, 20, 30, 40, 50], mode='u_kn'):
            N_k[k] is the number of samples generated from state k
 
         """
+
+        np.random.seed(seed)
+
         N_k = np.array(N_k, np.int32)
         if len(N_k) != self.n_states:
             raise Exception("N_k has %d states while self.n_states has %d states." % (len(N_k), self.n_states))
@@ -207,6 +212,6 @@ def evenly_spaced_oscillators(cls, n_states, n_samples_per_state, lower_O_k=1.0,
         N_k = (np.ones(n_states) * n_samples_per_state).astype('int')
 
         testsystem = cls(O_k, k_k)
-        x_n, u_kn, N_k_output, s_n = testsystem.sample(N_k, mode='u_kn')
+        x_n, u_kn, N_k_output, s_n = testsystem.sample(N_k, mode='u_kn', seed=seed)
 
         return name, testsystem, x_n, u_kn, N_k_output, s_n