From a4ee42429349493753f9a4ebf854cb18d6255743 Mon Sep 17 00:00:00 2001 From: Daisuke Oyama Date: Sun, 16 Aug 2026 13:58:18 +0900 Subject: [PATCH 1/4] DOC: Correct the state_values docstring in MarkovChain The attribute does not default to integers 0 through n-1; it is None when unset, in which case the states are represented by their indices (as the Attributes entry and the simulate docstring already state). Co-Authored-By: Claude Fable 5 --- quantecon/markov/core.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 ---------- From 5b1672af7d991590393722ae7d8869fe157c4f59 Mon Sep 17 00:00:00 2001 From: Daisuke Oyama Date: Sun, 16 Aug 2026 13:44:11 +0900 Subject: [PATCH 2/4] ENH: Refine DiscreteDP state_values; record in DPSolveResult Refine the state_values support added in #832: * Correct the docstring: the attribute is None when unset (in which case the states are represented by their indices), and the values may be 2-dimensional (row per state). * Record state_values in DPSolveResult. * Reorganize the tests into a class covering both formulations, 2-dimensional values, the solution pipeline, and carry-over by to_sa_pair_form/to_product_form. Cf. #248 Co-Authored-By: Claude Fable 5 --- quantecon/markov/ddp.py | 26 +++-- quantecon/markov/tests/test_ddp.py | 161 +++++++++++++++-------------- 2 files changed, 103 insertions(+), 84 deletions(-) diff --git a/quantecon/markov/ddp.py b/quantecon/markov/ddp.py index 32f46396..3b7a7acf 100644 --- a/quantecon/markov/ddp.py +++ b/quantecon/markov/ddp.py @@ -179,9 +179,11 @@ 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. Attributes ---------- @@ -420,11 +422,12 @@ 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() - self.epsilon = 1e-3 - self.max_iter = 250 # Call the setter method self.state_values = state_values + self.epsilon = 1e-3 + self.max_iter = 250 + # Linear equation solver to be used in evaluate_policy if self._sparse: import scipy.sparse.linalg @@ -872,6 +875,7 @@ 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, method='value iteration', epsilon=epsilon, max_iter=max_iter) @@ -914,6 +918,7 @@ 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, method='policy iteration', max_iter=max_iter) @@ -973,6 +978,7 @@ def midrange(z): sigma=sigma, num_iter=num_iter, mc=self.controlled_mc(sigma), + state_values=self.state_values, method='modified policy iteration', epsilon=epsilon, max_iter=max_iter, @@ -1011,6 +1017,7 @@ 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, method='linear programming', max_iter=max_iter) @@ -1028,7 +1035,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 +1059,11 @@ 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 method : str Method employed diff --git a/quantecon/markov/tests/test_ddp.py b/quantecon/markov/tests/test_ddp.py index 6a573c4e..49343d02 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,95 @@ 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) - - ddp = DiscreteDP(R, Q, beta) - assert_(ddp.state_values is None) - - 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) - - ddp.state_values = None - assert_(ddp.state_values is None) - - -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 - state_values = np.array(['state1']) - msg = 'state_values must be an array_like of length n' - - with assert_raises(ValueError) as exc_info: - DiscreteDP(R, Q, beta, state_values=state_values) - assert_(str(exc_info.exception) == msg) +class TestDiscreteDPStateValues: + def setup_method(self): + # From Puterman 2005, Section 3.1 + beta = 0.95 - ddp = DiscreteDP(R, Q, beta) - with assert_raises(ValueError) as exc_info: - ddp.state_values = state_values - assert_(str(exc_info.exception) == msg) + # 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 + # 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] -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' + 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] - assert_raises( - ValueError, DiscreteDP, R, Q, beta, state_values=state_values - ) + def test_state_values_stored(self): + for ddp in self.ddps: + assert_array_equal(ddp.state_values, self.state_values) - ddp = DiscreteDP(R, Q, beta) - assert_raises(ValueError, setattr, ddp, 'state_values', state_values) + 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) + 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 -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_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]) - with assert_raises(ValueError) as exc_info: - DiscreteDP(R, Q, beta, state_values=state_values) - assert_(str(exc_info.exception) == msg) + 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) - 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_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) + + 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) From 7a7e7fbaf960a1b64248afb0457cdaf975fc80e6 Mon Sep 17 00:00:00 2001 From: Daisuke Oyama Date: Sun, 16 Aug 2026 13:48:40 +0900 Subject: [PATCH 3/4] ENH: Add action_values to DiscreteDP Add an action_values keyword argument to DiscreteDP, symmetric to state_values. The values play no role in the solution algorithms. DPSolveResult records action_values, and gains a sigma_values property returning the optimal policy decoded to action values, i.e., action_values[sigma] (sigma itself if action_values is None). The semantics follow the "universal action set" approach (approach 2 of QuantEcon/QuantEcon.jl#94 (comment); cf. QuantEcon/QuantEcon.jl#402): action_values labels a common action set A with A(s) a subset of A for each state s. In the state-action pairs formulation, its length defines the number of actions and may exceed a_indices.max() + 1, so a universal action set strictly containing the union of the feasible sets is representable in both formulations. to_sa_pair_form/to_product_form carry action_values unchanged, with to_product_form restoring the full number of columns, so actions feasible at no state survive the round trip. Cf. #248 Co-Authored-By: Claude Fable 5 --- quantecon/markov/ddp.py | 103 +++++++++++++++++++++-- quantecon/markov/tests/test_ddp.py | 126 +++++++++++++++++++++++++++++ 2 files changed, 224 insertions(+), 5 deletions(-) diff --git a/quantecon/markov/ddp.py b/quantecon/markov/ddp.py index 3b7a7acf..bb5ab8bd 100644 --- a/quantecon/markov/ddp.py +++ b/quantecon/markov/ddp.py @@ -185,6 +185,16 @@ class DiscreteDP: 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 ---------- R, Q, beta : see Parameters. @@ -205,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 @@ -307,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: @@ -422,8 +435,15 @@ 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() - # Call the setter method + # 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 @@ -465,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, @@ -555,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): @@ -578,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)) @@ -588,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 @@ -876,6 +944,7 @@ def value_iteration(self, v_init=None, epsilon=None, max_iter=None): 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) @@ -919,6 +988,7 @@ def policy_iteration(self, v_init=None, max_iter=None): 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) @@ -979,6 +1049,7 @@ def midrange(z): 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, @@ -1018,6 +1089,7 @@ def linprog_simplex(self, v_init=None, max_iter=None): 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) @@ -1065,6 +1137,13 @@ class DPSolveResult(dict): 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 @@ -1075,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 49343d02..5d66cd7e 100644 --- a/quantecon/markov/tests/test_ddp.py +++ b/quantecon/markov/tests/test_ddp.py @@ -587,3 +587,129 @@ def test_to_sa_and_to_product_carry_state_values(self): 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) + + +class TestDiscreteDPActionValues: + def setup_method(self): + # From Puterman 2005, Section 3.1 + beta = 0.95 + + # 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 + + # 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] + + 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_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) + + 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) + + 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) From 2dd0aef98ebafeaffb07de3986b0786ccf8e11c7 Mon Sep 17 00:00:00 2001 From: Daisuke Oyama Date: Sun, 16 Aug 2026 13:49:11 +0900 Subject: [PATCH 4/4] ENH: Add util.index_dict Add `index_dict`, which builds a Numba typed dict mapping each value in an array of unique values (scalars or rows) to its index, usable both from the interpreter and inside jitted functions. This is the value-to-index translation layer for building DiscreteDP instances in value space (with `state_values`/`action_values` attached), and for the interface proposed in #228. Notable: Numba typed dicts hash floats by bit pattern, so -0.0 and 0.0 are distinct keys (unlike Python dicts); -0.0 is therefore rejected at build time. Co-Authored-By: Claude Fable 5 --- docs/source/util.rst | 1 + docs/source/util/indexing.rst | 7 ++ quantecon/__init__.py | 3 +- quantecon/util/__init__.py | 1 + quantecon/util/indexing.py | 155 ++++++++++++++++++++++++++ quantecon/util/tests/test_indexing.py | 140 +++++++++++++++++++++++ 6 files changed, 306 insertions(+), 1 deletion(-) create mode 100644 docs/source/util/indexing.rst create mode 100644 quantecon/util/indexing.py create mode 100644 quantecon/util/tests/test_indexing.py 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/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']))