ENH: Add action_values to DiscreteDP; add util.index_dict

docs/source/util.rst

@@ -8,6 +8,7 @@ Utilities
    util/combinatorics
    util/common_messages
    util/compat
+   util/indexing
    util/notebooks
    util/numba
    util/random

docs/source/util/indexing.rst

@@ -0,0 +1,7 @@
+indexing
+========
+
+.. automodule:: quantecon.util.indexing
+    :members:
+    :undoc-members:
+    :show-inheritance:

quantecon/__init__.py

@@ -53,4 +53,5 @@
 #<-
 from ._rank_nullspace import rank_est, nullspace
 from ._robustlq import RBLQ
-from .util import searchsorted, fetch_nb_dependencies, tic, tac, toc, Timer, timeit
+from .util import searchsorted, index_dict, fetch_nb_dependencies, \
+    tic, tac, toc, Timer, timeit

quantecon/markov/core.py

@@ -104,8 +104,8 @@ class MarkovChain:
 
     state_values : array_like(default=None)
         Array_like of length n containing the values associated with the
-        states, which must be homogeneous in type. If None, the values
-        default to integers 0 through n-1.
+        states, which must be homogeneous in type. If None, the states
+        are represented by their indices, 0 through n-1.
 
     Attributes
     ----------

quantecon/markov/ddp.py

@@ -179,9 +179,21 @@ class DiscreteDP:
         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.
+        Array_like of length n containing the values associated with
+        the states, which must be homogeneous in type. May be
+        2-dimensional, in which case row `state_values[s]` is the value
+        associated with state `s`. If None, the states are represented
+        by their indices, 0 through n-1.
+
+    action_values : array_like, optional(default=None)
+        Array_like of length m containing the values associated with
+        the actions, which must be homogeneous in type. May be
+        2-dimensional, in which case row `action_values[a]` is the
+        value associated with action `a`. If None, the actions are
+        represented by their indices, 0 through m-1. In the
+        state-action pairs formulation, the length defines the number
+        of actions m, and may exceed `a_indices.max() + 1` (actions
+        feasible at no state).
 
     Attributes
     ----------
@@ -203,6 +215,9 @@ class DiscreteDP:
     state_values : array_like or None
         Array of state values if set, None otherwise.
 
+    action_values : array_like or None
+        Array of action values if set, None otherwise.
+
     Notes
     -----
     DiscreteDP accepts beta=1 for convenience. In this case, infinite
@@ -305,7 +320,7 @@ class DiscreteDP:
 
     """
     def __init__(self, R, Q, beta, s_indices=None, a_indices=None,
-                 state_values=None):
+                 state_values=None, action_values=None):
         if not (0 <= beta <= 1):
             raise ValueError('beta must be in [0, 1]')
         if beta == 1:
@@ -420,10 +435,18 @@ def __init__(self, R, Q, beta, s_indices=None, a_indices=None,
         # Check that for every state, at least one action is feasible
         self._check_action_feasibility()
 
+        # Number of actions
+        if self._sa_pair:
+            self._num_actions = self.a_indices.max() + 1
+        else:
+            self._num_actions = self.R.shape[1]
+
+        # Call the setter methods
+        self.state_values = state_values
+        self.action_values = action_values
+
         self.epsilon = 1e-3
         self.max_iter = 250
-        # Call the setter method
-        self.state_values = state_values
 
         # Linear equation solver to be used in evaluate_policy
         if self._sparse:
@@ -462,6 +485,52 @@ def state_values(self, values):
                 )
             self._state_values = values
 
+    @property
+    def action_values(self):
+        return self._action_values
+
+    @action_values.setter
+    def action_values(self, values):
+        """
+        Set action values of the DiscreteDP.
+
+        Parameters
+        ----------
+        values : array_like or None
+            Array of action values, or None to unset. For the product
+            formulation, must be of length m; for the state-action
+            pairs formulation, of length at least `a_indices.max() +
+            1`, and the length defines the number of actions.
+        """
+        if values is None:
+            self._action_values = None
+            if self._sa_pair:
+                # Restore the default number of actions
+                self._num_actions = self.a_indices.max() + 1
+        else:
+            values = np.asarray(values)
+            if self._sa_pair:
+                if (values.ndim < 1) or \
+                        (values.shape[0] < self.a_indices.max() + 1):
+                    raise ValueError(
+                        'action_values must be an array_like of length '
+                        'at least a_indices.max() + 1'
+                    )
+            else:
+                if (values.ndim < 1) or \
+                        (values.shape[0] != self._num_actions):
+                    raise ValueError(
+                        'action_values must be an array_like of length m'
+                    )
+            if np.issubdtype(values.dtype, np.object_):
+                raise ValueError(
+                    'data in action_values must be homogeneous in type'
+                )
+            self._action_values = values
+            if self._sa_pair:
+                # The length defines the number of actions
+                self._num_actions = values.shape[0]
+
     def _check_action_feasibility(self):
         """
         Check that for every state, reward is finite for some action,
@@ -552,7 +621,8 @@ def to_sa_pair_form(self, sparse=True):
                 QL = self.Q[s_ind, a_ind]
             return DiscreteDP(
                 RL, QL, self.beta, s_ind, a_ind,
-                state_values=self.state_values
+                state_values=self.state_values,
+                action_values=self.action_values
             )
 
     def to_product_form(self):
@@ -575,7 +645,7 @@ def to_product_form(self):
         """
         if self._sa_pair:
             ns = self.num_states
-            na = self.a_indices.max() + 1
+            na = self._num_actions
             R = np.full((ns, na), -np.inf)
             R[self.s_indices, self.a_indices] = self.R
             Q = np.zeros((ns, na, ns))
@@ -585,7 +655,8 @@ def to_product_form(self):
             else:
                 _fill_dense_Q(self.s_indices, self.a_indices, self.Q, Q)
             return DiscreteDP(
-                R, Q, self.beta, state_values=self.state_values
+                R, Q, self.beta, state_values=self.state_values,
+                action_values=self.action_values
             )
         else:
             return self
@@ -872,6 +943,8 @@ def value_iteration(self, v_init=None, epsilon=None, max_iter=None):
                             sigma=sigma,
                             num_iter=num_iter,
                             mc=self.controlled_mc(sigma),
+                            state_values=self.state_values,
+                            action_values=self.action_values,
                             method='value iteration',
                             epsilon=epsilon,
                             max_iter=max_iter)
@@ -914,6 +987,8 @@ def policy_iteration(self, v_init=None, max_iter=None):
                             sigma=sigma,
                             num_iter=num_iter,
                             mc=self.controlled_mc(sigma),
+                            state_values=self.state_values,
+                            action_values=self.action_values,
                             method='policy iteration',
                             max_iter=max_iter)
 
@@ -973,6 +1048,8 @@ def midrange(z):
                             sigma=sigma,
                             num_iter=num_iter,
                             mc=self.controlled_mc(sigma),
+                            state_values=self.state_values,
+                            action_values=self.action_values,
                             method='modified policy iteration',
                             epsilon=epsilon,
                             max_iter=max_iter,
@@ -1011,6 +1088,8 @@ def linprog_simplex(self, v_init=None, max_iter=None):
                             sigma=sigma,
                             num_iter=num_iter,
                             mc=self.controlled_mc(sigma),
+                            state_values=self.state_values,
+                            action_values=self.action_values,
                             method='linear programming',
                             max_iter=max_iter)
 
@@ -1028,7 +1107,8 @@ def controlled_mc(self, sigma):
         Returns
         -------
         mc : MarkovChain
-            Controlled Markov chain.
+            Controlled Markov chain, with `state_values` attached if
+            set for this instance.
 
         """
         _, Q_sigma = self.RQ_sigma(sigma)
@@ -1051,7 +1131,18 @@ class DPSolveResult(dict):
         Number of iterations
 
     mc : MarkovChain
-        Controlled Markov chain
+        Controlled Markov chain, with the `state_values` attached if
+        set
+
+    state_values : ndarray or None
+        State values of the `DiscreteDP` instance solved
+
+    action_values : ndarray or None
+        Action values of the `DiscreteDP` instance solved
+
+    sigma_values : ndarray
+        Computed optimal policy function, decoded to action values
+        (`sigma` itself if `action_values` is None)
 
     method : str
         Method employed
@@ -1063,6 +1154,20 @@ class DPSolveResult(dict):
         Maximum number of iterations
 
     """
+    @property
+    def sigma_values(self):
+        """
+        Return the optimal policy function decoded to action values,
+        i.e., the array whose s-th element is
+        `action_values[sigma[s]]`. If `action_values` is None, return
+        `sigma` itself.
+
+        """
+        action_values = self.get('action_values')
+        if action_values is None:
+            return self['sigma']
+        return action_values[self['sigma']]
+
     # This is sourced from sicpy.optimize.OptimizeResult.
     def __getattr__(self, name):
         try: