diff --git a/docs/source/util.rst b/docs/source/util.rst index 72742a99..794c8c22 100644 --- a/docs/source/util.rst +++ b/docs/source/util.rst @@ -8,6 +8,7 @@ Utilities util/combinatorics util/common_messages util/compat + util/indexing util/notebooks util/numba util/random diff --git a/docs/source/util/indexing.rst b/docs/source/util/indexing.rst new file mode 100644 index 00000000..6f1d4f20 --- /dev/null +++ b/docs/source/util/indexing.rst @@ -0,0 +1,7 @@ +indexing +======== + +.. automodule:: quantecon.util.indexing + :members: + :undoc-members: + :show-inheritance: diff --git a/quantecon/__init__.py b/quantecon/__init__.py index 8375508a..b88d762f 100644 --- a/quantecon/__init__.py +++ b/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 diff --git a/quantecon/markov/core.py b/quantecon/markov/core.py index ec844a4e..73dbd40f 100644 --- a/quantecon/markov/core.py +++ b/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 ---------- diff --git a/quantecon/markov/ddp.py b/quantecon/markov/ddp.py index 32f46396..bb5ab8bd 100644 --- a/quantecon/markov/ddp.py +++ b/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: diff --git a/quantecon/markov/tests/test_ddp.py b/quantecon/markov/tests/test_ddp.py index 6a573c4e..5d66cd7e 100644 --- a/quantecon/markov/tests/test_ddp.py +++ b/quantecon/markov/tests/test_ddp.py @@ -447,7 +447,6 @@ 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) @@ -455,7 +454,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, state_values=state_values) + ddp = DiscreteDP(R, Q, beta) ddp_sa = ddp.to_sa_pair_form() ddp_sa2 = ddp_sa.to_sa_pair_form() ddp_sa3 = ddp.to_sa_pair_form(sparse=False) @@ -469,7 +468,6 @@ def test_ddp_to_sa_and_to_product(): # allclose doesn't work on sparse assert_(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 @@ -482,13 +480,11 @@ 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) @@ -499,84 +495,221 @@ def test_ddp_to_sa_and_to_product(): 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) +class TestDiscreteDPStateValues: + def setup_method(self): + # From Puterman 2005, Section 3.1 + beta = 0.95 - ddp = DiscreteDP(R, Q, beta) - assert_(ddp.state_values is None) + # Formulation with R: n x m, Q: n x m x n + n, m = 2, 2 # number of states, number of actions + R = [[5, 10], [-1, -np.inf]] + Q = np.empty((n, m, n)) + Q[0, 0, :] = 0.5, 0.5 + Q[0, 1, :] = 0, 1 + Q[1, :, :] = 0, 1 - new_state_values = ['new_state1', 'new_state2'] - ddp.state_values = new_state_values - assert_(isinstance(ddp.state_values, np.ndarray)) - assert_array_equal(ddp.state_values, new_state_values) + # Formulation with state-action pairs + s_indices = [0, 0, 1] + a_indices = [0, 1, 0] + R_sa = [R[0][0], R[0][1], R[1][0]] + Q_sa = np.asarray(Q)[s_indices, a_indices] - ddp.state_values = None - assert_(ddp.state_values is None) + self.state_values = np.array([1.2, 3.4]) + ddp0 = DiscreteDP(R, Q, beta, state_values=self.state_values) + ddp_sa = DiscreteDP(R_sa, Q_sa, beta, s_indices, a_indices, + state_values=self.state_values) + self.ddps = [ddp0, ddp_sa] + def test_state_values_stored(self): + for ddp in self.ddps: + assert_array_equal(ddp.state_values, self.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]) + def test_state_values_default_none(self): + R = [[5, 10], [-1, -np.inf]] + Q = [[(0.5, 0.5), (0, 1)], [(0, 1), (0.5, 0.5)]] + ddp = DiscreteDP(R, Q, 0.95) + assert_(ddp.state_values is None) - ddp = DiscreteDP(R, Q, beta, state_values=state_values) - mc = ddp.controlled_mc(sigma) + def test_state_values_setter(self): + for ddp in self.ddps: + ddp.state_values = [5, 6] + assert_array_equal(ddp.state_values, [5, 6]) + ddp.state_values = None + assert_(ddp.state_values is None) + ddp.state_values = self.state_values + + def test_state_values_2dim(self): + # Values may be vectors (e.g., (asset, productivity) pairs) + values_2d = np.array([[0., 0.1], [1., 0.1]]) + for ddp in self.ddps: + ddp.state_values = values_2d + assert_array_equal(ddp.state_values, values_2d) + ddp.state_values = self.state_values - assert_array_equal(mc.state_values, state_values) + def test_state_values_invalid(self): + msg_length = 'state_values must be an array_like of length n' + msg_dtype = 'data in state_values must be homogeneous in type' + for ddp in self.ddps: + for values in [[1.2, 3.4, 5.6], # Wrong length + 'state0', # Scalar + [(0,), (0, 1)], # Non-homogeneous + np.array(['state0', 1], dtype=object)]: + assert_raises(ValueError, setattr, ddp, 'state_values', + values) + with assert_raises(ValueError) as exc_info: + ddp.state_values = [1.2, 3.4, 5.6] + assert_(str(exc_info.exception) == msg_length) + with assert_raises(ValueError) as exc_info: + ddp.state_values = np.array(['state0', 1], dtype=object) + assert_(str(exc_info.exception) == msg_dtype) + R = [[5, 10], [-1, -np.inf]] + Q = [[(0.5, 0.5), (0, 1)], [(0, 1), (0.5, 0.5)]] + assert_raises(ValueError, DiscreteDP, R, Q, 0.95, + state_values=[1.2]) + def test_controlled_mc_state_values(self): + for ddp in self.ddps: + for method in ['vi', 'pi', 'mpi']: + res = ddp.solve(method=method) + assert_array_equal(res.mc.state_values, self.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 - state_values = np.array(['state1']) - msg = 'state_values must be an array_like of length n' + def test_state_values_recorded_in_result(self): + for ddp in self.ddps: + res = ddp.solve(method='pi') + assert_array_equal(res.state_values, self.state_values) + R = [[5, 10], [-1, -np.inf]] + Q = [[(0.5, 0.5), (0, 1)], [(0, 1), (0.5, 0.5)]] + res = DiscreteDP(R, Q, 0.95).solve(method='pi') + assert_(res.state_values is None) - with assert_raises(ValueError) as exc_info: - DiscreteDP(R, Q, beta, state_values=state_values) - assert_(str(exc_info.exception) == msg) + def test_to_sa_and_to_product_carry_state_values(self): + ddp0, ddp_sa = self.ddps + for ddp_new in [ddp0.to_sa_pair_form(), ddp0.to_sa_pair_form(False), + ddp_sa.to_product_form()]: + assert_array_equal(ddp_new.state_values, self.state_values) - ddp = DiscreteDP(R, Q, beta) - with assert_raises(ValueError) as exc_info: - ddp.state_values = state_values - assert_(str(exc_info.exception) == msg) +class TestDiscreteDPActionValues: + def setup_method(self): + # From Puterman 2005, Section 3.1 + beta = 0.95 -def test_ddp_scalar_state_values(): - R = np.array([[1, 2], [3, 4]]) - Q = np.full((2, 2, 2), 0.5) - beta = 0.95 - state_values = 'state1' + # Formulation with R: n x m, Q: n x m x n + n, m = 2, 2 # number of states, number of actions + R = [[5, 10], [-1, -np.inf]] + Q = np.empty((n, m, n)) + Q[0, 0, :] = 0.5, 0.5 + Q[0, 1, :] = 0, 1 + Q[1, :, :] = 0, 1 - assert_raises( - ValueError, DiscreteDP, R, Q, beta, state_values=state_values - ) + # Formulation with state-action pairs + s_indices = [0, 0, 1] + a_indices = [0, 1, 0] + R_sa = [R[0][0], R[0][1], R[1][0]] + Q_sa = np.asarray(Q)[s_indices, a_indices] - ddp = DiscreteDP(R, Q, beta) - assert_raises(ValueError, setattr, ddp, 'state_values', state_values) + self.action_values = np.array([-9.9, -8.8]) + ddp0 = DiscreteDP(R, Q, beta, action_values=self.action_values) + ddp_sa = DiscreteDP(R_sa, Q_sa, beta, s_indices, a_indices, + action_values=self.action_values) + self.ddps = [ddp0, ddp_sa] + def test_action_values_stored(self): + for ddp in self.ddps: + assert_array_equal(ddp.action_values, self.action_values) -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) - msg = 'data in state_values must be homogeneous in type' + def test_action_values_default_none(self): + R = [[5, 10], [-1, -np.inf]] + Q = [[(0.5, 0.5), (0, 1)], [(0, 1), (0.5, 0.5)]] + ddp = DiscreteDP(R, Q, 0.95) + assert_(ddp.action_values is None) - with assert_raises(ValueError) as exc_info: - DiscreteDP(R, Q, beta, state_values=state_values) - assert_(str(exc_info.exception) == msg) + def test_action_values_setter(self): + for ddp in self.ddps: + ddp.action_values = [5, 6] + assert_array_equal(ddp.action_values, [5, 6]) + ddp.action_values = None + assert_(ddp.action_values is None) + ddp.action_values = self.action_values + + def test_action_values_2dim(self): + # Values may be vectors + values_2d = np.array([[0., 0.1], [1., 0.1]]) + for ddp in self.ddps: + ddp.action_values = values_2d + assert_array_equal(ddp.action_values, values_2d) + ddp.action_values = self.action_values + + def test_action_values_invalid(self): + msg_length_prod = 'action_values must be an array_like of length m' + msg_length_sa = 'action_values must be an array_like of length ' + \ + 'at least a_indices.max() + 1' + msg_dtype = 'data in action_values must be homogeneous in type' + for ddp in self.ddps: + for values in [[1.2], # Too short + 'action0', # Scalar + np.array(['action0', 1], dtype=object)]: + assert_raises(ValueError, setattr, ddp, 'action_values', + values) + with assert_raises(ValueError) as exc_info: + ddp.action_values = np.array(['action0', 1], dtype=object) + assert_(str(exc_info.exception) == msg_dtype) + ddp0, ddp_sa = self.ddps + assert_raises(ValueError, setattr, ddp0, 'action_values', + [1.2, 3.4, 5.6]) # Wrong length (product form) + assert_raises(ValueError, setattr, ddp0, 'action_values', + [(0,), (0, 1)]) # Non-homogeneous + with assert_raises(ValueError) as exc_info: + ddp0.action_values = [1.2] + assert_(str(exc_info.exception) == msg_length_prod) + with assert_raises(ValueError) as exc_info: + ddp_sa.action_values = [1.2] + assert_(str(exc_info.exception) == msg_length_sa) + R = [[5, 10], [-1, -np.inf]] + Q = [[(0.5, 0.5), (0, 1)], [(0, 1), (0.5, 0.5)]] + assert_raises(ValueError, DiscreteDP, R, Q, 0.95, + action_values=[1.2]) + + def test_action_values_longer_sa(self): + # In the sa formulation, the length of action_values defines + # the number of actions, which may exceed a_indices.max() + 1 + # (actions feasible at no state) + _, ddp_sa = self.ddps + values_longer = np.array([-9.9, -8.8, -7.7]) + ddp_sa.action_values = values_longer + assert_array_equal(ddp_sa.action_values, values_longer) + ddp_prod = ddp_sa.to_product_form() + assert_(ddp_prod.R.shape[1] == len(values_longer)) + assert_(np.isneginf(ddp_prod.R[:, -1]).all()) + assert_array_equal(ddp_prod.action_values, values_longer) + # Round trip back to sa form preserves the values + ddp_sa2 = ddp_prod.to_sa_pair_form() + assert_array_equal(ddp_sa2.action_values, values_longer) + ddp_sa.action_values = self.action_values + + def test_action_values_recorded_in_result(self): + for ddp in self.ddps: + res = ddp.solve(method='pi') + assert_array_equal(res.action_values, self.action_values) + R = [[5, 10], [-1, -np.inf]] + Q = [[(0.5, 0.5), (0, 1)], [(0, 1), (0.5, 0.5)]] + res = DiscreteDP(R, Q, 0.95).solve(method='pi') + assert_(res.action_values is None) - ddp = DiscreteDP(R, Q, beta) - with assert_raises(ValueError) as exc_info: - ddp.state_values = state_values - assert_(str(exc_info.exception) == msg) + def test_sigma_values(self): + for ddp in self.ddps: + for method in ['vi', 'pi', 'mpi']: + res = ddp.solve(method=method) + assert_array_equal(res.sigma_values, + self.action_values[res.sigma]) + # With action_values unset, sigma_values is sigma itself + R = [[5, 10], [-1, -np.inf]] + Q = [[(0.5, 0.5), (0, 1)], [(0, 1), (0.5, 0.5)]] + res = DiscreteDP(R, Q, 0.95).solve(method='pi') + assert_array_equal(res.sigma_values, res.sigma) + + def test_to_sa_and_to_product_carry_action_values(self): + ddp0, ddp_sa = self.ddps + for ddp_new in [ddp0.to_sa_pair_form(), ddp0.to_sa_pair_form(False), + ddp_sa.to_product_form()]: + assert_array_equal(ddp_new.action_values, self.action_values) diff --git a/quantecon/util/__init__.py b/quantecon/util/__init__.py index 8ecb232d..489cacdc 100644 --- a/quantecon/util/__init__.py +++ b/quantecon/util/__init__.py @@ -4,6 +4,7 @@ """ from .array import searchsorted +from .indexing import index_dict from .notebooks import fetch_nb_dependencies from .random import check_random_state, rng_integers from .timing import tic, tac, toc, loop_timer, Timer, timeit diff --git a/quantecon/util/indexing.py b/quantecon/util/indexing.py new file mode 100644 index 00000000..c83adcf0 --- /dev/null +++ b/quantecon/util/indexing.py @@ -0,0 +1,155 @@ +""" +Indexing Utilities + +Utilities +--------- +index_dict + +""" +import numpy as np +from numba import jit, types +from numba.typed import Dict + + +@jit(nopython=True, cache=True) +def _fill_1d(dd, arr): + for i in range(arr.shape[0]): + dd[arr[i]] = i + + +# Cache of jitted fill functions for 2-dimensional values arrays, keyed +# by the number of columns (= the length of the tuple keys), which must +# be a compile-time constant for tuple construction under Numba +_fill_nd_cache = {} + + +def _make_fill_nd(d): + fill = _fill_nd_cache.get(d) + if fill is None: + elems = ", ".join(f"arr[i, {k}]" for k in range(d)) + code = ( + "def _fill(dd, arr):\n" + " for i in range(arr.shape[0]):\n" + f" dd[({elems},)] = i\n" + ) + ns = {} + exec(code, ns) + fill = jit(nopython=True)(ns["_fill"]) + _fill_nd_cache[d] = fill + return fill + + +def _diagnose_nonfinite(arr): + isfinite = np.isfinite(arr) + if arr.ndim == 2: + isfinite = isfinite.all(axis=1) + i = int(np.argmin(isfinite)) + v = arr[i] if arr.ndim == 1 else tuple(arr[i]) + raise ValueError(f"values contains non-finite value {v} at index {i}") + + +def _diagnose_negative_zero(arr): + is_negzero = np.signbit(arr) & (arr == 0) + if arr.ndim == 2: + is_negzero = is_negzero.any(axis=1) + i = int(np.argmax(is_negzero)) + raise ValueError( + f"values contains -0.0 at index {i}; use 0.0 instead" + ) + + +def _diagnose_duplicates(arr): + seen = {} + for i in range(arr.shape[0]): + key = arr[i] if arr.ndim == 1 else tuple(arr[i]) + if key in seen: + raise ValueError( + f"duplicate value {key} at indices {seen[key]} and {i}" + ) + seen[key] = i + + +def index_dict(values): + """ + Build a Numba typed dictionary mapping each value in `values` to its + index. + + Parameters + ---------- + values : array_like + Array of unique values, of shape (n,) for scalar values or + (n, d) for vector values. dtype must be of float or integer + kind (the values are stored as float64 or int64 keys); values + must be finite. + + Returns + ------- + numba.typed.Dict + Dict mapping value -> index, with int64 indices. For an input + of shape (n,), keys are scalars (float64 or int64, following + the dtype kind of `values`); for shape (n, d), keys are + d-tuples thereof. Usable both from the interpreter and inside + Numba-jitted functions, and passable to jitted functions as an + argument. + + Raises + ------ + ValueError + If `values` contains duplicate, non-finite, or negative-zero + values, or is not 1- or 2-dimensional. + + Notes + ----- + Lookup is by exact equality: query keys must be values drawn from + `values` itself (or arithmetic that reproduces them bitwise). For + approximate location on a sorted grid, use `np.searchsorted` + instead. + + Numba typed dicts hash floats by bit pattern, so `-0.0` and `0.0` + are distinct keys (unlike in a Python dict). `values` therefore + must not contain `-0.0`; a query key that may have been computed as + `-0.0` can be normalized by adding `0.0` to it. + + Examples + -------- + >>> d = index_dict(np.array([0.1, 1.0])) + >>> d[1.0] + 1 + >>> d2 = index_dict(np.array([[0., 0.1], [0., 1.], [0.5, 0.1]])) + >>> d2[(0.5, 0.1)] + 2 + + """ + arr = np.asarray(values) + if arr.ndim not in (1, 2): + raise ValueError("values must be 1- or 2-dimensional") + if arr.ndim == 2 and arr.shape[1] == 0: + raise ValueError("values must have at least one column") + + if arr.dtype.kind == 'f': + arr = arr.astype(np.float64, copy=False) + scalar_type = types.float64 + if not np.isfinite(arr).all(): + _diagnose_nonfinite(arr) + if np.any(np.signbit(arr) & (arr == 0)): + _diagnose_negative_zero(arr) + elif arr.dtype.kind in 'iu': + arr = arr.astype(np.int64, copy=False) + scalar_type = types.int64 + else: + raise ValueError(f"unsupported dtype: {arr.dtype}") + + if arr.ndim == 1: + key_type = scalar_type + fill = _fill_1d + else: + key_type = types.UniTuple(scalar_type, arr.shape[1]) + fill = _make_fill_nd(arr.shape[1]) + + dd = Dict.empty(key_type, types.int64) + fill(dd, arr) + + if len(dd) != arr.shape[0]: + _diagnose_duplicates(arr) + + return dd diff --git a/quantecon/util/tests/test_indexing.py b/quantecon/util/tests/test_indexing.py new file mode 100644 index 00000000..ead79d5f --- /dev/null +++ b/quantecon/util/tests/test_indexing.py @@ -0,0 +1,140 @@ +""" +Tests for Indexing Utilities + +Functions +--------- +index_dict + +""" +import numpy as np +from numpy.testing import assert_ +from numba import njit +import pytest +from quantecon.util import index_dict + + +class TestIndexDict1D: + def setup_method(self): + self.vals = np.array([0.5, -1.0, 2.5]) + self.d = index_dict(self.vals) + + def test_lookup(self): + for i, v in enumerate(self.vals): + assert_(self.d[v] == i) + + def test_length(self): + assert_(len(self.d) == len(self.vals)) + + def test_use_in_jitted_function(self): + @njit + def f(dd, v): + return dd[v] + + assert_(f(self.d, -1.0) == 1) + + def test_pass_as_argument_and_modify(self): + @njit + def contains(dd, v): + return v in dd + + assert_(contains(self.d, 2.5)) + assert_(not contains(self.d, 10.0)) + + +class TestIndexDict2D: + def setup_method(self): + self.vals = np.array([[0.0, 0.1], [0.0, 1.0], [0.5, 0.1], + [0.5, 1.0]]) + self.d = index_dict(self.vals) + + def test_lookup(self): + for i in range(self.vals.shape[0]): + assert_(self.d[tuple(self.vals[i])] == i) + + def test_use_in_jitted_function(self): + @njit + def f(dd, row): + return dd[(row[0], row[1])] + + assert_(f(self.d, self.vals[2]) == 2) + + +def test_index_dict_3_columns(): + vals = np.arange(12.).reshape(4, 3) + d = index_dict(vals) + for i in range(vals.shape[0]): + assert_(d[tuple(vals[i])] == i) + + +def test_index_dict_int_dtype(): + vals = np.array([2, 0, 5]) + d = index_dict(vals) + for i, v in enumerate(vals): + assert_(d[v] == i) + + +def test_index_dict_int_dtype_2d(): + vals = np.array([[0, 1], [1, 0]]) + d = index_dict(vals) + assert_(d[(1, 0)] == 1) + + +def test_index_dict_array_like(): + d = index_dict([0.1, 0.2]) + assert_(d[0.2] == 1) + + +def test_index_dict_empty(): + d = index_dict(np.empty(0)) + assert_(len(d) == 0) + + +def test_index_dict_duplicates_raise(): + with pytest.raises(ValueError, match="duplicate value"): + index_dict(np.array([0.5, 1.0, 0.5])) + + +def test_index_dict_duplicates_raise_2d(): + with pytest.raises(ValueError, match="duplicate value"): + index_dict(np.array([[0.5, 1.0], [0.5, 1.0]])) + + +def test_index_dict_negative_zero_raises(): + # Numba typed dicts hash floats by bit pattern, so -0.0 and 0.0 + # would be distinct keys; -0.0 is rejected at build time + with pytest.raises(ValueError, match="-0.0"): + index_dict(np.array([1.0, -0.0])) + with pytest.raises(ValueError, match="-0.0"): + index_dict(np.array([[1.0, -0.0]])) + + +def test_index_dict_positive_zero_ok(): + d = index_dict(np.array([0.0, 1.0])) + assert_(d[0.0] == 0) + + +def test_index_dict_nan_raises(): + with pytest.raises(ValueError, match="non-finite"): + index_dict(np.array([0.5, np.nan])) + + +def test_index_dict_inf_raises(): + with pytest.raises(ValueError, match="non-finite"): + index_dict(np.array([[0.5, np.inf]])) + + +def test_index_dict_invalid_ndim(): + with pytest.raises(ValueError, match="1- or 2-dimensional"): + index_dict(np.zeros((2, 2, 2))) + with pytest.raises(ValueError, match="1- or 2-dimensional"): + index_dict(np.float64(1.0)) + + +def test_index_dict_zero_columns(): + with pytest.raises(ValueError, match="at least one column"): + index_dict(np.empty((3, 0))) + + +def test_index_dict_invalid_dtype(): + with pytest.raises(ValueError, match="unsupported dtype"): + index_dict(np.array(['a', 'b']))