FEAT: Add option to supply a random seed (issue #153) (#346)

* FEAT: Add random_state option to discrete_rv.py

* TEST: Add test for discrete_rv.py random state option

* FEAT: Add random_state option to lqcontrol.py

* TEST: Add test for lqcontrol.py random state option

* FEAT: Add random_state option to lqnash.py

* FIX: Make formatting of test_lqnash.py consistent with other tests

* FEAT: Add random_state option to lss.py

* TEST: Add test for lss.py random state option

* FEAT: Add random_state option to quad.py

* TEST: Add test for quad.py random state option

* FIX: Correct typo

* FIX: Correct the docstring for the random_state options

quantecon/discrete_rv.py

@@ -11,7 +11,7 @@
 import numpy as np
 from numpy import cumsum
 from numpy.random import uniform
-
+from .util import check_random_state
 
 class DiscreteRV:
     """
@@ -58,7 +58,7 @@ def q(self, val):
         self._q = np.asarray(val)
         self.Q = cumsum(val)
 
-    def draw(self, k=1):
+    def draw(self, k=1, random_state=None):
         """
         Returns k draws from q.
 
@@ -70,10 +70,18 @@ def draw(self, k=1):
         k : scalar(int), optional
             Number of draws to be returned
 
+        random_state : int or np.random.RandomState, optional
+            Random seed (integer) or np.random.RandomState instance to set
+            the initial state of the random number generator for
+            reproducibility. If None, a randomly initialized RandomState is
+            used.
+
         Returns
         -------
         array_like(int)
             An array of k independent draws from q
 
         """
-        return self.Q.searchsorted(uniform(0, 1, size=k))
+        random_state = check_random_state(random_state)
+
+        return self.Q.searchsorted(random_state.uniform(0, 1, size=k))

quantecon/lqcontrol.py

@@ -12,6 +12,7 @@
 from numpy import dot
 from scipy.linalg import solve
 from .matrix_eqn import solve_discrete_riccati
+from .util import check_random_state
 
 
 class LQ:
@@ -241,7 +242,7 @@ def stationary_values(self):
 
         return P, F, d
 
-    def compute_sequence(self, x0, ts_length=None):
+    def compute_sequence(self, x0, ts_length=None, random_state=None):
         """
         Compute and return the optimal state and control sequences
         :math:`x_0, ..., x_T` and :math:`u_0,..., u_T`  under the
@@ -255,6 +256,12 @@ def compute_sequence(self, x0, ts_length=None):
         ts_length : scalar(int)
             Length of the simulation -- defaults to T in finite case
 
+        random_state : int or np.random.RandomState, optional
+            Random seed (integer) or np.random.RandomState instance to set
+            the initial state of the random number generator for
+            reproducibility. If None, a randomly initialized RandomState is
+            used.
+
         Returns
         ========
         x_path : array_like(float)
@@ -282,11 +289,12 @@ def compute_sequence(self, x0, ts_length=None):
             self.stationary_values()
 
         # == Set up initial condition and arrays to store paths == #
+        random_state = check_random_state(random_state)
         x0 = np.asarray(x0)
         x0 = x0.reshape(self.n, 1)  # Make sure x0 is a column vector
         x_path = np.empty((self.n, T+1))
         u_path = np.empty((self.k, T))
-        w_path = np.random.randn(self.j, T+1)
+        w_path = random_state.randn(self.j, T+1)
         Cw_path = dot(C, w_path)
 
         # == Compute and record the sequence of policies == #

quantecon/lqnash.py

@@ -3,10 +3,11 @@
 import numpy as np
 from numpy import dot, eye
 from scipy.linalg import solve
+from .util import check_random_state
 
 
 def nnash(A, B1, B2, R1, R2, Q1, Q2, S1, S2, W1, W2, M1, M2,
-          beta=1.0, tol=1e-8, max_iter=1000):
+          beta=1.0, tol=1e-8, max_iter=1000, random_state=None):
     r"""
     Compute the limit of a Nash linear quadratic dynamic game. In this
     problem, player i minimizes
@@ -65,6 +66,11 @@ def nnash(A, B1, B2, R1, R2, Q1, Q2, S1, S2, W1, W2, M1, M2,
         This is the tolerance level for convergence
     max_iter : scalar(int), optional(default=1000)
         This is the maximum number of iteratiosn allowed
+    random_state : int or np.random.RandomState, optional
+        Random seed (integer) or np.random.RandomState instance to set
+        the initial state of the random number generator for
+        reproducibility. If None, a randomly initialized RandomState is
+        used.
 
     Returns
     -------
@@ -102,12 +108,13 @@ def nnash(A, B1, B2, R1, R2, Q1, Q2, S1, S2, W1, W2, M1, M2,
     else:
         k_2 = B2.shape[1]
 
+    random_state = check_random_state(random_state)
     v1 = eye(k_1)
     v2 = eye(k_2)
     P1 = np.zeros((n, n))
     P2 = np.zeros((n, n))
-    F1 = np.random.randn(k_1, n)
-    F2 = np.random.randn(k_2, n)
+    F1 = random_state.randn(k_1, n)
+    F2 = random_state.randn(k_2, n)
 
     for it in range(max_iter):
         # update

quantecon/lss.py

@@ -10,6 +10,7 @@
 from numpy.random import multivariate_normal
 from scipy.linalg import solve
 from numba import jit
+from .util import check_random_state
 
 
 @jit
@@ -148,7 +149,7 @@ def convert(self, x):
         """
         return np.atleast_2d(np.asarray(x, dtype='float'))
 
-    def simulate(self, ts_length=100):
+    def simulate(self, ts_length=100, random_state=None):
         r"""
         Simulate a time series of length ts_length, first drawing
 
@@ -158,9 +159,13 @@ def simulate(self, ts_length=100):
 
         Parameters
         ----------
-
         ts_length : scalar(int), optional(default=100)
             The length of the simulation
+        random_state : int or np.random.RandomState, optional
+            Random seed (integer) or np.random.RandomState instance to set
+            the initial state of the random number generator for
+            reproducibility. If None, a randomly initialized RandomState is
+            used.
 
         Returns
         -------
@@ -170,21 +175,23 @@ def simulate(self, ts_length=100):
             A k x ts_length array, where the t-th column is :math:`y_t`
 
         """
+        random_state = check_random_state(random_state)
+
         x0 = multivariate_normal(self.mu_0.flatten(), self.Sigma_0)
-        w = np.random.randn(self.m, ts_length-1)
+        w = random_state.randn(self.m, ts_length-1)
         v = self.C.dot(w)  # Multiply each w_t by C to get v_t = C w_t
         # == simulate time series == #
         x = simulate_linear_model(self.A, x0, v, ts_length)
 
         if self.H is not None:
-            v = np.random.randn(self.l, ts_length)
+            v = random_state.randn(self.l, ts_length)
             y = self.G.dot(x) + self.H.dot(v)
         else:
             y = self.G.dot(x)
 
         return x, y
 
-    def replicate(self, T=10, num_reps=100):
+    def replicate(self, T=10, num_reps=100, random_state=None):
         r"""
         Simulate num_reps observations of :math:`x_T` and :math:`y_T` given
         :math:`x_0 \sim N(\mu_0, \Sigma_0)`.
@@ -195,6 +202,11 @@ def replicate(self, T=10, num_reps=100):
             The period that we want to replicate values for
         num_reps : scalar(int), optional(default=100)
             The number of replications that we want
+        random_state : int or np.random.RandomState, optional
+            Random seed (integer) or np.random.RandomState instance to set
+            the initial state of the random number generator for
+            reproducibility. If None, a randomly initialized RandomState is
+            used.
 
         Returns
         -------
@@ -207,12 +219,14 @@ def replicate(self, T=10, num_reps=100):
             observation of :math:`y_T`
 
         """
+        random_state = check_random_state(random_state)
+
         x = np.empty((self.n, num_reps))
         for j in range(num_reps):
-            x_T, _ = self.simulate(ts_length=T+1)
+            x_T, _ = self.simulate(ts_length=T+1, random_state=random_state)
             x[:, j] = x_T[:, -1]
         if self.H is not None:
-            v = np.random.randn(self.l, num_reps)
+            v = random_state.randn(self.l, num_reps)
             y = self.G.dot(x) + self.H.dot(v)
         else:
             y = self.G.dot(x)

quantecon/quad.py

@@ -22,6 +22,7 @@
 from scipy.special import gammaln
 import sympy as sym
 from .ce_util import ckron, gridmake
+from .util import check_random_state
 
 __all__ = ['qnwcheb', 'qnwequi', 'qnwlege', 'qnwnorm', 'qnwlogn',
            'qnwsimp', 'qnwtrap', 'qnwunif', 'quadrect', 'qnwbeta',
@@ -73,7 +74,7 @@ def qnwcheb(n, a=1, b=1):
     return _make_multidim_func(_qnwcheb1, n, a, b)
 
 
-def qnwequi(n, a, b, kind="N", equidist_pp=None):
+def qnwequi(n, a, b, kind="N", equidist_pp=None, random_state=None):
     """
     Generates equidistributed sequences with property that averages
     value of integrable function evaluated over the sequence converges
@@ -105,6 +106,12 @@ def qnwequi(n, a, b, kind="N", equidist_pp=None):
     equidist_pp : array_like, optional(default=None)
         TODO: I don't know what this does
 
+    random_state : int or np.random.RandomState, optional
+        Random seed (integer) or np.random.RandomState instance to set
+        the initial state of the random number generator for
+        reproducibility. If None, a randomly initialized RandomState is
+        used.
+
     Returns
     -------
     nodes : np.ndarray(dtype=float)
@@ -124,6 +131,8 @@ def qnwequi(n, a, b, kind="N", equidist_pp=None):
     Economics and Finance, MIT Press, 2002.
 
     """
+    random_state = check_random_state(random_state)
+
     if equidist_pp is None:
         equidist_pp = np.sqrt(np.array(list(sym.primerange(0, 7920))))
 
@@ -153,7 +162,7 @@ def qnwequi(n, a, b, kind="N", equidist_pp=None):
         nodes = np.outer(i * (i+1) / 2, j)
         nodes = (nodes - np.fix(nodes)).squeeze()
     elif kind.upper() == "R":  # pseudo-random
-        nodes = np.random.rand(n, d).squeeze()
+        nodes = random_state.rand(n, d).squeeze()
     else:
         raise ValueError("Unknown sequence requested")
 

quantecon/tests/test_discrete_rv.py

@@ -56,3 +56,12 @@ def test_draw_lln(self):
         counts = (counts / counts.sum()).values
         assert max(np.abs(counts - self.drv.q)) < 1e-2
 
+    def test_draw_with_seed(self):
+        x = np.array([0.03326189, 0.60713005, 0.84514831, 0.28493183,
+                      0.12393182, 0.35308009, 0.70371579, 0.81728178,
+                      0.21294538, 0.05358209])
+        draws = DiscreteRV(x).draw(k=10, random_state=5)
+
+        expected_output = np.array([1, 2, 1, 2, 1, 1, 2, 1, 1, 1])
+
+        assert_allclose(draws, expected_output)

quantecon/tests/test_lqcontrol.py

@@ -61,6 +61,14 @@ def test_scalar_sequences(self):
         assert_allclose(u_0, u_seq, rtol=1e-4)
         assert_allclose(x_1, x_seq[0, -1], rtol=1e-4)
 
+    def test_scalar_sequences_with_seed(self):
+        lq_scalar = self.lq_scalar
+        x0 = 2
+        x_seq, u_seq, w_seq = lq_scalar.compute_sequence(x0, 10, 5)
+
+        expected_output = np.array([[ 0.44122749, -0.33087015]])
+
+        assert_allclose(w_seq, expected_output)
 
     def test_mat_sequences(self):
 
@@ -88,8 +96,6 @@ def test_stationary_mat(self):
         assert_allclose(val_func_lq, val_func_answer, atol=1e-3)
 
 
-
-
 if __name__ == '__main__':
     suite = unittest.TestLoader().loadTestsFromTestCase(TestLQControl)
     unittest.TextTestRunner(verbosity=2, stream=sys.stderr).run(suite)