Permalink
Browse files

ddp: Export sa_indices (#255)

  • Loading branch information...
1 parent fb42379 commit e73730ceb934ed9dda603d7f6e17e57b5026db02 @oyamad oyamad committed with mmcky May 18, 2016
@@ -9,3 +9,4 @@ Markov
markov/ddp
markov/gth_solve
markov/random
+ markov/utilities
@@ -0,0 +1,7 @@
+utilities
+=========
+
+.. automodule:: quantecon.markov.utilities
+ :members:
+ :undoc-members:
+ :show-inheritance:
@@ -6,7 +6,6 @@ Utilities
util/array
util/common_messages
- util/external
util/notebooks
util/random
util/timing
@@ -9,3 +9,4 @@
random_discrete_dp
from .approximation import tauchen
from .ddp import DiscreteDP, backward_induction
+from .utilities import sa_indices
@@ -12,6 +12,7 @@
from .core import MarkovChain
from .ddp import DiscreteDP
+from .utilities import sa_indices
from ..util import check_random_state
from ..random import probvec, sample_without_replacement
@@ -210,27 +211,11 @@ def random_discrete_dp(num_states, num_actions, beta=None,
beta = random_state.random_sample()
if sa_pair:
- s_indices, a_indices = _sa_indices(num_states, num_actions)
+ s_indices, a_indices = sa_indices(num_states, num_actions)
else:
s_indices, a_indices = None, None
R.shape = (num_states, num_actions)
Q.shape = (num_states, num_actions, num_states)
ddp = DiscreteDP(R, Q, beta, s_indices, a_indices)
return ddp
-
-@jit
-def _sa_indices(num_states, num_actions):
- L = num_states * num_actions
- s_indices = np.empty(L, dtype=int)
- a_indices = np.empty(L, dtype=int)
-
- i = 0
- for s in range(num_states):
- for a in range(num_actions):
- s_indices[i] = s
- a_indices[i] = a
- i += 1
-
- return s_indices, a_indices
-
@@ -0,0 +1,30 @@
+"""
+Filename: test_utilities.py
+Author: Daisuke Oyama
+
+Tests for markov/utilities.py
+
+"""
+from numpy.testing import assert_array_equal
+from quantecon.markov import sa_indices
+
+
+def test_sa_indices():
+ num_states, num_actions = 3, 4
+ s_expected = [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]
+ a_expected = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]
+
+ s, a = sa_indices(num_states, num_actions)
+
+ for indices, indices_expected in zip([s, a], [s_expected, a_expected]):
+ assert_array_equal(indices, indices_expected)
+
+
+if __name__ == '__main__':
+ import sys
+ import nose
+
+ argv = sys.argv[:]
+ argv.append('--verbose')
+ argv.append('--nocapture')
+ nose.main(argv=argv, defaultTest=__file__)
@@ -0,0 +1,51 @@
+"""
+Utility routines for the markov submodule
+
+"""
+import numpy as np
+from numba import jit
+
+@jit(nopython=True)
+def sa_indices(num_states, num_actions):
+ """
+ Generate `s_indices` and `a_indices` for `DiscreteDP`, for the case
+ where all the actions are feasible at every state.
+
+ Parameters
+ ----------
+ num_states : scalar(int)
+ Number of states.
+
+ num_actions : scalar(int)
+ Number of actions.
+
+ Returns
+ -------
+ s_indices : ndarray(int, ndim=1)
+ Array containing the state indices.
+
+ a_indices : ndarray(int, ndim=1)
+ Array containing the action indices.
+
+ Examples
+ --------
+ >>> s_indices, a_indices = qe.markov.sa_indices(4, 3)
+ >>> s_indices
+ array([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3])
+ >>> a_indices
+ array([0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2])
+
+ """
+ L = num_states * num_actions
+ dtype = np.int_
+ s_indices = np.empty(L, dtype=dtype)
+ a_indices = np.empty(L, dtype=dtype)
+
+ i = 0
+ for s in range(num_states):
+ for a in range(num_actions):
+ s_indices[i] = s
+ a_indices[i] = a
+ i += 1
+
+ return s_indices, a_indices

0 comments on commit e73730c

Please sign in to comment.