Merge pull request #193 from QuantEcon/mc-fix

Fix: MarkovChain.simulate

Author: mmcky

quantecon/markov/core.py

@@ -85,6 +85,7 @@
 
 """
 from __future__ import division
+import numbers
 import numpy as np
 from scipy import sparse
 from fractions import gcd
@@ -348,10 +349,18 @@ def simulate(self, ts_length, init=None, num_reps=None, random_state=None):
         random_state = check_random_state(random_state)
         dim = 1  # Dimension of the returned array: 1 or 2
 
+        msg_out_or_range = 'index {init} is out of the state space'
+
         try:
             k = len(init)  # init is an array
             dim = 2
             init_states = np.asarray(init, dtype=int)
+            # Check init_states are in the state space
+            if (init_states >= self.n).any() or (init_states < -self.n).any():
+                idx = np.where(
+                    (init_states >= self.n) + (init_states < -self.n)
+                )[0][0]
+                raise ValueError(msg_out_or_range.format(init=idx))
             if num_reps is not None:
                 k *= num_reps
                 init_states = np.tile(init_states, num_reps)
@@ -362,7 +371,10 @@ def simulate(self, ts_length, init=None, num_reps=None, random_state=None):
                 k = num_reps
             if init is None:
                 init_states = random_state.randint(self.n, size=k)
-            elif isinstance(init, int):
+            elif isinstance(init, numbers.Integral):
+                # Check init is in the state space
+                if init >= self.n or init < -self.n:
+                    raise ValueError(msg_out_or_range.format(init=init))
                 init_states = np.ones(k, dtype=int) * init
             else:
                 raise ValueError(
@@ -526,7 +538,7 @@ def mc_sample_path(P, init=0, sample_size=1000, random_state=None):
     """
     random_state = check_random_state(random_state)
 
-    if isinstance(init, int):
+    if isinstance(init, numbers.Integral):
         X_0 = init
     else:
         cdf0 = np.cumsum(init)

quantecon/markov/tests/test_core.py

@@ -267,6 +267,21 @@ def test_simulate_init_array_num_reps():
     assert_array_equal(X[:, 0], init*num_reps)
 
 
+def test_simulate_init_type():
+    P = [[0.4, 0.6], [0.2, 0.8]]
+    mc = MarkovChain(P)
+
+    seed = 0
+    ts_length = 3
+    init = 0  # int
+    X = mc.simulate(ts_length, init=init, random_state=seed)
+
+    inits_np_int = [t(init) for t in [np.int32, np.int64]]
+    for init in inits_np_int:
+        X_np_int = mc.simulate(ts_length, init=init, random_state=seed)
+        assert_array_equal(X_np_int, X)
+
+
 def test_simulate_dense_vs_sparse():
     n = 5
     a = 1/3
@@ -393,6 +408,18 @@ def test_raises_value_error_non_sum_one():
     assert_raises(ValueError, MarkovChain, sparse.csr_matrix(P))
 
 
+def test_raises_value_error_simulate_init_out_of_range():
+    P = [[0.4, 0.6], [0.2, 0.8]]
+    mc = MarkovChain(P)
+
+    n = mc.n
+    ts_length = 3
+    assert_raises(ValueError, mc.simulate, ts_length, init=n)
+    assert_raises(ValueError, mc.simulate, ts_length, init=-(n+1))
+    assert_raises(ValueError, mc.simulate, ts_length, init=[0, n])
+    assert_raises(ValueError, mc.simulate, ts_length, init=[0, -(n+1)])
+
+
 if __name__ == '__main__':
     import sys
     import nose