FEAT: Add random_state option to arma.py with tests (#329)

quantecon/arma.py

@@ -10,6 +10,7 @@
 import numpy as np
 from numpy import conj, pi
 from scipy.signal import dimpulse, freqz, dlsim
+from .util import check_random_state
 
 
 class ARMA:
@@ -230,7 +231,7 @@ def autocovariance(self, num_autocov=16):
         # num_autocov should be <= len(acov) / 2
         return acov[:num_autocov]
 
-    def simulation(self, ts_length=90):
+    def simulation(self, ts_length=90, random_state=None):
         """
         Compute a simulated sample path assuming Gaussian shocks.
 
@@ -239,14 +240,22 @@ def simulation(self, ts_length=90):
         ts_length : scalar(int), optional(default=90)
             Number of periods to simulate for
 
+        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
         -------
         vals : array_like(float)
             A simulation of the model that corresponds to this class
 
         """
+        random_state = check_random_state(random_state)
+
         sys = self.ma_poly, self.ar_poly, 1
-        u = np.random.randn(ts_length, 1) * self.sigma
+        u = random_state.randn(ts_length, 1) * self.sigma
         vals = dlsim(sys, u)[1]
 
         return vals.flatten()

quantecon/tests/test_arma.py

@@ -11,12 +11,11 @@
 import os
 import unittest
 import numpy as np
-from numpy.testing import assert_allclose
+from numpy.testing import assert_array_equal
 from quantecon.arma import ARMA
 
 
 class TestARMA(unittest.TestCase):
-
     def setUp(self):
         # Initial Values
         phi = np.array([.95, -.4, -.4])
@@ -26,7 +25,6 @@ def setUp(self):
 
         self.lp = ARMA(phi, theta, sigma)
 
-
     def tearDown(self):
         del self.lp
 
@@ -37,13 +35,22 @@ def test_simulate(self):
 
         self.assertTrue(sim.size==250)
 
+    def test_simulate_with_seed(self):
+        lp = self.lp
+        seed = 5
+        sim0 = lp.simulation(ts_length=10, random_state=seed)
+        sim1 = lp.simulation(ts_length=10, random_state=seed)
+
+        assert_array_equal(sim0, sim1)
+
     def test_impulse_response(self):
         lp = self.lp
 
         imp_resp = lp.impulse_response(impulse_length=75)
 
         self.assertTrue(imp_resp.size==75)
 
+
 if __name__ == '__main__':
     suite = unittest.TestLoader().loadTestsFromTestCase(TestARMA)
     unittest.TextTestRunner(verbosity=2, stream=sys.stderr).run(suite)