ENH: Add state_values to DiscreteDP

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):

quantecon/markov/tests/test_ddp.py

@@ -447,14 +447,15 @@ 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)
     _Q[0, 0] = 0
     _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
+    )