Merge pull request #121 from marscher/traj_gen_random_state

added random_state handling to traj generation

Author: marscher

msmtools/generation/api.py

@@ -42,7 +42,7 @@ class MarkovChainSampler(object):
 
     """
 
-    def __init__(self, P, dt=1):
+    def __init__(self, P, dt=1, random_state=None):
         """
         Constructs a sampling object with transition matrix P. The results will be produced every dt'th time step
 
@@ -71,12 +71,26 @@ def __init__(self, P, dt=1):
         # initialize mu
         self.mudist = None
 
+        self.random_state = random_state
+
         # generate discrete random value generators for each line
-        self.rgs = np.ndarray((self.n), dtype=object)
+        self.rgs = np.ndarray(self.n, dtype=object)
         from scipy.stats import rv_discrete
-        for i in range(self.n):
-            nz = np.nonzero(self.P[i])
-            self.rgs[i] = rv_discrete(values=(nz, self.P[i, nz]))
+        for i, row in enumerate(self.P):
+            nz = row.nonzero()[0]
+            self.rgs[i] = rv_discrete(values=(nz, row[nz]))
+
+    def _get_start_state(self):
+        if self.mudist is None:
+            # compute mu, the stationary distribution of P
+            from ..analysis import stationary_distribution
+            from scipy.stats import rv_discrete
+
+            mu = stationary_distribution(self.P)
+            self.mudist = rv_discrete(values=(np.arange(self.n), mu))
+        # sample starting point from mu
+        start = self.mudist.rvs(random_state=self.random_state)
+        return start
 
     def trajectory(self, N, start=None, stop=None):
         """
@@ -97,22 +111,12 @@ def trajectory(self, N, start=None, stop=None):
         stop = types.ensure_int_vector_or_None(stop, require_order=False)
 
         if start is None:
-            if self.mudist is None:
-                # compute mu, the stationary distribution of P
-                from ..analysis import stationary_distribution
-                from scipy.stats import rv_discrete
-
-                mu = stationary_distribution(self.P)
-                self.mudist = rv_discrete(values=(np.arange(self.n), mu))
-            # sample starting point from mu
-            start = self.mudist.rvs()
+          start = self._get_start_state()
 
         # evaluate stopping set
-        stopat = np.ndarray((self.n), dtype=bool)
-        stopat[:] = False
-        if (stop is not None):
-            for s in np.array(stop):
-                stopat[s] = True
+        stopat = np.zeros(self.n, dtype=bool)
+        if stop is not None:
+            stopat[np.array(stop)] = True
 
         # result
         traj = np.zeros(N, dtype=int)
@@ -122,9 +126,10 @@ def trajectory(self, N, start=None, stop=None):
             return traj[:1]
         # else run until end or stopping state
         for t in range(1, N):
-            traj[t] = self.rgs[traj[t - 1]].rvs()
+            traj[t] = self.rgs[traj[t - 1]].rvs(random_state=self.random_state)
             if stopat[traj[t]]:
-                return traj[:t+1]
+                traj = np.resize(traj, t + 1)
+                break
         # return
         return traj
 
@@ -149,7 +154,7 @@ def trajectories(self, M, N, start=None, stop=None):
         return trajs
 
 
-def generate_traj(P, N, start=None, stop=None, dt=1):
+def generate_traj(P, N, start=None, stop=None, dt=1, random_state=None):
     """
     Generates a realization of the Markov chain with transition matrix P.
 
@@ -167,18 +172,22 @@ def generate_traj(P, N, start=None, stop=None, dt=1):
     dt : int
         trajectory will be saved every dt time steps.
         Internally, the dt'th power of P is taken to ensure a more efficient simulation.
+    random_state : None or int or numpy.random.RandomState instance, optional
+        This parameter defines the RandomState object to use for drawing random variates.
+        If None, the global np.random state is used. If integer, it is used to seed the local RandomState instance.
+        Default is None.
 
     Returns
     -------
     traj_sliced : (N/dt, ) ndarray
         A discrete trajectory with length N/dt
 
     """
-    sampler = MarkovChainSampler(P, dt=dt)
+    sampler = MarkovChainSampler(P, dt=dt, random_state=random_state)
     return sampler.trajectory(N, start=start, stop=stop)
 
 
-def generate_trajs(P, M, N, start=None, stop=None, dt=1):
+def generate_trajs(P, M, N, start=None, stop=None, dt=1, random_state=None):
     """
     Generates multiple realizations of the Markov chain with transition matrix P.
 
@@ -198,14 +207,18 @@ def generate_trajs(P, M, N, start=None, stop=None, dt=1):
     dt : int
         trajectory will be saved every dt time steps.
         Internally, the dt'th power of P is taken to ensure a more efficient simulation.
+    random_state : None or int or numpy.random.RandomState instance, optional
+        This parameter defines the RandomState object to use for drawing random variates.
+        If None, the global np.random state is used. If integer, it is used to seed the local RandomState instance.
+        Default is None.
 
     Returns
     -------
     traj_sliced : (N/dt, ) ndarray
         A discrete trajectory with length N/dt
 
     """
-    sampler = MarkovChainSampler(P, dt=dt)
+    sampler = MarkovChainSampler(P, dt=dt, random_state=random_state)
     return sampler.trajectories(M, N, start=start, stop=stop)
 
 
@@ -235,12 +248,12 @@ def transition_matrix_metropolis_1d(E, d=1.0):
 
     """
     # check input
-    if (d <= 0 or d > 1):
+    if d <= 0 or d > 1:
         raise ValueError('Diffusivity must be in (0,1]. Trying to set the invalid value', str(d))
     # init
     n = len(E)
     P = np.zeros((n, n))
-    # set offdiagonals
+    # set off diagonals
     P[0, 1] = 0.5 * d * min(1.0, math.exp(-(E[1] - E[0])))
     for i in range(1, n - 1):
         P[i, i - 1] = 0.5 * d * min(1.0, math.exp(-(E[i - 1] - E[i])))

tests/generation/test_generation.py

@@ -1,4 +1,3 @@
-
 # This file is part of MSMTools.
 #
 # Copyright (c) 2015, 2014 Computational Molecular Biology Group
@@ -25,66 +24,63 @@
 import msmtools.estimation as msmest
 import msmtools.analysis as msmana
 
-class Test(unittest.TestCase):
 
-    def setUp(self):
-        """Safe random state"""
-        self.state = np.random.get_state()
-        """Set seed to enforce deterministic behavior"""
-        np.random.seed(42)
+class TestTrajGeneration(unittest.TestCase):
 
-    def tearDown(self):
-        """Reset state"""
-        np.random.set_state(self.state)
+    @classmethod
+    def setUpClass(cls):
+        cls.P = np.array([[0.9, 0.1],
+                          [0.1, 0.9]])
+
+    def setUp(self):
+        self.random_state = np.random.RandomState(42)
 
     def test_trajectory(self):
-        P = np.array([[0.9,0.1],
-                      [0.1,0.9]])
         N = 1000
-        traj = msmgen.generate_traj(P, N, start=0)
+        traj = msmgen.generate_traj(self.P, N, start=0, random_state=self.random_state)
 
         # test shapes and sizes
         assert traj.size == N
         assert traj.min() >= 0
         assert traj.max() <= 1
 
         # test statistics of transition matrix
-        C = msmest.count_matrix(traj,1)
+        C = msmest.count_matrix(traj, 1)
         Pest = msmest.transition_matrix(C)
-        assert np.max(np.abs(Pest - P)) < 0.025
-
+        assert np.max(np.abs(Pest - self.P)) < 0.025
 
     def test_trajectories(self):
-        P = np.array([[0.9,0.1],
-                      [0.1,0.9]])
-
         # test number of trajectories
         M = 10
         N = 10
-        trajs = msmgen.generate_trajs(P, M, N, start=0)
+        trajs = msmgen.generate_trajs(self.P, M, N, start=0, random_state=self.random_state)
         assert len(trajs) == M
 
+    def test_stats(self):
         # test statistics of starting state
-        trajs = msmgen.generate_trajs(P, 1000, 1)
+        N = 5000
+        trajs = msmgen.generate_trajs(self.P, N, 1, random_state=self.random_state)
         ss = np.concatenate(trajs).astype(int)
-        pi = msmana.stationary_distribution(P)
-        piest = msmest.count_states(ss) / 1000.0
-        assert np.max(np.abs(pi - piest)) < 0.025
+        pi = msmana.stationary_distribution(self.P)
+        piest = msmest.count_states(ss) / float(N)
+        np.testing.assert_allclose(piest, pi, atol=0.025)
 
-        # test stopping state = starting state
+    def test_stop_eq_start(self):
         M = 10
-        trajs = msmgen.generate_trajs(P, M, N, start=0, stop=0)
+        N = 10
+        trajs = msmgen.generate_trajs(self.P, M, N, start=0, stop=0, random_state=self.random_state)
         for traj in trajs:
             assert traj.size == 1
 
+    def test_stop(self):
         # test if we always stop at stopping state
         M = 100
+        N = 10
         stop = 1
-        trajs = msmgen.generate_trajs(P, M, N, start=0, stop=stop)
+        trajs = msmgen.generate_trajs(self.P, M, N, start=0, stop=stop, random_state=self.random_state)
         for traj in trajs:
             assert traj.size == N or traj[-1] == stop
             assert stop not in traj[:-1]
 
 if __name__ == "__main__":
-    # import sys;sys.argv = ['', 'Test.testName']
     unittest.main()