diff --git a/quantecon/arma.py b/quantecon/arma.py index e202c028b..e6a7e7a31 100644 --- a/quantecon/arma.py +++ b/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() diff --git a/quantecon/tests/test_arma.py b/quantecon/tests/test_arma.py index 31bc23d35..7b9d461c9 100644 --- a/quantecon/tests/test_arma.py +++ b/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,6 +35,14 @@ 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 @@ -44,6 +50,7 @@ def test_impulse_response(self): self.assertTrue(imp_resp.size==75) + if __name__ == '__main__': suite = unittest.TestLoader().loadTestsFromTestCase(TestARMA) unittest.TextTestRunner(verbosity=2, stream=sys.stderr).run(suite)