diff --git a/quantecon/markov/ddp.py b/quantecon/markov/ddp.py index fb0f035d..6ef33287 100644 --- a/quantecon/markov/ddp.py +++ b/quantecon/markov/ddp.py @@ -178,6 +178,11 @@ class DiscreteDP: a_indices : array_like(int, ndim=1), optional(default=None) Array containing the indices of the actions. + state_values : array_like, optional(default=None) + Array_like of length num_states containing the values associated with + the states, which must be homogeneous in type. If None, the values + default to integers 0 through num_states-1. + Attributes ---------- R, Q, beta : see Parameters. @@ -195,6 +200,9 @@ class DiscreteDP: max_iter : scalar(int), default=250 Default value for the maximum number of iterations. + state_values : array_like or None + Array of state values if set, None otherwise. + Notes ----- DiscreteDP accepts beta=1 for convenience. In this case, infinite @@ -296,7 +304,8 @@ class DiscreteDP: 4 """ - def __init__(self, R, Q, beta, s_indices=None, a_indices=None): + def __init__(self, R, Q, beta, s_indices=None, a_indices=None, + state_values=None): if not (0 <= beta <= 1): raise ValueError('beta must be in [0, 1]') if beta == 1: @@ -413,6 +422,21 @@ def __init__(self, R, Q, beta, s_indices=None, a_indices=None): self.epsilon = 1e-3 self.max_iter = 250 + # State labels are for mapping only, not numerical computations + if state_values is None: + self.state_values = None + else: + state_values = np.asarray(state_values) + if (state_values.ndim < 1 or + state_values.shape[0] != self.num_states): + raise ValueError( + 'state_values must be an array_like of length num_states' + ) + if np.issubdtype(state_values.dtype, np.object_): + raise ValueError( + 'data in state_values must be homogeneous in type' + ) + self.state_values = state_values # Linear equation solver to be used in evaluate_policy if self._sparse: @@ -511,7 +535,10 @@ def to_sa_pair_form(self, sparse=True): QL = sp.csr_matrix(self.Q[s_ind, a_ind]) else: QL = self.Q[s_ind, a_ind] - return DiscreteDP(RL, QL, self.beta, s_ind, a_ind) + return DiscreteDP( + RL, QL, self.beta, s_ind, a_ind, + state_values=self.state_values + ) def to_product_form(self): """ @@ -542,7 +569,9 @@ def to_product_form(self): self.Q.toarray(), Q) else: _fill_dense_Q(self.s_indices, self.a_indices, self.Q, Q) - return DiscreteDP(R, Q, self.beta) + return DiscreteDP( + R, Q, self.beta, state_values=self.state_values + ) else: return self @@ -988,7 +1017,7 @@ def controlled_mc(self, sigma): """ _, Q_sigma = self.RQ_sigma(sigma) - return MarkovChain(Q_sigma) + return MarkovChain(Q_sigma, state_values=self.state_values) class DPSolveResult(dict): diff --git a/quantecon/markov/tests/test_ddp.py b/quantecon/markov/tests/test_ddp.py index 2997ec23..8335cf6e 100644 --- a/quantecon/markov/tests/test_ddp.py +++ b/quantecon/markov/tests/test_ddp.py @@ -447,6 +447,7 @@ def test_ddp_to_sa_and_to_product(): Q[0, 0, 0] = 0 Q[0, 0, 1] = 2/n beta = 0.95 + state_values = np.array(['state0', 'state1', 'state2']) sparse_R = np.array([0, 1, 1, 0, 1]) _Q = np.full((5, 3), 1/3) @@ -454,7 +455,7 @@ def test_ddp_to_sa_and_to_product(): _Q[0, 1] = 2/n sparse_Q = sparse.coo_matrix(_Q) - ddp = DiscreteDP(R, Q, beta) + ddp = DiscreteDP(R, Q, beta, state_values=state_values) ddp_sa = ddp.to_sa_pair_form() ddp_sa2 = ddp_sa.to_sa_pair_form() ddp_sa3 = ddp.to_sa_pair_form(sparse=False) @@ -468,6 +469,7 @@ def test_ddp_to_sa_and_to_product(): # allcose doesn't work on sparse np.max(np.abs((sparse_Q - ddp_s.Q))) < 1e-15 assert_allclose(ddp_s.beta, beta) + assert_array_equal(ddp_s.state_values, state_values) # these two will have probability 0 in state 2, action 0 b/c # of the infeasiability in R @@ -480,11 +482,13 @@ def test_ddp_to_sa_and_to_product(): assert_allclose(ddp_f.R, ddp.R) assert_allclose(ddp_f.Q, funky_Q) assert_allclose(ddp_f.beta, ddp.beta) + assert_array_equal(ddp_f.state_values, state_values) # this one is just the original one. assert_allclose(ddp4.R, ddp.R) assert_allclose(ddp4.Q, ddp.Q) assert_allclose(ddp4.beta, ddp.beta) + assert_array_equal(ddp4.state_values, state_values) for method in ["pi", "vi", "mpi"]: sol1 = ddp.solve(method=method) @@ -493,3 +497,59 @@ def test_ddp_to_sa_and_to_product(): for k in ["v", "sigma", "num_iter"]: assert_allclose(sol1[k], sol2[k]) + + +def test_ddp_state_values(): + R = np.array([[1, 2], [3, 4]]) + Q = np.full((2, 2, 2), 0.5) + beta = 0.95 + state_values = np.array(['state1', 'state2']) + + ddp = DiscreteDP(R, Q, beta, state_values=state_values) + + assert_array_equal(ddp.state_values, state_values) + + +def test_ddp_controlled_mc_state_values(): + R = np.array([[1, 2], [3, 4]]) + Q = np.full((2, 2, 2), 0.5) + beta = 0.95 + state_values = np.array(['state1', 'state2']) + sigma = np.array([0, 1]) + + ddp = DiscreteDP(R, Q, beta, state_values=state_values) + mc = ddp.controlled_mc(sigma) + + assert_array_equal(mc.state_values, state_values) + + +def test_ddp_state_values_wrong_length(): + R = np.array([[1, 2], [3, 4]]) + Q = np.full((2, 2, 2), 0.5) + beta = 0.95 + + assert_raises( + ValueError, DiscreteDP, R, Q, beta, + state_values=np.array(['state1']) + ) + + +def test_ddp_scalar_state_values(): + R = np.array([[1, 2], [3, 4]]) + Q = np.full((2, 2, 2), 0.5) + beta = 0.95 + + assert_raises( + ValueError, DiscreteDP, R, Q, beta, state_values='state1' + ) + + +def test_ddp_object_dtype_state_values(): + R = np.array([[1, 2], [3, 4]]) + Q = np.full((2, 2, 2), 0.5) + beta = 0.95 + state_values = np.array(['state1', 2], dtype=object) + + assert_raises( + ValueError, DiscreteDP, R, Q, beta, state_values=state_values + )