From 95f3efb9a1141a32da891f64839bb4b5c5d3cd5a Mon Sep 17 00:00:00 2001 From: QBatista Date: Wed, 23 Aug 2017 17:16:45 +0900 Subject: [PATCH] FEAT: Add random_state option to arma.py with tests --- quantecon/arma.py | 13 +++++++++++-- quantecon/tests/test_arma.py | 13 ++++++++++--- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/quantecon/arma.py b/quantecon/arma.py index f9c5d508f..c7de330a4 100644 --- a/quantecon/arma.py +++ b/quantecon/arma.py @@ -11,6 +11,7 @@ from numpy import conj, pi import matplotlib.pyplot as plt from scipy.signal import dimpulse, freqz, dlsim +from .util import check_random_state class ARMA: @@ -231,7 +232,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. @@ -240,14 +241,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)